STGNN交通流预测实战：从数据集预处理到模型训练完整指南（PyTorch版）

STGNN交通流预测实战：从数据集预处理到模型训练完整指南（PyTorch版）

交通流预测一直是智慧城市建设的核心挑战之一。想象一下，当你早晨打开导航app，它能准确预测未来一小时的交通拥堵情况——这背后很可能就运行着类似STGNN（时空图神经网络）的算法。不同于传统时序预测模型，STGNN能同时捕捉路网的空间拓扑关系和交通流的时间动态特性，在真实场景中往往能实现更精准的预测。本文将带你完整实现一个基于PyTorch的STGNN交通流预测项目，从原始数据清洗到模型调参，手把手解决每个环节可能遇到的"坑"。

1. 环境准备与数据获取

1.1 基础环境配置

推荐使用conda创建隔离的Python环境（3.8+版本），核心依赖包括：

conda create -n stgnn python=3.8 conda activate stgnn pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113 pip install numpy pandas scipy h5py scikit-learn tqdm

提示：CUDA版本需与本地GPU驱动兼容，可通过nvidia-smi查询最高支持的CUDA版本

1.2 数据集下载与解析

METR-LA和PEMS-BAY是交通预测领域的基准数据集，其结构对比如下：

下载后原始数据结构示例：

with h5py.File('metr-la.h5', 'r') as f: print(f['df'].keys()) # 输出：['axis0', 'axis1', 'block0_items', 'block0_values'] data = f['df']['block0_values'][:] # 形状(时间步长, 传感器数)

2. 数据预处理实战

2.1 时空数据标准化

交通流数据通常需要两种归一化处理：

1. 空间维度标准化：对每个传感器的历史流量单独进行z-score标准化

   def z_score_normalize(data): mean = np.nanmean(data, axis=0) std = np.nanstd(data, axis=0) return (data - mean) / std, mean, std

2. 时间维度填充：采用线性插值处理缺失值

   from scipy.interpolate import interp1d def temporal_interpolation(data): x = np.arange(data.shape[0]) for i in range(data.shape[1]): mask = ~np.isnan(data[:, i]) f = interp1d(x[mask], data[mask, i], kind='linear', fill_value="extrapolate") data[:, i] = f(x) return data

2.2 图结构构建

路网拓扑关系通过邻接矩阵表达，常用阈值高斯核函数计算空间相关性：

def build_adjacency(coords, threshold=0.1): """ coords: (N, 2) 的经纬度坐标数组 threshold: 相关性阈值，小于该值的边将被丢弃 """ dist_mx = np.zeros((len(coords), len(coords))) for i in range(len(coords)): for j in range(len(coords)): dist = np.linalg.norm(coords[i] - coords[j]) dist_mx[i][j] = np.exp(-dist**2 / 0.1) # 高斯核 dist_mx[dist_mx < threshold] = 0 return dist_mx

注意：实际项目中建议使用路网真实连接关系（如有）替代几何距离

3. STGNN模型深度定制

3.1 模型架构关键修改

原始STGNN代码常需调整以下核心参数：

class STGNN(nn.Module): def __init__(self, infea, outfea, L, d, P, Q): super(STGNN, self).__init__() # 时空块堆叠层数 self.L = L # 隐藏层维度 self.d = d # 历史时间步长 self.P = P # 预测时间步长 self.Q = Q self.st_blocks = nn.ModuleList([ STBlock(infea if i==0 else d, d, P) for i in range(L) ]) self.output = nn.Linear(d, outfea)

典型参数配置经验值：

3.2 多GPU训练适配

修改训练脚本支持DataParallel：

if torch.cuda.device_count() > 1: print(f"Using {torch.cuda.device_count()} GPUs") model = nn.DataParallel(model) model.to(device)

需特别注意：

• 自定义Module的forward参数需保持张量在相同设备
• 损失函数计算需在主GPU上聚合

4. 训练优化与错误排查

4.1 学习率调度策略

采用warmup+余弦退火组合策略：

from torch.optim.lr_scheduler import CosineAnnealingLR optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) warmup_epochs = 10 scheduler = CosineAnnealingLR(optimizer, T_max=max_epoch-warmup_epochs) for epoch in range(max_epoch): if epoch < warmup_epochs: lr = base_lr * (epoch + 1) / warmup_epochs for param_group in optimizer.param_groups: param_group['lr'] = lr else: scheduler.step()

4.2 常见错误解决方案

错误1：维度不匹配导致GRU报错

RuntimeError: input.size(-1) must be equal to input_size

解决方法： 检查STBlock中GRU的input_size与前一层的输出维度是否一致，特别是修改了d参数后需要同步更新：

self.gru = nn.GRU(input_size=d, hidden_size=d)

错误2：邻接矩阵NaN值

ValueError: Input contains NaN

解决方法： 在构建邻接矩阵后添加校验：

adj = np.nan_to_num(adj, nan=0.0) adj = torch.FloatTensor(adj)

5. 模型评估与结果可视化

5.1 多指标评估体系

实现完整的评估流程：

def evaluate(true, pred): # 反归一化 true = scaler.inverse_transform(true) pred = scaler.inverse_transform(pred) metrics = { 'MAE': np.mean(np.abs(true - pred)), 'RMSE': np.sqrt(np.mean((true - pred)**2)), 'MAPE': np.mean(np.abs((true - pred)/true)) * 100 } return metrics

典型baseline对比结果（METR-LA数据集）：

5.2 预测结果可视化

使用Matplotlib动态展示预测效果：

import matplotlib.animation as animation fig, ax = plt.subplots() line_true, = ax.plot([], [], 'r-', label='True') line_pred, = ax.plot([], [], 'b--', label='Pred') def animate(i): ax.set_title(f'Time step {i}') line_true.set_data(range(207), true[i]) line_pred.set_data(range(207), pred[i]) return line_true, line_pred ani = animation.FuncAnimation(fig, animate, frames=24, interval=200) plt.legend() plt.show()

在实际项目中，我发现调整P（历史时间步长）和Q（预测步长）的比例对结果影响显著——当P/Q≈2时模型表现最佳。例如对于15分钟预测（Q=3），使用30分钟历史数据（P=6）比单纯增加网络深度更有效。