简介这份文档面向已具备PyTorch与深度学习基础的交通预测学习者与研究者围绕IJCAI 2018论文《Spatio-Temporal Graph Convolutional Networks》展开STGCN的复现与代码解析。内容从论文思路、摘要与引言梳理入手依次讲解PeMSD7(M)等数据集由网格转图结构的处理方式、以切比雪夫近似谱图卷积提取空间特征的图卷积模块、用GLU实现一维因果卷积的门控时间卷积模块以及由两个时间卷积与一个空间卷积构成的“三明治”式ST-Conv Block并说明L2损失、时空块堆叠与输出层如何生成下一时刻速度图同时整理了GCN、时空卷积块与深度学习在交通研究中的应用等知识点。文中还提到图由距离关系数学构建而非实际路网纯卷积结构训练更快、参数更少在中长期交通预测上更具优势。资源包仅1个PDF文件约408KB结构紧凑便于通读。目前已有4486人学习下载适合希望吃透STGCN原理、对照论文与源码理解模型细节的读者。1. 从 ARIMA 到 STGCN交通速度预测为什么需要图卷积很多做交通速度预测的工程师第一反应是上 LSTM但 IJCAI 2018 的 STGCN 论文给了一个反直觉结论纯卷积结构训练速度比基于 RNN 的模型快 10 倍以上参数还更少。它要解决的是短期和中长期交通预测里的两个硬骨头路网不是规则网格CNN 的滑动窗口没法直接套RNN 按时间步迭代误差会累积还难并行。论文在 PeMSD7(M) 和 METR-LA 上验证把交通路网当成图用切比雪夫一阶近似的图卷积抓空间依赖用门控一维卷积抓时间依赖。适合正在做时空序列、图卷积神经网络、交通流预测的读者如果你刚完成 pytorch 安装和 anaconda 配置 pytorch 环境这个复现比只读 pytorch 教程更能建立 pytorch 张量基础。2. Chebyshev 一阶近似下的空间图卷积从路网距离矩阵到 A_hat2.1 交通路网不是规则网格为什么 CNN 核在图上失效交通速度预测里每个检测器或路段是一个节点节点之间通过道路连通或空间距离产生关联。规则图像里卷积核按 3x3 邻域滑动路网没有上下左右只有邻接关系。STGCN 采用 Kipf Welling 2017 的近似谱图卷积把图信号变换到谱域再滤波最终化简成一阶形式只用到邻接矩阵和度矩阵。论文里图的构建不是基于实际路网而是通过数学方法根据距离关系建图。PeMSD7(M) 有 228 个节点METR-LA 有 207 个节点每个节点只取速度一个特征。常见做法是距离阈值高斯核设节点 i、j 距离 d_ij若 d_ij κ则权重 exp(-(d_ij/σ)^2)否则 0。这样得到的 A 是稀疏矩阵既能表达空间邻近又避免全连接。与自适应图卷积相比这种静态图构建更轻但需要道路距离或节点坐标。下面表格列出几种图构建方式的差异。图构建方式输入优点限制距离阈值高斯核节点坐标/路段距离稀疏、物理含义清楚阈值 κ 和 σ 要调实际路网邻接拓扑连接真实连通性数据难获取自适应图卷积可学习节点嵌入不依赖先验图参数和显存增加单位阵邻接已有 A实现简单忽略距离强度2.2 距离图构建与归一化邻接矩阵的代码实现源码里get_normalized_adj处理邻接矩阵返回度归一化后的 A_wave。它的作用是把原始邻接矩阵变成图卷积直接可用的形式同时把自身节点也纳入聚合范围。import numpy as np def get_normalized_adj(A): Returns the degree normalized adjacency matrix. A A np.diag(np.ones(A.shape[0], dtypenp.float32)) D np.array(np.sum(A, axis1)).reshape((-1,)) D[D 10e-5] 10e-5 # Prevent infs diag np.reciprocal(np.sqrt(D)) A_wave np.multiply(np.multiply(diag.reshape((-1, 1)), A), diag.reshape((1, -1))) return A_wave这段代码先执行A A I表示卷积时不仅看邻居也看节点自身状态。D是度矩阵diag是D^{-1/2}最后用两次np.multiply完成左右归一化等价于D^{-1/2} (AI) D^{-1/2}。参数上A.shape[0]是节点数METR-LA 里是 207PeMSD7(M) 里是 22810e-5用来防止孤立节点的度为零导致inf。返回矩阵是float32和后面 PyTorch 张量保持一致。如果手里只有节点坐标或路段距离常见做法是先构造距离矩阵再用高斯核转成权重def distance_gaussian_kernel(dist, kappa0.5, sigma0.5): # dist: (N, N) 距离矩阵单位与 kappa 一致 W np.exp(-(dist / sigma) ** 2) W[dist kappa] 0.0 np.fill_diagonal(W, 1.0) return W.astype(np.float32)kappa控制邻域半径超过就断开sigma控制衰减速度越小则远处节点权重掉得越快。注意距离矩阵最好先做归一化否则sigma在不同城市、不同单位下没有可比性。np.fill_diagonal(W, 1.0)是给自身边一个明确权重后面还会再加单位阵所以这里也可以置 0只要保持逻辑一致。2.3 Chebyshev 多项式近似与一阶化简的工程取舍原始谱图卷积需要特征分解复杂度高图一大就难算。Chebyshev 近似用多项式T_k逼近滤波器截断到 K 阶。Kipf Welling 取 K1 并重归一化得到gθ * x ≈ θ (I D^{-1/2} A D^{-1/2}) x。STGCN 直接把A_hat与输入特征相乘再用Theta矩阵做通道变换。源码里的图卷积核心只有两行lfs torch.einsum(ij,jklm-kilm, [A_hat, t.permute(1, 0, 2, 3)]) t2 F.relu(torch.matmul(lfs, self.Theta1))A_hat形状是(N, N)t形状是(B, N, T, C)先permute(1, 0, 2, 3)变成(N, B, T, C)。einsum(ij,jklm-kilm)用i, j做节点维聚合输出(B, N, T, C)也就是每个节点把邻居特征加权求和。接着matmul(lfs, self.Theta1)在最后一维做通道变换Theta1形状是(C, spatial_channels)把 64 通道压到 16 通道再进入下一个时间卷积。spatial_channels16是瓶颈策略的一部分目的是减少参数和计算量同时让空间与时间模块的通道数对齐。这个一阶近似只保留局部邻域换更远的依赖要靠堆叠 ST-Conv Block 来扩大感受野。3. PyTorch 数据管道Z-score、滑窗与 [B,H,W,C] 的张量契约3.1 METR-LA 原始张量的读取与 Z-score 归一化METR-LA 的node_values.npy常见形状是(34272, 207, 2)表示 34272 个时间片、207 个节点、2 个特征。源码用transpose((1, 2, 0))转成(207, 2, 34272)让节点维在前方便后面按节点采样。归一化用 Z-score对每个特征维度单独减均值、除标准差。def load_metr_la_data(): if (not os.path.isfile(data/adj_mat.npy) or not os.path.isfile(data/node_values.npy)): with zipfile.ZipFile(data/METR-LA.zip, r) as zip_ref: zip_ref.extractall(data/) A np.load(data/adj_mat.npy) X np.load(data/node_values.npy).transpose((1, 2, 0)) X X.astype(np.float32) means np.mean(X, axis(0, 2)) X X - means.reshape(1, -1, 1) stds np.std(X, axis(0, 2)) X X / stds.reshape(1, -1, 1) return A, X, means, stdsaxis(0, 2)表示在节点维和时间维上求统计量保留特征维所以means和stds形状都是(2,)。means.reshape(1, -1, 1)把(2,)变成(1, 2, 1)才能和(207, 2, 34272)做广播。两个统计量必须保存下来因为评估时要反变换回速度单位。注意X.astype(np.float32)放在归一化前避免后续除法产生float64张量。变量形状含义node_values.npy原始(34272, 207, 2)时间, 节点, 特征X转置后(207, 2, 34272)节点, 特征, 时间means/stds(2,)每个特征的均值和标准差A(207, 207)邻接矩阵3.2 generate_dataset 的滑窗切分与样本数量推导generate_dataset用滑动窗口把长时间序列切成监督学习样本。输入窗口num_timesteps_input12预测窗口num_timesteps_output3步长为 1。def generate_dataset(X, num_timesteps_input, num_timesteps_output): indices [(i, i (num_timesteps_input num_timesteps_output)) for i in range(X.shape[2] - (num_timesteps_input num_timesteps_output) 1)] features, target [], [] for i, j in indices: features.append(X[:, :, i: i num_timesteps_input].transpose((0, 2, 1))) target.append(X[:, 0, i num_timesteps_input: j]) return torch.from_numpy(np.array(features)), torch.from_numpy(np.array(target))indices的个数是X.shape[2] - (12 3) 1。训练集取 60%所以时间片数是34272 * 0.6 20563样本数为20563 - 15 1 20549。features里每个样本先取X[:, :, i:i12]形状(207, 2, 12)再transpose((0, 2, 1))变成(207, 12, 2)target只取第 0 个特征通道即速度形状(207, 3)。最终X返回(20549, 207, 12, 2)Y返回(20549, 207, 3)。这里预测目标只有速度特征里的另一个通道可以是流量或占有率训练时作为输入辅助信息。3.3 DataLoader 的 batch 组装与维度重排batch_size50每个 batch 的X是(50, 207, 12, 2)Y是(50, 207, 3)。PyTorch 的Conv2d要求输入是(N, C, H, W)所以TimeBlock内部先做X.permute(0, 3, 1, 2)变成(50, 2, 207, 12)。这里的H是节点数 207W是时间长度 12卷积核(1, kernel_size3)只沿时间维滑动不混合节点。如果保持(50, 207, 12, 2)直接卷积通道维变成最后一个维度卷积会沿节点和时间混合语义完全错。阶段形状含义原始X(207, 2, 34272)节点, 特征, 时间训练集X(207, 2, 20563)60% 时间片样本X(20549, 207, 12, 2)样本, 节点, 历史步, 特征batchX(50, 207, 12, 2)batch, 节点, 历史步, 特征TimeBlock输入(50, 2, 207, 12)batch, 通道, 节点, 时间batchY(50, 207, 3)batch, 节点, 未来 3 步提示把num_timesteps_input和num_timesteps_output写进配置文件不要散落在函数调用里。后面改预测窗口时generate_dataset、STGCN的Linear输入、评估反变换要一起改漏掉一个就会形状不匹配。4. STGCN 三明治结构TimeBlock 门控卷积、Theta 空间变换与 STGCNBlock 前向4.1 GLU 与一维因果卷积TimeBlock 为什么不用 RNNSTGCN 的时间模块叫TimeBlock核心是三个Conv2d对每个节点独立做一维时间卷积。论文用 GLU 实现门控源码里的写法是conv1作为主体conv2经过sigmoid做门conv3再参与残差式相加。class TimeBlock(nn.Module): def __init__(self, in_channels, out_channels, kernel_size3): super(TimeBlock, self).__init__() self.conv1 nn.Conv2d(in_channels, out_channels, (1, kernel_size)) self.conv2 nn.Conv2d(in_channels, out_channels, (1, kernel_size)) self.conv3 nn.Conv2d(in_channels, out_channels, (1, kernel_size)) def forward(self, X): # X: (batch_size, num_nodes, num_timesteps, num_features) X X.permute(0, 3, 1, 2) # (B, C, N, T) temp self.conv1(X) torch.sigmoid(self.conv2(X)) out F.relu(temp self.conv3(X)) out out.permute(0, 2, 3, 1) # (B, N, T, C) return out三个卷积核形状都是(1, kernel_size)也就是说只在时间维W上滑动节点维H不参与。kernel_size3且没有 padding每过一层时间长度减少kernel_size - 1 2。conv2的sigmoid输出在 0 到 1 之间相当于给conv1的结果做门控conv3再补一条通路F.relu保留非线性。与 RNN 相比这种结构不依赖上一步输出整个时间窗内的卷积可以并行算训练速度优势就来自这里。4.2 Theta 参数矩阵与 einsum 实现图卷积STGCNBlock把时间卷积、空间图卷积、再一个时间卷积串成三明治。空间部分没有用nn.Conv2d而是一个可学习的Theta1参数矩阵配合A_hat做节点间聚合。class STGCNBlock(nn.Module): def __init__(self, in_channels, spatial_channels, out_channels, num_nodes): super(STGCNBlock, self).__init__() self.temporal1 TimeBlock(in_channelsin_channels, out_channelsout_channels) self.Theta1 nn.Parameter(torch.FloatTensor(out_channels, spatial_channels)) self.temporal2 TimeBlock(in_channelsspatial_channels, out_channelsout_channels) self.batch_norm nn.BatchNorm2d(num_nodes) self.reset_parameters() def reset_parameters(self): stdv 1. / math.sqrt(self.Theta1.shape[1]) self.Theta1.data.uniform_(-stdv, stdv) def forward(self, X, A_hat): t self.temporal1(X) lfs torch.einsum(ij,jklm-kilm, [A_hat, t.permute(1, 0, 2, 3)]) t2 F.relu(torch.matmul(lfs, self.Theta1)) t3 self.temporal2(t2) return self.batch_norm(t3)Theta1形状(out_channels, spatial_channels)初始化用均匀分布范围由1 / sqrt(spatial_channels)决定。einsum完成节点维聚合matmul完成通道变换F.relu增加非线性。通道数从out_channels64压到spatial_channels16再恢复到 64这是瓶颈策略空间变换的参数量从64*64降到64*16同时迫使模型在低维空间里保留关键空间模式。BatchNorm2d(num_nodes)把节点数当作通道数训练时按节点维统计均值和方差如果自己改网络结构建议先permute成(B, C, N, T)再 BN或者换成LayerNorm否则容易在 batch 较小时出现统计抖动。4.3 两个 STGCNBlock 加输出层的维度推演整体STGCN堆两个STGCNBlock再接一个TimeBlock和全连接层。第一个 block 的输入通道是 2输出通道 64空间通道 16第二个 block 输入 64输出 64空间通道 16最后的TimeBlock把通道统一到 64再展平时间维和通道维。class STGCN(nn.Module): def __init__(self, num_nodes, num_features, num_timesteps_input, num_timesteps_output): super(STGCN, self).__init__() self.block1 STGCNBlock(in_channelsnum_features, out_channels64, spatial_channels16, num_nodesnum_nodes) self.block2 STGCNBlock(in_channels64, out_channels64, spatial_channels16, num_nodesnum_nodes) self.last_temporal TimeBlock(in_channels64, out_channels64) self.fully nn.Linear((num_timesteps_input - 2 * 5) * 64, num_timesteps_output) def forward(self, A_hat, X): out1 self.block1(X, A_hat) out2 self.block2(out1, A_hat) out3 self.last_temporal(out2) out4 self.fully(out3.reshape((out3.shape[0], out3.shape[1], -1))) return out4以batch_size50、num_timesteps_input12、num_timesteps_output3为例维度变化如下表。每个STGCNBlock里有两个TimeBlock每个减少 2所以一个 block 时间维减少 4两个 block 减少 8最后last_temporal再减少 2总减少 10剩下 2。模块输入输出时间维变化block1的temporal1(50,207,12,2)(50,207,10,64)12 → 10block1的Theta1(50,207,10,64)(50,207,10,16)通道 64 → 16block1的temporal2(50,207,10,16)(50,207,8,64)10 → 8block2整体(50,207,8,64)(50,207,4,64)8 → 4last_temporal(50,207,4,64)(50,207,2,64)4 → 2reshape(50,207,2,64)(50,207,128)展平时间与通道fully(50,207,128)(50,207,3)输出未来 3 步Linear输入写的是(num_timesteps_input - 2 * 5) * 64其中2 * 5对应 5 个TimeBlock各减少 2。out3.reshape((out3.shape[0], out3.shape[1], -1))把(B, N, 2, 64)展平成(B, N, 128)再对每个节点做全连接输出(B, N, 3)。损失用 L2训练时直接对归一化后的速度做回归。论文强调纯卷积参数少、训练快这个结构里最耗参的其实就是几个Conv2d和最后的Linear没有 RNN 的循环权重。5. 训练与收敛排查kernel_size3 的边界、BN 顺序和全连接层形状匹配5.1 从 W 维度的减少量反推 Linear 输入全连接层的输入维度不是拍脑袋写的它由时间维减少量决定。源码里Linear((num_timesteps_input - 2 * 5) * 64, num_timesteps_output)只在kernel_size3时成立。5 个TimeBlock分别是两个STGCNBlock里的 4 个加上最后 1 个每个减少kernel_size - 1。把公式写清楚num_timesteps_input 12 kernel_size 3 num_time_blocks 5 reduction num_time_blocks * (kernel_size - 1) assert num_timesteps_input - reduction 0 linear_in (num_timesteps_input - reduction) * 64 print(linear_in) # 128如果改成kernel_size5每个 TimeBlock 减少 45 个共减少 2012 步输入直接变成负数Conv2d在时间维会先报形状错误。此时要么把num_timesteps_input增大到 20 以上要么给时间卷积加 padding要么减少 block 数量。我一般会先把num_time_blocks和kernel_size改成变量训练前跑一遍形状断言避免训练到一半才发现Linear维度对不上。5.2 显存与批大小调整的取舍batch_size50在 METR-LA 上能跑但换到节点更多的 PeMSD7(M) 或更长历史窗口时容易爆显存。中间激活最大的一块来自block1的第一个时间卷积形状(50, 64, 207, 10)约 660 万个浮点数单精度正向激活约 26MB反向还要再存一份。如果 OOM优先降batch_size再降spatial_channels最后才动num_timesteps_input因为历史窗口一变全连接层输入也要重算。调整项影响建议batch_size50 → 16显存线性下降BN 统计更抖先固定随机种子对比 lossspatial_channels16 → 8参数和激活下降精度可能掉 1% 到 3%kernel_size3 → 2每个 TimeBlock 只减 1全连接输入从 128 变 448num_timesteps_input12 → 24历史更长显存上升注意 W 减少后仍为正注意BatchNorm2d(num_nodes)在 batch 较小时统计量不稳定loss 抖动不一定来自学习率。先把batch_size拉回 32 以上或者把 BN 换成LayerNorm再看曲线。5.3 损失曲线抖动时先查归一化反变换训练在 Z-score 空间里算 L2评估时必须反变换回速度单位否则 RMSE 看起来很小但不可解释。反变换用保存的means和stdsdef inverse_transform(X, means, stds): # X: (B, N, T) 或 (B, N, T, C) return X * stds.reshape(1, -1, 1) means.reshape(1, -1, 1)如果只预测速度通道取stds[0]和means[0]即可如果多通道一起预测形状要对齐到(1, C, 1)。排查 loss 抖动时先拿 200 个样本过拟合关掉 shuffle把batch_size设为 16训练几十个 epoch如果训练 loss 都降不下去问题通常在A_hat的归一化或Theta1初始化如果训练 loss 降了但验证 RMSE 反变换后异常大检查stds是否接近零、有没有把特征通道取错。先用小 batch 过拟合 200 个样本验证全连接层 reshape 是否对齐再放开全量训练。本文还有配套的精品资源点击获取