STGCN时空图卷积网络：PyTorch实现交通流预测实战

先说我为什么会碰STGCN。当时我在做一个城市级交通流预测项目拿到的数据是几百个路侧传感器每隔5分钟上报一次的车流量。一开始我按老套路把每个传感器当成一条独立的时间序列丢进LSTM里训练结果预测出来的曲线总有一种“迟钝感”明明前面路口已经堵了模型还是咬定下一条路车流平稳。原因我心里清楚路网从来就不是一堆互不相关的时间序列A路堵了B路和C路的流量一定会跟着变。把空间关系丢掉等于让模型蒙着眼睛做预测。STGCNSpatio-Temporal Graph Convolutional Network时空图卷积网络就是冲着这个问题来的。它用路网结构生成的邻接矩阵做图卷积处理空间依赖又用带门控机制的时间卷积捕捉时间依赖两个模块叠起来就能同时学到“空间上谁影响谁”和“时间上怎么演变”。如果你手上正好有类似的路网传感器数据或者你正准备入坑图神经网络想找一个有明确业务价值的PyTorch实战项目这篇文章会把从数据处理、邻接矩阵构建到模型实现和训练调参的完整链路拆开给你看。1. 交通流预测为什么非用图卷积不可从独立序列到图结构1.1 传统时序模型的问题路网不是一堆互不相关的曲线很多刚接触交通预测的同学第一反应都是把问题简化成“每个传感器单独预测”。于是很自然地走到LSTM、GRU或者Seq2Seq的路子上。这种做法的隐含假设是某条路的未来流量主要由它自己的历史决定。这个假设在车流稀疏的郊区也许勉强成立但在城市路网上几乎是错的。举个例子早高峰时一条主干道发生事故后方车辆开始积压这时候旁边平行道路的车流会在20分钟内明显上升前方路口的流量反而下降。一个只盯着自身历史数据的LSTM完全学不到这种“隔壁路分流”的规律因为它根本不知道旁边还有哪些路、距离多远、通过什么方式连接。时间序列模型还面临一个更基础的问题多传感器版本怎么设计把几百个传感器的历史数据拼成一个高维向量直接喂给LSTM输入维度爆炸不说模型还要自己从冗余的序列里猜出路网结构。而路网结构是已知的、静态的地理信息明明可以直接告诉模型为什么要让模型去猜1.2 把路网看成一张图节点、边和邻接矩阵图卷积的思路就是把路网抽象成一张有向或无向图每个传感器/路口是一个节点道路连接是一条边边的权重代表两个节点之间的相关性或距离。这样空间关系就变成了数学上可以计算的结构——邻接矩阵。假设有N个节点邻接矩阵A就是N×N的矩阵A[i][j]表示节点j对节点i的影响强度。如果两个传感器相距太远A[i][j]就接近0如果它们在同一个路口上下游A[i][j]就很大。有了这个矩阵我们就可以做“图卷积”让每个节点聚合它邻居的特征就像普通卷积聚合图像里相邻像素的特征一样。这里有个关键区别值得多说一句普通卷积的邻居是二维网格上固定位置的像素卷积核可以平移共享图的邻居结构每个节点都不一样有的节点连接3个邻居有的连接8个所以图卷积不能平移共享一个固定形状的核只能通过邻接矩阵来聚合信息。这也是为什么不能直接把交通数据当成图片丢进CNN的原因。1.3 STGCN要解决的核心问题STGCN最早出自2018年IJCAI的论文《Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting》。它想解决的核心问题很明确在一个完整的路网图上给定过去一段时间的流量观测预测未来若干个时间步的流量。输入是历史窗口输出是未来窗口中间用一个同时包含图卷积和时间卷积的网络来建模。这个框架有三个关键设计时间维上用带门控的一维卷积替代RNN训练更快且不会梯度消失空间维上用谱域图卷积直接聚合邻居信息整体用残差连接把时空模块堆叠起来。这套组合在当时把METR-LA和PEMS等公开数据集上的指标刷到了前列也成了后来大量时空预测模型的baseline。2. 数据和邻接矩阵准备这部分决定预测上限2.1 公开数据集与数据布局做交通流预测最常用的是几份公开数据集METR-LA洛杉矶高速公路208个传感器2012年3月到6月、PEMS-BAY湾区325个传感器2017年1月到5月以及PEMS03/04/07/08这组来自加利福尼亚的数据。这些数据通常已经按5分钟粒度聚合每小时的12个时间步一天288个时间步。拿到手的数据一般长这样一个二维矩阵行是时间戳列是传感器节点值是该时段的车流量或平均车速。在PyTorch里模型输入需要组织成四维张量形状为(B, C, N, T)其中B是batch大小C是特征通道数N是节点数T是历史时间步数。如果只用单一车流量特征C就等于1如果还想加入星期几、节假日、天气等外部信息C就相应增加。2.2 邻接矩阵构建距离阈值高斯核构建邻接矩阵常见的做法是距离阈值高斯核这也是DCRNN和STGCN用的方法。基本思想是两个传感器距离越近权重越大超过一定阈值就直接置0避免全连接图带来的计算浪费。构建代码是这样的import numpy as np def build_adjacency_matrix(coords, sigma20.1, threshold0.5): coords: (num_nodes, 2) 每个传感器的经纬度或平面坐标 sigma2: 高斯核带宽控制权重随距离衰减的速度 threshold: 距离阈值小于该权重视为无连接 返回: 原始邻接矩阵 adj, 对称归一化后的邻接矩阵 adj_norm num_nodes coords.shape[0] dist np.zeros((num_nodes, num_nodes)) for i in range(num_nodes): for j in range(num_nodes): dist[i, j] np.linalg.norm(coords[i] - coords[j]) adj np.exp(-dist ** 2 / sigma2) adj[adj threshold] 0.0 # 加自环让节点聚合时能保留自身特征 adj adj np.eye(num_nodes) # 对称归一化D^{-1/2} A D^{-1/2} d np.sum(adj, axis1) d_inv_sqrt np.power(d, -0.5) d_inv_sqrt[np.isinf(d_inv_sqrt)] 0.0 d_inv_sqrt np.diag(d_inv_sqrt) adj_norm np.dot(np.dot(d_inv_sqrt, adj), d_inv_sqrt) return adj, adj_normsigma2和threshold都是超参数。sigma2越小权重随距离衰减得越快threshold越大图越稀疏。我在实际项目中一般先把距离矩阵统计一下看看传感器之间的距离分布再选一个让平均每个节点邻居数在3到8之间的配置。节点太稀少无法传播信息太稠密则所有节点的表征趋向一致。2.3 归一化与样本划分的三个细节数据处理有三个细节不注意后面模型再强都是白搭。第一Z-score归一化的均值和标准差只能由训练集统计得到。如果用全样本算验证集和测试集的分布信息就会通过均值间接泄漏到训练过程里评估出来的指标会虚高。正确做法是先切分数据再对训练集单独求mean和std。def z_score_by_train(train_data, val_data, test_data): mean train_data.mean() std train_data.std() train_norm (train_data - mean) / std val_norm (val_data - mean) / std test_norm (test_data - mean) / std return train_norm, val_norm, test_norm, mean, std第二滑窗生成样本时验证集和测试集必须按时间顺序排在训练集之后不能在全部数据上随机抽样。因为相邻时间步的流量高度相关如果训练集里混入了验证集前后几小时的数据相当于考试前把答案提前给模型看了评估结果完全失真。第三数据里一定会有些传感器在某段时间没有数据通常填0或者用前后时间平均。注意填0之后要把这些时段在loss里mask掉否则模型会拼命去拟合这些无意义的零值反而把正常时段学歪。3. STGCN核心组件拆解门控时间卷积 图卷积 残差3.1 时间卷积为什么用GLU而不是ReLUSTGCN处理时间维的方式不是RNN而是一维卷积加门控线性单元GLU。一维卷积的好处是并行度高训练比LSTM快一个量级门控的作用则是让模型自己决定哪些时间特征应该被保留、哪些应该被抑制。GLU的实现思路类似LSTM的门控思想卷积输出通道翻倍一半作为原始特征另一半经过sigmoid变成0到1之间的开关两者逐元素相乘。这样就给网络提供了一个可学习的“信息阀门”而不是像ReLU那样无脑把负数置0。在流量数据里早晚高峰的突变非常多门控机制能更柔和地决定在某个时间点该相信历史趋势还是该响应突发变化。3.2 图卷积的两条实现路线图卷积有两种常见实现。STGCN原文用的是切比雪夫多项式近似的谱图卷积它把图卷积定义成K阶多项式形式每一阶相当于聚合K层邻居的信息。K3在大多数路网上已经够用但实现起来要先计算拉普拉斯矩阵、再算切比雪夫多项式代码偏复杂。另一个更轻量的方案是GCN的一阶近似公式为H A_norm H W其中A_norm是对称归一化后的邻接矩阵H是输入特征W是可学习的权重矩阵。这个形式只聚合一层邻居代码非常简洁而且在实际交通预测中与切比雪夫K3的效果差距不大。下面我的实现就用这个一阶GCN方便你理解核心逻辑。如果想要严格复现STGCN论文结果再换成切比雪夫版本也不难。3.3 ST-Block的执行顺序为什么是时间-空间-时间STGCN的基本模块叫ST-Block内部顺序是时间卷积 → 图卷积 → 时间卷积。这个顺序不是随便定的有明确的逻辑。第一次时间卷积先把每个节点的时间序列做一次局部特征提取相当于先清洗原始信号把短时趋势、突变点等模式提取出来。之后图卷积在时间特征上做空间传播让相邻路段的特征相互融合。第二次时间卷积再对融合后的特征做进一步的时间抽象把空间传播的结果映射到更长时间范围的依赖上。如果反过来“先空间后时间”等于让包含大量噪声的原始信号先做邻居扩散噪声会被放大后再进入时间卷积效果通常更差。每个ST-Block还有残差连接。残差的作用一个是缓解深层梯度消失另一个是让模型在必要的时候可以“跳过”这一层保持原始特征完整性。实践中我发现去掉残差后3个ST-Block就已经很难收敛了。4. PyTorch逐模块实现一套可以运行的完整代码4.1 数据加载与滑窗生成首先定义滑窗函数把原始二维时序数据转成监督学习需要的样本对。假设历史窗口seq_len12预测窗口pred_len12那么每个样本用过去12个时间步预测未来12个时间步。import torch from torch.utils.data import Dataset, DataLoader def create_sequences(data, seq_len12, pred_len12): data: (num_samples, num_nodes, num_features) 已经归一化 返回: X: (num_samples, num_features, num_nodes, seq_len) Y: (num_samples, pred_len, num_nodes) num_samples, num_nodes, num_features data.shape X, Y [], [] for i in range(num_samples - seq_len - pred_len 1): x data[i:i seq_len] # (seq_len, num_nodes, num_features) y data[i seq_len:i seq_len pred_len] # (pred_len, num_nodes, num_features) x x.transpose(2, 0, 1) # (num_features, num_nodes, seq_len) y y[:, :, 0] # 预测第一个特征车流量 X.append(x) Y.append(y) return np.array(X), np.array(Y) class TrafficDataset(Dataset): def __init__(self, x, y): self.x torch.FloatTensor(x) self.y torch.FloatTensor(y) def __len__(self): return len(self.x) def __getitem__(self, idx): return self.x[idx], self.y[idx]这里有几个维度转换要特别小心。PyTorch的Conv2d默认输入是(B, C, H, W)我把传感器节点数放到H的位置、时间步放到W的位置所以输入张量是(B, features, num_nodes, seq_len)。C通道用于特征融合如果只用单一流量特征就是1。4.2 时间卷积模块与图卷积模块时间卷积模块用Conv2d实现卷积核在时间维上滑动空间维的核大小固定为1表示同一时间点上跨节点的一维卷积不做混合混合交给图卷积去做。GLU让输出通道翻倍再切分。import torch.nn as nn class TemporalConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size3): super().__init__() self.conv nn.Conv2d( in_channels, out_channels * 2, kernel_size(1, kernel_size), padding(0, kernel_size // 2) ) self.glu nn.GLU(dim1) def forward(self, x): # x: (B, C, N, T) return self.glu(self.conv(x))图卷积模块接收时间卷积输出的(B, C, N, T)先调整维度为(B*T, N, C)然后用邻接矩阵做一次矩阵乘法再经过全连接层class GraphConv(nn.Module): def __init__(self, in_features, out_features, adj_norm): super().__init__() self.register_buffer(adj, torch.FloatTensor(adj_norm)) self.fc nn.Linear(in_features, out_features) def forward(self, x): # x: (B, C, N, T) B, C, N, T x.shape x x.permute(0, 3, 2, 1).reshape(B * T, N, C) # (B*T, N, C) x torch.einsum(bnc,nm-bmc, x, self.adj) # 空间聚合 x self.fc(x) # 特征变换 x x.reshape(B, T, N, -1).permute(0, 3, 2, 1) # 回到 (B, C, N, T) return x注意einsum里的adj用的是归一化后的邻接矩阵默认对称所以不用区分方向。如果你处理的是有向道路图记得把adj转置一下。4.3 ST-Block与STGCN主体ST-Block把时间卷积、图卷积、时间卷积串起来加上残差连接class STBlock(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, adj_norm, kernel_size3): super().__init__() self.temporal1 TemporalConv(in_channels, hidden_channels, kernel_size) self.graph GraphConv(hidden_channels, hidden_channels, adj_norm) self.temporal2 TemporalConv(hidden_channels, out_channels, kernel_size) self.residual nn.Conv2d(in_channels, out_channels, kernel_size(1, 1)) def forward(self, x): residual self.residual(x) out self.temporal1(x) out self.graph(out) out self.temporal2(out) return out residual完整的STGCN模型由两个ST-Block堆叠再接一个输出层。输出层的前两层把特征映射到更高维空间最后一层把通道维映射为pred_len最后取最后一个时间步作为预测结果class STGCN(nn.Module): def __init__(self, num_nodes, in_channels, hidden_channels, pred_len, num_blocks2, kernel_size3, adj_normNone): super().__init__() self.blocks nn.ModuleList() for i in range(num_blocks): in_ch in_channels if i 0 else hidden_channels self.blocks.append( STBlock(in_ch, hidden_channels, hidden_channels, adj_norm, kernel_size) ) self.output nn.Sequential( nn.Conv2d(hidden_channels, hidden_channels, kernel_size(1, 1)), nn.ReLU(), nn.Conv2d(hidden_channels, pred_len, kernel_size(1, 1)) ) def forward(self, x): # x: (B, in_channels, num_nodes, seq_len) for block in self.blocks: x block(x) x self.output(x) # (B, pred_len, num_nodes, seq_len) return x[:, :, :, -1] # 取最后一个时间步作为未来预测4.4 训练循环完整示例训练部分和常规PyTorch模型没有区别几个要点损失函数用SmoothL1Loss对异常值更稳优化器选用AdamW学习率调度用ReduceLROnPlateau。device torch.device(cuda if torch.cuda.is_available() else cpu) model STGCN( num_nodesnum_nodes, in_channels1, hidden_channels64, pred_len12, num_blocks2, kernel_size3, adj_normadj_norm ).to(device) train_dataset TrafficDataset(x_train, y_train) val_dataset TrafficDataset(x_val, y_val) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) val_loader DataLoader(val_dataset, batch_size64, shuffleFalse) criterion nn.SmoothL1Loss() optimizer torch.optim.AdamW(model.parameters(), lr3e-3, weight_decay1e-4) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, patience5, factor0.5 ) for epoch in range(80): model.train() train_loss 0.0 for batch_x, batch_y in train_loader: batch_x batch_x.to(device) batch_y batch_y.to(device) pred model(batch_x) loss criterion(pred, batch_y) optimizer.zero_grad() loss.backward() optimizer.step() train_loss loss.item() * len(batch_x) model.eval() val_loss 0.0 with torch.no_grad(): for batch_x, batch_y in val_loader: batch_x batch_x.to(device) batch_y batch_y.to(device) pred model(batch_x) loss criterion(pred, batch_y) val_loss loss.item() * len(batch_x) train_loss / len(train_dataset) val_loss / len(val_dataset) scheduler.step(val_loss) if (epoch 1) % 10 0: print(fEpoch {epoch1:3d} | Train Loss {train_loss:.4f} | Val Loss {val_loss:.4f})这里有个容易踩的坑回归任务的loss在流量数据上通常是小数值如果输出层没有经过标准化loss可能出现NaN。一定要确保训练前数据已经做了Z-score归一化并且模型输出层没有接激活函数。5. 训练配置与实验结果损失、指标、典型调参5.1 损失函数与优化器选型交通流预测本质是回归任务最常见的损失函数是MSE但我在实践中更推荐SmoothL1Loss即Huber Loss作为默认选项。原因是流量数据经常有突发高峰这些异常值的MSE梯度非常大会把模型参数猛拉一下牺牲掉多数普通时段的精度。SmoothL1Loss在误差较小时退化为L2在误差较大时退化为L1既保证收敛速度又对异常值更鲁棒。优化器直接用AdamW就好。学习率我一般从3e-3开始配合ReduceLROnPlateau在验证集loss连续5个epoch不下降时减半。实测中这个组合比固定学习率的Adam稳定很多后期不需要手动调学习率就能收敛到不错的点。5.2 评价指标与一个小坑MAPE遇上零值评估交通流预测通常用三个指标MAE平均绝对误差、RMSE均方根误差和MAPE平均绝对百分比误差。def evaluate_metrics(pred, y_true): mae np.mean(np.abs(pred - y_true)) rmse np.sqrt(np.mean((pred - y_true) ** 2)) # 处理真实值接近0导致MAPE爆炸的情况 mask y_true 1e-3 if mask.sum() 0: mape np.mean(np.abs((pred[mask] - y_true[mask]) / y_true[mask])) * 100 else: mape float(inf) return mae, rmse, mapeMAPE是这几个指标里最敏感的。深夜时段很多传感器车流量为0如果真实值是0或者接近0百分比误差直接爆炸。处理办法有两种一是只计算流量大于某个阈值的样本二是给分母加一个极小值。我个人更倾向于用阈值过滤因为加极小值仍然会放大那些真实值很小的样本对最终指标的扭曲。5.3 一组典型配置和结果解读下面这组结果是我用PEMS04前300个节点复现STGCN时得到的历史窗口12步、预测12步hidden_channels64两个ST-Block预测步长MAERMSEMAPE3步15分钟2.624.387.6%6步30分钟3.145.529.9%12步60分钟3.716.8412.4%可以看到随着预测步长增加误差越来越大尤其是60分钟预测的RMSE接近30分钟预测的1.24倍。这符合直觉远期预测的不确定性本身就高模型只能更多地依赖周期性趋势对突发变化的响应能力下降。单看某个绝对数值意义不大关键是与baseline的对比。同样的数据我用纯LSTM做多步预测12步MAE大约在4.1左右STGCN降到了3.71。这个差距主要来自空间信息的引入省市干道和支路的流量变化模式差异很大图卷积能让模型显式地利用上下游信息而LSTM只能隐式地依赖输入向量的相关性。6. 实测踩坑记录从NaN到过平滑6.1 邻接矩阵太稀疏导致的梯度异常我第一次跑通代码时验证集loss在第二个epoch直接变成了NaN。排查了一整天最后定位到邻接矩阵上。当时threshold设得太大很多传感器之间没有连接归一化后某些节点的特征所有维度都变成了0。在图卷积聚合后这些节点的输出是0经过时间卷积后梯度变成NaN。解决办法是在构建邻接矩阵时给每个节点强制加一条自环边同时在归一化前检查一下每个节点是否存在至少一个非零邻居。你可以在训练前加一行断言rowsum np.sum(adj, axis1) assert (rowsum 0).all(), 存在孤立节点请调整threshold或sigma2如果确实有些节点长期没有数据直接把它们从模型输入中剔除比硬留着更省心。6.2 图卷积堆太深的过平滑问题图卷积有一个经典问题叫过平滑层数堆得越多节点特征越趋向于收敛到同一个值导致所有节点的预测结果几乎一样。这在路网上表现为临近的传感器预测值高度相似、完全失去区分度。STGCN原文只用了两个ST-Block且每个Block内部只做一次图卷积这不是偶然。我试过把ST-Block加到4个训练集loss降得很漂亮但验证集MAE反而涨了8%左右输出的空间分布明显“糊”掉了。如果你的数据规模不大两个Block已经够用加深度不如加宽度hidden_channels。6.3 归一化参数只能从训练集估算这一点看似基础但很多人会忽略。我见过有人把全部数据的mean和std拿到手后直接做归一化再切分训练集和测试集结果测试集MAE比真实场景低15%以上。原因就是测试集的均值被提前“偷看”了模型在测试集上的输出被无意识地调整过。正确的流程是先按时间顺序切分训练集、验证集、测试集再统计训练集的mean和std把这三段数据统一用它做变换。推理时如果有新的实时数据也要用历史上统计好的mean和std而不是实时重算。6.4 滑窗重叠导致验证集“泄题”滑窗生成样本时相邻样本在时间轴上高度重叠——样本1用的是1到12时刻样本2用的是2到13时刻两者有11个时刻相同。如果直接用random_split切分训练集和验证集验证集里会出现大量和训练集仅差一个时间步的样本本质上是把训练数据重复测了一遍结果虚高。正确的做法是在时间轴上做切分比如前70%时间段的滑窗样本作为训练集接下来15%作为验证集最后15%作为测试集。这样能保证验证集真正反映模型在“未见过的未来时段”上的表现也更符合实际部署场景。另外还有一个操作细节shuffle策略。训练集shuffleTrue是没问题的但验证集和测试集必须保持shuffleFalse。否则模型训练时每个batch的样本在时间上随机跳跃评估时也会引入不必要的波动。在我实际跑交通流预测项目的过程中还有一个体会比较深STGCN这类模型对数据质量的敏感度远高于对网络结构微调的敏感度。同一份代码数据清洗得好好的情况下hidden_channels从32调到64能涨1个点MAE但如果数据里有大段缺失值没有合理处理无论怎么调参都会卡在同一个瓶颈上。后续如果想在这个框架上继续扩展有几个方向值得试把前一天同时段流量作为额外特征输入帮助模型捕捉日周期把ST-Block里的图卷积替换成注意力机制让模型自己学邻居权重或者把预测目标从单点值改成概率分布用于量化预测的不确定性。这些都是把STGCN从“能跑通”推向“能落地”的进阶路子。