Social LSTM实战：用Python复现行人轨迹预测模型（附代码）

Social LSTM实战从零构建行人轨迹预测系统行人轨迹预测一直是计算机视觉和机器人导航领域的核心挑战。想象一下当你走在拥挤的商场里会不自觉地调整步伐和路线避开迎面而来的人群——这种看似简单的行为背后隐藏着复杂的社交互动模式。本文将带你用Python实现一个能够理解这种社交行为的深度学习模型Social LSTM。1. 环境准备与数据加载在开始构建模型前我们需要配置合适的开发环境并准备训练数据。Social LSTM对计算资源有一定要求建议使用配备GPU的工作站或云服务器。基础环境配置conda create -n social_lstm python3.8 conda activate social_lstm pip install torch1.10.0cu113 torchvision0.11.1cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html pip install numpy pandas matplotlib scikit-learn我们将使用ETH和UCY这两个公开的行人轨迹数据集。这些数据集包含了真实场景中行人运动的坐标序列是轨迹预测任务的基准测试集。数据预处理关键步骤坐标归一化将所有轨迹坐标转换到相对坐标系序列分割按8帧观察12帧预测的标准划分样本社交关系构建基于空间距离确定行人间的交互关系def load_eth_dataset(data_path): 加载并预处理ETH数据集 raw_data pd.read_csv(data_path) # 坐标归一化 scene_min raw_data[[x, y]].min() scene_max raw_data[[x, y]].max() raw_data[x] (raw_data[x] - scene_min[x]) / (scene_max[x] - scene_min[x]) raw_data[y] (raw_data[y] - scene_min[y]) / (scene_max[y] - scene_min[y]) return raw_data注意实际应用中应考虑数据增强技术如随机旋转和平移以提高模型泛化能力。2. 模型架构设计Social LSTM的核心创新在于其社交池化层该层使模型能够捕捉行人之间的交互模式。我们将使用PyTorch实现这一复杂架构。2.1 基础LSTM模块首先构建标准的LSTM单元用于学习单个行人的运动模式import torch.nn as nn class TrajectoryLSTM(nn.Module): def __init__(self, input_dim2, embedding_dim64, hidden_dim128): super().__init__() self.embedding nn.Linear(input_dim, embedding_dim) self.lstm nn.LSTM(embedding_dim, hidden_dim, batch_firstTrue) def forward(self, x): embedded torch.relu(self.embedding(x)) output, (hidden, cell) self.lstm(embedded) return output, hidden, cell2.2 社交池化层实现社交池化层是Social LSTM区别于普通LSTM的关键组件它负责聚合邻近行人的隐藏状态class SocialPooling(nn.Module): def __init__(self, pool_size32, hidden_dim128): super().__init__() self.pool_size pool_size self.hidden_dim hidden_dim self.pool_embedding nn.Linear(pool_size*pool_size*hidden_dim, hidden_dim) def forward(self, hidden_states, positions, neighbor_masks): batch_size, num_peds hidden_states.size(0), hidden_states.size(1) # 初始化社交张量 social_tensor torch.zeros(batch_size, num_peds, self.pool_size, self.pool_size, self.hidden_dim).to(hidden_states.device) # 填充社交张量 for b in range(batch_size): for i in range(num_peds): for j in range(num_peds): if neighbor_masks[b,i,j] 1: # 是邻居 rel_pos positions[b,j] - positions[b,i] grid_x int((rel_pos[0] 1) * (self.pool_size - 1) / 2) grid_y int((rel_pos[1] 1) * (self.pool_size - 1) / 2) if 0 grid_x self.pool_size and 0 grid_y self.pool_size: social_tensor[b,i,grid_x,grid_y] hidden_states[b,j] # 嵌入聚合后的社交信息 social_tensor social_tensor.view(batch_size, num_peds, -1) social_embedded torch.relu(self.pool_embedding(social_tensor)) return social_embedded提示实际实现时可使用向量化操作替代循环显著提升计算效率。3. 完整Social LSTM实现将基础LSTM与社交池化层结合构建完整的Social LSTM模型class SocialLSTM(nn.Module): def __init__(self, args): super().__init__() self.args args self.traj_lstm TrajectoryLSTM() self.social_pooling SocialPooling() self.position_predictor nn.Linear(2*args.hidden_dim, 5) # 预测二元高斯参数 def forward(self, traj_batch): # 处理输入轨迹 obs_traj traj_batch[observed] batch_size, num_peds, seq_len, _ obs_traj.size() # 初始化隐藏状态 hidden_states torch.zeros(batch_size, num_peds, self.args.hidden_dim).to(obs_traj.device) # 处理观察序列 for t in range(seq_len): current_pos obs_traj[:,:,t,:] _, hiddens, _ self.traj_lstm(current_pos.unsqueeze(1)) hidden_states hiddens.squeeze(0) # 计算社交信息 neighbor_masks self._get_neighbor_masks(current_pos) social_embeddings self.social_pooling(hidden_states, current_pos, neighbor_masks) # 更新隐藏状态 combined torch.cat([hidden_states, social_embeddings], dim-1) hidden_states self.traj_lstm.lstm_cell(combined, hidden_states) # 预测未来轨迹 pred_traj [] last_pos obs_traj[:,:,-1,:] for _ in range(self.args.pred_len): # 预测下一步位置 gaussian_params self.position_predictor(hidden_states) next_pos self._sample_from_gaussian(gaussian_params) pred_traj.append(next_pos) # 更新状态 _, hiddens, _ self.traj_lstm(next_pos.unsqueeze(1)) hidden_states hiddens.squeeze(0) # 更新社交信息 neighbor_masks self._get_neighbor_masks(next_pos) social_embeddings self.social_pooling(hidden_states, next_pos, neighbor_masks) combined torch.cat([hidden_states, social_embeddings], dim-1) hidden_states self.traj_lstm.lstm_cell(combined, hidden_states) return torch.stack(pred_traj, dim2)4. 模型训练与优化Social LSTM的训练需要特别注意批处理策略和损失函数设计因为每个场景中的行人数量可能不同。4.1 自定义数据加载器from torch.utils.data import Dataset, DataLoader class TrajectoryDataset(Dataset): def __init__(self, data, obs_len8, pred_len12): self.data data self.obs_len obs_len self.pred_len pred_len def __len__(self): return len(self.data[scene_ids]) def __getitem__(self, idx): scene_id self.data[scene_ids][idx] frame_ids self.data[scene_frame_map][scene_id] num_peds len(self.data[scene_ped_map][scene_id]) # 提取轨迹序列 obs_traj torch.zeros(num_peds, self.obs_len, 2) pred_traj torch.zeros(num_peds, self.pred_len, 2) for i, ped_id in enumerate(self.data[scene_ped_map][scene_id]): full_traj self.data[trajectories][ped_id] obs_traj[i] torch.from_numpy(full_traj[:self.obs_len]) pred_traj[i] torch.from_numpy(full_traj[self.obs_len:self.obs_lenself.pred_len]) return { observed: obs_traj, predicted: pred_traj, scene_id: scene_id }4.2 训练循环实现def train(model, dataloader, optimizer, device): model.train() total_loss 0 for batch in dataloader: obs_traj batch[observed].to(device) pred_traj batch[predicted].to(device) # 前向传播 pred_output model(batch) # 计算负对数似然损失 loss gaussian_2d_loss(pred_output, pred_traj) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() return total_loss / len(dataloader) def gaussian_2d_loss(pred_params, true_pos): 计算二元高斯分布的负对数似然 mu_x pred_params[..., 0] mu_y pred_params[..., 1] sigma_x torch.exp(pred_params[..., 2]) sigma_y torch.exp(pred_params[..., 3]) rho torch.tanh(pred_params[..., 4]) x true_pos[..., 0] y true_pos[..., 1] z_x ((x - mu_x) / sigma_x)**2 z_y ((y - mu_y) / sigma_y)**2 z_xy (x - mu_x)*(y - mu_y)/(sigma_x*sigma_y) nll torch.log(2*np.pi*sigma_x*sigma_y*torch.sqrt(1-rho**2)) \ 1/(2*(1-rho**2)) * (z_x z_y - 2*rho*z_xy) return torch.mean(nll)5. 评估与结果分析训练完成后我们需要定量评估模型性能并与基线方法进行比较。使用三个标准指标平均位移误差(ADE)所有预测时间步位置误差的平均最终位移误差(FDE)预测终点与真实终点的距离非线性位移误差(NL-ADE)轨迹非线性区域的误差评估代码实现def evaluate(model, dataloader, device): model.eval() metrics {ade: 0, fde: 0, nl_ade: 0} with torch.no_grad(): for batch in dataloader: obs_traj batch[observed].to(device) pred_traj batch[predicted].to(device) pred_output model(batch) pred_pos pred_output[..., :2] # 取均值作为预测位置 # 计算各项指标 metrics[ade] torch.mean(torch.norm(pred_pos - pred_traj, dim-1)).item() metrics[fde] torch.mean(torch.norm(pred_pos[:,:,-1] - pred_traj[:,:,-1], dim-1)).item() # 识别非线性区域 is_nonlinear identify_nonlinear_regions(pred_traj) nl_ade torch.mean(torch.norm(pred_pos[is_nonlinear] - pred_traj[is_nonlinear], dim-1)) metrics[nl_ade] nl_ade.item() for k in metrics: metrics[k] / len(dataloader) return metrics典型训练结果对比模型ADE (m)FDE (m)NL-ADE (m)线性模型1.332.941.89普通LSTM0.982.121.42Social LSTM0.611.240.836. 实际应用与优化建议将训练好的Social LSTM模型部署到实际系统中时还需要考虑以下优化方向实时性优化使用TensorRT加速推理实现增量式预测避免每帧重新计算多模态预测同时生成多条可能轨迹为每条轨迹分配概率场景融合结合场景语义信息障碍物、出入口等考虑行人姿态和视线方向def predict_real_time(model, current_obs, past_statesNone): 实时预测接口 if past_states is None: past_states initialize_states() # 处理当前观测 processed_obs preprocess(current_obs) # 预测未来轨迹 with torch.no_grad(): pred_traj, new_states model(processed_obs, past_states) return pred_traj, new_states在实际机器人导航系统中集成Social LSTM时发现将预测模块与路径规划器松耦合比端到端方案更稳定。具体实现中我们维护一个预测结果缓存区每100ms更新一次预测而规划器则以更高频率(如10ms)重新规划路径。