简介该资源是PyTorch复现STGCN的配套学习资料面向具备一定深度学习基础、从事交通流预测或时空数据建模的研究者与工程师。内容围绕IJCAI 2018论文《Spatio-Temporal Graph Convolutional Networks》展开系统梳理近似谱图卷积、GLU门控一维因果卷积、时空卷积块ST-Conv Block的瓶颈结构与归一化设计并说明PeMSD7(M)数据集上228个顶点、以速度作为单一特征的图数据构建方式以及模型输出下一时刻路网速度图的完整流程。压缩包共1个PDF文件约408KB篇幅精炼便于集中阅读论文思路与复现代码要点。目前已有4486人学习适合作为交通速度预测入门、图卷积时空建模原理理解与PyTorch实现对照的参考材料也可用于梳理GCN、STGCN在交通研究中的应用脉络与常见挑战。1. 为什么交通速度预测需要 STGCN 而不是单路 LSTM交通速度预测不是把一条时间序列往后平移。路网里相邻路段的拥堵会沿拓扑传播早高峰一个匝道限流三公里外的速度也可能在十五分钟后掉下来。单路 LSTM 或 Transformer 能把时间依赖拟合得很好却常把每个路段当成独立样本切断了空间依赖。STGCN 的做法是把路网看作图节点是传感器或路段边权是距离或连通性用图卷积聚合邻居信息再用门控时间卷积抓局部时变模式。Pytorch 复现 STGCN 的价值在于你能把图卷积、时间卷积、残差连接逐层拆开改邻接矩阵、改切比雪夫阶数、换多步输出都能在训练循环里看到损失和 MAE 的变化。适合已经会写 Pytorch Dataset 和训练循环、想把交通速度预测从单序列推到路网时空建模的人。2. 从图卷积到时空块STGCN 的结构拆解与 Pytorch 张量准备STGCN 的复现难点不在反向传播而在维度。交通数据同时有路网节点、时间步和特征通道三层结构图卷积最怕邻接矩阵对错节点时间卷积最怕把批次维和时间维混在一起。先把手算的图信号变成 Pytorch 张量再谈堆 block。2.1 路网图与交通速度张量邻接矩阵和特征矩阵怎么对齐假设有 N 个传感器每个传感器记录速度、流量、占有率三个通道采样间隔 5 分钟。原始数据可以整理成(T, N, C)T 是总时间步C 通常取 1 或 3。邻接矩阵A是(N, N)第 i 行第 j 列表示路段 i 和 j 的连通强度。常见做法是用距离的高斯核构造距离越近权重越大超过阈值直接置零避免远距离节点互相干扰。import numpy as np import torch def gaussian_adj(distance, sigma0.1, threshold0.5): # distance: (N, N) 节点间距离对角线为 0 w np.exp(-np.square(distance) / sigma**2) w[w threshold] 0.0 np.fill_diagonal(w, 0.0) return w.astype(np.float32) def normalized_laplacian(adj): # adj: (N, N) 非负邻接矩阵 d adj.sum(axis1) d_inv_sqrt np.power(d, -0.5, whered 0) d_inv_sqrt[np.isinf(d_inv_sqrt)] 0.0 D_inv_sqrt np.diag(d_inv_sqrt) I np.eye(adj.shape[0], dtypenp.float32) L I - D_inv_sqrt adj D_inv_sqrt return L.astype(np.float32) adj gaussian_adj(distance, sigma0.1, threshold0.5) L normalized_laplacian(adj)sigma控制权重衰减速度太小会让邻接矩阵接近单位阵太大则所有节点都变邻居。threshold用来切断弱连接通常先看距离分布再选 0.1 到 0.5 之间的值。归一化拉普拉斯L是后面切比雪夫近似的输入它让图卷积在数值上稳定。注意对角线置零不要把节点自身算进邻居权重。张量形状含义原始速度矩阵(T, N, C)T 个时间步N 个节点C 个特征邻接矩阵(N, N)节点间空间权重归一化拉普拉斯(N, N)图卷积的传播算子训练输入(B, C_in, T_in, N)批次、输入通道、历史步、节点训练标签(B, T_out, N)批次、预测步、节点2.2 图卷积神经网络在 STGCN 里的切比雪夫近似图卷积神经网络在谱域的定义依赖拉普拉斯矩阵的特征分解但特征分解是 O(N^3)大路网不划算。STGCN 采用切比雪夫多项式近似把卷积核限制在 K 阶邻域内。公式上T_0(L)IT_1(L)LT_k(L)2L T_{k-1}(L)-T_{k-2}(L)。K 通常取 2 到 4对应每个节点最多看到 K 跳邻居。def cheb_polynomial(L, K): # L: (N, N) 归一化拉普拉斯 N L.shape[0] cheb [np.eye(N, dtypenp.float32), L.astype(np.float32)] for k in range(2, K): cheb.append(2 * L cheb[-1] - cheb[-2]) return cheb[:K] # 返回 K 个 (N, N) 矩阵 cheb_polys cheb_polynomial(L, K3)K越大感受野越广但参数量和过拟合风险也上升。交通速度预测里 K2 或 K3 往往够用因为三四跳以外的节点影响已经很小。切比雪夫多项式在训练前一次性算好作为固定矩阵传给模型不参与梯度更新。如果路网动态变化再考虑自适应图卷积而不是把 K 堆到 8。2.3 时空卷积块门控时间卷积、空间图卷积与残差一个 STGCN 块可以记成“时间卷积 → 空间图卷积 → 时间卷积”外面套残差。时间卷积用门控线性单元把通道分成两半一半做线性输出一半过 sigmoid 做门。空间图卷积就是 2.2 的切比雪夫卷积。残差负责把输入通道映射到输出通道避免堆叠后梯度消失。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, 2 * out_channels, kernel_size(kernel_size, 1), padding(kernel_size // 2, 0)) def forward(self, x): # x: (B, C_in, T, N) h self.conv(x) a, b h.chunk(2, dim1) return a * torch.sigmoid(b)kernel_size控制时间卷积看几个相邻时刻取 3 比较平衡。padding保持时间长度不变方便残差相加。空间图卷积前要把(B, C, T, N)变成(B*T, C, N)因为图卷积在节点维聚合时间维可以并入批次。这个变换顺序写错模型不会报维度错误但学到的全是错节点关系。2.4 Pytorch 张量基础维度变换和批处理Dataset 返回的样本通常是(N, C, T)DataLoader 堆成(B, N, C, T)。模型内部要统一成(B, C, T, N)图卷积再临时拆成(B*T, C, N)。用permute而不是view因为通道和时间需要交换内存不连续时view会报错。# 假设 x 来自 DataLoader形状 (B, N, C, T) x x.permute(0, 2, 3, 1).contiguous() # (B, C, T, N) # 进入图卷积前 B, C, T, N x.shape x_graph x.permute(0, 2, 1, 3).reshape(B * T, C, N) # 图卷积后恢复 x x_graph.reshape(B, T, C, N).permute(0, 2, 1, 3)contiguous()在 permute 后调用保证后面的 reshape 或 view 不会因为内存步长报错。批大小 B 可以变T 和 N 必须和预计算的切比雪夫矩阵一致。如果 N 对不上先检查邻接矩阵是不是用了另一个数据集的节点数。3. 用 PyTorch 搭出 STGCN环境、数据管道与训练循环环境、数据、模型、训练循环这四件事分开做排错会快很多。Anaconda 配置 PyTorch 环境适合反复切换 CPU 和 GPU 版本交通速度数据用滑动窗口生成样本模型把第 2 章的组件拼起来训练循环里注意损失和标签的形状对齐。3.1 Anaconda 配置 PyTorch 环境与 CUDA 12.0 版本选择PyTorch 安装最稳妥的方式是去官网看命令但本地用 Anaconda 管理环境依然方便。CUDA 12.0 对应的 PyTorch 版本通常选 2.1 及以上具体命令以官网为准。如果只是先在 CPU 上验证模型结构装 CPU 版就够不必和显卡驱动较劲。conda create -n stgcn python3.10 -y conda activate stgcn # 有 GPU 时按官网选择对应 CUDA 版本的命令例如 # conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia # 只跑 CPU 时 conda install pytorch torchvision torchaudio cpuonly -c pytorch python -c import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())最后一行用来确认三件事PyTorch 版本、CUDA 版本、显卡是否可用。如果输出None或False先别改模型回到环境排查。Windows 上如果出现c10.dll初始化失败常见原因是 CUDA 版和驱动不匹配或者多个环境混用。建一个干净环境重装对应版本比在原环境里反复卸载更省时间。3.2 交通速度数据集滑动窗口与 Dataset 实现交通速度数据先按通道做 z-score 归一化。用历史 12 步预测未来 3 步即 1 小时输入、15 分钟输出。切割时不要让训练集和测试集共享窗口否则测试集会泄漏训练集末尾的信息。import torch from torch.utils.data import Dataset class TrafficWindowDataset(Dataset): def __init__(self, data, input_len12, pred_len3): # data: (T, N, C) 已归一化 self.data torch.as_tensor(data, dtypetorch.float32) self.input_len input_len self.pred_len pred_len def __len__(self): return len(self.data) - self.input_len - self.pred_len 1 def __getitem__(self, idx): x self.data[idx: idx self.input_len] # (T_in, N, C) y self.data[idx self.input_len: idx self.input_len self.pred_len] # (T_out, N, C) x x.permute(1, 2, 0) # (N, C, T_in) y y.permute(1, 2, 0) # (N, C, T_out) return x, yinput_len和pred_len是最直接的时间窗口参数。输入越长模型能看到更早的趋势但也会稀释近期变化。交通速度预测里 12 步是常见起点预测 3 或 6 步。如果速度序列有缺失值归一化前先插值或掩码不要用 0 填充后直接算 MAPE。3.3 STGCN 模型代码ChebConv、TemporalConv、STGCNBlock把 2.2 和 2.3 的组件写成模块。ChebConv 的权重形状是(K, C_in, C_out)对每个切比雪夫项分别做节点聚合再求和。注意cheb_polys要放在模型同一设备上通常注册成 buffer。import torch import torch.nn as nn class ChebConv(nn.Module): def __init__(self, in_channels, out_channels, K): super().__init__() self.K K self.weight nn.Parameter(torch.empty(K, in_channels, out_channels)) self.bias nn.Parameter(torch.zeros(out_channels)) nn.init.xavier_uniform_(self.weight) def forward(self, x, cheb_polys): # x: (B, C_in, N) # cheb_polys: list of (N, N)长度 K out 0 for k in range(self.K): x_k torch.einsum(bcn,nm-bcm, x, cheb_polys[k]) out out torch.einsum(bcn,co-bon, x_k, self.weight[k]) return out self.bias.view(1, -1, 1) class STGCNBlock(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, K, kernel_size3): super().__init__() self.tconv1 TemporalConv(in_channels, hidden_channels, kernel_size) self.cheb ChebConv(hidden_channels, hidden_channels, K) self.tconv2 TemporalConv(hidden_channels, out_channels, kernel_size) if in_channels ! out_channels: self.residual nn.Conv2d(in_channels, out_channels, kernel_size(1, 1)) else: self.residual nn.Identity() self.relu nn.ReLU() def forward(self, x, cheb_polys): # x: (B, C_in, T, N) res self.residual(x) h self.tconv1(x) B, C, T, N h.shape h h.permute(0, 2, 1, 3).reshape(B * T, C, N) h self.cheb(h, cheb_polys) h h.reshape(B, T, C, N).permute(0, 2, 1, 3) h self.tconv2(h) return self.relu(h res) class STGCN(nn.Module): def __init__(self, num_nodes, in_channels, pred_len, cheb_polys, hidden64, K3): super().__init__() self.register_buffer(cheb_polys, torch.stack(cheb_polys, dim0)) # (K, N, N) self.block1 STGCNBlock(in_channels, hidden, hidden, K) self.block2 STGCNBlock(hidden, hidden, hidden, K) self.head nn.Conv2d(hidden, pred_len, kernel_size(1, 1)) def forward(self, x): # x: (B, C_in, T_in, N) polys [self.cheb_polys[k] for k in range(self.cheb_polys.shape[0])] h self.block1(x, polys) h self.block2(h, polys) out self.head(h) # (B, pred_len, T_in, N) out out.mean(dim2) # 在时间维平均输出 (B, pred_len, N) return outhidden控制中间通道数64 是常见起点数据量大可以加到 128。head把通道映射到预测步数最后在时间维取平均是为了把(B, pred_len, T_in, N)压成(B, pred_len, N)。也可以用最后时刻代替平均但对交通速度来说平均能降低个别时间步的噪声。3.4 训练循环、损失函数与学习率调度损失用 MSE优化器用 Adam学习率从 1e-3 开始。梯度裁剪防止早期爆炸。标签在 Dataset 里是(B, N, C, T_out)进模型前要变成(B, T_out, N)。def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0.0 for x, y in loader: x x.to(device) # (B, N, C, T_in) y y.to(device) # (B, N, C, T_out) x x.permute(0, 2, 3, 1).contiguous() # (B, C, T_in, N) y y.permute(0, 2, 3, 1).mean(dim2) # (B, T_out, N) optimizer.zero_grad() pred model(x) # (B, T_out, N) loss criterion(pred, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() * x.size(0) return total_loss / len(loader.dataset)clip_grad_norm_的max_norm取 5.0是经验值梯度范数持续很大时再降到 1.0。学习率调度可以用StepLR每 20 个 epoch 乘以 0.5。如果验证集 MAE 连续 5 轮不降就停别硬跑。超参数建议起点作用input_len12历史时间窗pred_len3预测步数K3切比雪夫阶数hidden64中间通道数kernel_size3时间卷积核lr1e-3初始学习率batch_size64批大小4. 训练 STGCN 时的指标、调参与排错训练能跑通不代表模型学到了路网结构。交通速度里有零值和缺失值指标要加掩码。超参数里 K、时间核和通道数互相牵制。报错大多集中在环境、维度和邻接矩阵三处。4.1 评价指标MAE、RMSE、MAPE 与零值掩码交通速度数据里 0 可能代表缺失直接算 MAPE 会炸。用掩码把无效位置排除并重新归一化掩码均值。MAE 和 RMSE 看绝对误差MAPE 看相对误差但 MAPE 对接近零的真实值仍然敏感上线前最好只作为参考。import numpy as np def masked_mae(pred, true, null_val0.0): mask (true ! null_val).astype(np.float32) mask / mask.mean() mask np.nan_to_num(mask) return np.mean(np.abs(pred - true) * mask) def masked_rmse(pred, true, null_val0.0): mask (true ! null_val).astype(np.float32) mask / mask.mean() mask np.nan_to_num(mask) return np.sqrt(np.mean(np.square(pred - true) * mask)) def masked_mape(pred, true, null_val0.0, eps1e-3): mask (true ! null_val).astype(np.float32) mask / mask.mean() mask np.nan_to_num(mask) return np.mean(np.abs((pred - true) / np.maximum(np.abs(true), eps)) * mask)null_val按数据集定义常见是 0。eps防止除零。掩码均值归一化后无效位置不会拉低整体指标。如果验证集 MAE 正常但 MAPE 很大先查是不是很多真实速度接近零。4.2 关键超参数切比雪夫阶数 K、时间核大小、通道数、学习率K从 2 加到 4验证集 MAE 通常会先降后升。kernel_size从 3 加到 5能看更长的时间片段但参数量增加。hidden从 64 加到 128训练集损失会降验证集不一定。学习率 1e-3 太大时损失震荡可以先用 5e-4 跑几轮。参数调大影响建议K空间感受野变大易过拟合从 3 开始数据少用 2kernel_size时间感受野变大参数增加3 或 5别超过 7hidden拟合能力增强显存上升64 起步按验证集加lr收敛快可能震荡1e-3 或 5e-4batch_size梯度稳定内存吃紧64显存小用 32调参时一次只动一个记录验证集 MAE。如果训练损失降、验证损失升先加 dropout 或减 hidden不要先加 K。4.3 常见报错CUDA 版本、DLL 初始化失败、邻接矩阵维度不匹配CUDA 版本不对torch.cuda.is_available()会返回 False但代码可能仍然在 CPU 上跑只是慢。Windows 上OSError: [WinError 1114] 动态链接库(DLL)初始化例程失败常和c10.dll有关优先重装干净环境。邻接矩阵维度不匹配通常在torch.einsum或矩阵乘法时报错检查 N 是否等于节点数。# 确认环境 python -c import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available()) # 确认邻接矩阵形状 python -c import numpy as np; anp.load(adj.npy); print(a.shape)如果adj.npy是(207, 207)模型里的num_nodes也必须是 207。切比雪夫多项式数量必须等于K少一个就会在循环里索引越界。GPU 和 CPU 切换时记得把cheb_polys也移到同一设备。4.4 自适应图卷积要不要上与预定义邻接矩阵的对比预定义邻接矩阵来自距离或连通关系物理意义强但只能表达静态边。自适应图卷积让模型自己学节点嵌入再算节点间相似度能补出隐藏依赖。代价是参数量增加小数据集上容易过拟合。稳妥做法是把两者融合而不是直接替换。方案优点风险预定义邻接可解释训练稳缺失动态边自适应图卷积能学隐藏依赖过拟合可解释性下降融合兼顾两者需要调融合权重如果在验证集上自适应版本没有稳定提升先保留预定义邻接把精力放在时间窗口和归一化上。5. 进阶技巧自适应图卷积与多步预测的落地细节5.1 自适应图卷积的实现E E^T 与预定义邻接融合自适应图卷积用节点嵌入矩阵 E 计算相似度再和预定义邻接相加。E 的形状是(N, embed_dim)E E.T得到(N, N)过 ReLU 和 softmax 后作为动态邻接。class AdaptiveAdj(nn.Module): def __init__(self, num_nodes, embed_dim10, alpha0.5): super().__init__() self.E nn.Parameter(torch.randn(num_nodes, embed_dim) * 0.01) self.alpha alpha def forward(self, pre_adj): # pre_adj: (N, N) 预定义邻接矩阵 adp torch.softmax(torch.relu(self.E self.E.T), dim1) return pre_adj self.alpha * adpembed_dim常取 10 到 32alpha控制动态邻接的权重。融合后的矩阵要重新做归一化拉普拉斯再算切比雪夫多项式。训练时 E 会更新所以切比雪夫多项式不能只算一次最好每个 epoch 或每若干步重算。5.2 预测结果反归一化与可视化验证模型输出是归一化后的速度反归一化要用训练集的均值和标准差。可视化时挑几个节点把真实曲线和预测曲线画在一起重点看峰值是否有滞后。def inverse_transform(pred, mean, std): # pred: (B, T_out, N)mean/std: (N,) 或标量 return pred * std mean如果峰值总是晚一步说明时间卷积的感受野不够可以加大kernel_size。如果峰值被削平可能是损失对高峰不敏感考虑加权 MSE 或换 Huber。5.3 多步预测的滚动策略与误差校正直接输出多个未来步训练稳定但长步误差会累积。滚动预测每次只输出下一步再把预测值拼回输入窗口适合在线场景但误差也会滚雪球。一个折中做法是模型一次输出 3 步滚动时每 3 步用真实值校正一次输入窗口。这样既保留多步输出的稳定性又不会让误差无限传播。上线前用验证集比较“直接多输出”和“滚动 3 步校正”的 MAE再决定用哪种。本文还有配套的精品资源点击获取