简介这是一份基于PyTorch的时空图卷积网络STGCN实现代码源自IJCAI 2018论文官方工程面向人体行为分析、动作识别等方向的研究者与开发者解决骨骼序列数据中空间拓扑关系与时间动态规律的联合建模问题。压缩包共12个文件以Python源码、Markdown文档和备份文件为主整体约28.71MB目前已有121人学习下载。资源提供了标准化的数据处理流程包括关节坐标归一化与坐标系转换并采用模块化方式搭建图卷积层、时序卷积层及全连接分类器训练过程使用交叉熵损失和Adam优化器辅以动态学习率调整。通过阅读源码与运行实验可深入理解时空特征并行提取的实现技巧掌握从数据预处理到模型评估的完整PyTorch项目开发范式为智能监控、人机交互等应用场景提供可直接扩展的算法基础。1. STGCN 是什么用图卷积吃透路网拓扑不是又一种“CNN 换皮”有人说 STGCN 就是把 CNN 换了个卷积核这句话误导了不少人。STGCN 时空图卷积网络Spatio-Temporal Graph Convolutional Network解决的是路网交通流预测这类自带拓扑结构的回归问题每个监测点是一条时间序列点与点之间又隔着一条真实路网。普通卷积神经网络把规则网格当输入天然吃不下这种“节点 边”的数据。STGCN 用图卷积吃掉空间依赖用门控时间卷积吃掉时序依赖两条线叠成 ST-Conv Block在 PyTorch 里几十行代码就能搭出可用版本。这篇笔记写给要自己写实现、调参数、复现基线的人默认你已经会跑 PyTorch但不一定懂图卷积。读完你能照着代码跑通一个最小模型也知道真正上线前该在哪里较劲。2. 先把网络拆开看STGCN 的图卷积与时间卷积在 PyTorch 里各是什么形状STGCN 的代码实现难不在 PyTorch 语法难在把“图上的卷积”和“时序上的卷积”的维度对齐。我见过太多人把两套卷积直接拼在一起维度报错后就开始瞎 permute。先把两条支线拆清楚后面写模型就是按图索骥。2.1 图卷积支线邻接矩阵、归一化拉普拉斯与 Chebyshev 展开普通卷积依赖平移不变性但路网没有这个概念每个节点的邻居数量不一样邻居是谁也不一样。图卷积的常见做法是用邻接矩阵 A 描述节点连接关系然后通过归一化控制消息传递的尺度。最常用的形式是先把 A 加上自环得到 \tilde A再算对称归一化拉普拉斯L I - D^{-1/2} \tilde A D^{-1/2}D 是 \tilde A 的度矩阵。这一步把节点特征从“绝对数值”变成“与邻居的差异”避免度数大的节点把消息放大。STGCN 原文用的是 Chebyshev 多项式近似把卷积核 g_θ 近似成 K 阶 Chebyshev 多项式避免对拉普拉斯做昂贵的特征分解。展开式是 T_0(x)x、T_1(x)\tilde L x、T_k(x)2\tilde L T_{k-1}(x)-T_{k-2}(x)这里的 \tilde L 是特征值缩放到 [-1, 1] 的拉普拉斯矩阵。在 PyTorch 里图卷积的输入特征通常长成 (B, C, T, N)B 是批次C 是通道数T 是时间步N 是节点数。图卷积只该作用于 N 这一维所以要把张量换到 (B, T, N, C) 或者 (B·T, N, C)用矩阵乘法把 L 和节点特征乘起来时间维和通道维全程不动。K 阶里 K2 到 3 就够用继续加大很容易过平滑——信号在图上传来传去最后所有节点趋向同一个值。2.2 时间卷积支线门控一维卷积为什么比 LSTM 更适合短序列STGCN 的时间支线不是 RNN而是一维卷积加门控。代码实现里通常写成Conv2d卷积核是 (Kt, 1)也就是只在时间维上滑节点维不动。门控机制采用 GLU卷积输出拆成两半 P 和 Q最终输出 P ⊙ σ(Q)。σ 是 sigmoid它的取值范围给模型一个“选择保留多少信息”的能力比单纯线性卷积表达能力更强。有人问为什么不直接用 LSTM。短序列场景下比如用过去 12 步预测未来 12 步LSTM 的迭代式推理训练慢、梯度路径长还不容易复现。卷积的优点是并行度高、行为可预期。跑实验时你会发现同一套超参在不同 seed 下 LSTM 结果能差好几个点STGCN 的波动小得多。做基线和上线都更省心。这里有一个容易忽略的关键点普通 Conv1d 默认是“非因果”的填充会让卷积核看到未来时刻。做短程预测时 padding 设为 Kt // 2序列长度不变影响不大。但要改成自回归式预测就得把 padding 换成左填充或者手工 mask 掉未来位置否则推理时信息泄漏线上表现会翻车。2.3 ST-Conv Block 的拼装逻辑残差、瓶颈与张量形状流转STGCN 主体由几个 ST-Conv Block 堆叠。每个 Block 的结构是“时间卷积 → 图卷积 → 时间卷积”中间夹残差连接。图卷积放在两个时间卷积中间出发点很朴素先让每个节点在时间上把自己理清楚再沿图结构融合邻居信息最后再在时间维度上把融合后的特征还原。用我常用的参数排列说明张量流转。输入 (B, C_in, T, N)第一层时间卷积输出 (B, C_hidden, T, N)图卷积保持 (B, C_hidden, T, N)第二层时间卷积输出 (B, C_out, T, N)。残差连接从输入直接接到输出如果 C_in 不等于 C_out就用一个 1x1 卷积先对齐通道。这里可以做一个瓶颈设计第一层时间卷积把通道压到 hidden最后再放大到 out让中间的图卷积在较低维度下工作。参数含义常见取值KChebyshev 阶数决定空间感受野2 ~ 3Kt时间卷积核宽度3 ~ 5C_hidden图卷积输入通道32 ~ 64num_pred预测步数6 / 12这个结构在 PyTorch 里没有魔法只有一组被反复验证过的 Conv2d 和 matmul。把形状对齐了别的都好说。3. 搭一个最小可运行的 STGCN 模型PyTorch 实现代码与逐段说明这一章直接给可复现代码。我按“输入形状约定 → 邻接矩阵处理 → 模型定义 → 前向验证”的顺序写。先把形状约定说死后面所有代码都以这个为准。3.1 输入输出形状约定先把 (B, C, T, N) 这四维对齐我习惯统一成 (B, C, T, N)其中 B 是 batchC 是每个时刻的输入通道单特征流量就是 1T 是历史时间窗口长度N 是监测点数量。模型输出是 (B, num_pred, N)表示每个节点未来 num_pred 步的预测值。如果你的原始数据是 (T_total, N, 1)构造样本时要做两步先用滑动窗口切出 X 和 Y再把维度从 (B, T, N, F) 转成 (B, F, T, N)。这一步很多新手漏掉PyTorch 的 Conv2d 默认第四维是空间宽度直接把 (B, T, N, F) 丢进去卷积会在时间和节点两个维度上同时滑动结果完全不对。def to_samples(raw, hist_steps12, pred_steps12): # raw: (T_total, N, F) 按时间顺序排列 total raw.shape[0] x, y [], [] for i in range(total - hist_steps - pred_steps 1): x.append(raw[i: i hist_steps]) # (hist_steps, N, F) y.append(raw[i hist_steps: i hist_steps pred_steps]) x torch.tensor(np.array(x)) # (B, T, N, F) y torch.tensor(np.array(y)) # (B, pred_steps, N, F) x x.permute(0, 3, 1, 2) # (B, F, T, N) y y.permute(0, 3, 1, 2) # (B, F, pred_steps, N) return x, y这段代码关键点是最后的 permute。原始切片里每个样本是 (T, N, F)堆叠后是 (B, T, N, F)转成 (B, F, T, N) 才能喂给后面的 Conv2d。如果你的输入有多个特征比如流量、速度、占有率F 1就保留这一维如果只有流量F1后续取 y[:, 0] 即可。3.2 邻接矩阵归一化代码这一步错了后面全错图卷积的质量基本由邻接矩阵决定。常见的数据集给的是检测点之间的距离矩阵先按阈值转成 0/1 邻接矩阵再算归一化拉普拉斯。不归一化直接喂给 Chebyshev 展开特征值范围会超过 [-1,1]多项式递归几轮后数值直接爆炸Loss 变成 NaN。import numpy as np import torch def build_scaled_laplacian(adj): adj: (N, N) 的 0/1 邻接矩阵先加自环再做对称归一化 返回: (N, N) 的 torch float32 张量特征值范围落在 [-1, 1] 附近 n adj.shape[0] adj adj np.eye(n) # 加自环 deg adj.sum(axis1) deg_inv_sqrt np.power(deg, -0.5) deg_inv_sqrt[np.isinf(deg_inv_sqrt)] 0.0 # 孤立点保护 d_mat np.diag(deg_inv_sqrt) lap np.eye(n) - d_mat adj d_mat # 对称归一化拉普拉斯 lam_max np.linalg.eigvals(lap).real.max() # 求最大特征值 lap_scaled 2.0 * lap / (lam_max 1e-6) - np.eye(n) return torch.tensor(lap_scaled, dtypetorch.float32)两点说明。第一自环必须加否则一个节点的更新完全忽略自身信息模型等于只学邻居置换。第二用np.linalg.eigvals求最大特征值在小图上没问题N 到几千时很慢可以用幂迭代法替代只求最大的那个特征值。实际数据里相邻检测点距离远大于阈值时邻接矩阵会非常稀疏这种稀疏性后面可以转成 PyTorch 稀疏张量省显存。3.3 图卷积层、门控时间卷积层与网络主干的完整代码先把两个基础模块写出来。ChebConv 实现 K 阶 Chebyshev 图卷积TemporalConv 实现 GLU 门控时间卷积。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.randn(K, in_channels, out_channels) * 0.05) self.bias nn.Parameter(torch.zeros(out_channels)) def forward(self, x, lap): # x: (B, C, T, N) lap: (N, N) b, c, t, n x.shape x x.permute(0, 2, 3, 1) # (B, T, N, C) if self.K 1: xs [x] else: x0 x x1 torch.matmul(lap, x) # 一阶邻居消息 xs [x0, x1] for _ in range(2, self.K): x2 2.0 * torch.matmul(lap, x1) - x0 xs.append(x2) x0, x1 x1, x2 out sum( torch.einsum(btnc,co-btno, xk, self.weight[k]) for k, xk in enumerate(xs) ) out out self.bias return out.permute(0, 3, 1, 2) # (B, C_out, T, N)这里torch.matmul(lap, x)的广播规则是lap 是 (N, N)x 是 (B, T, N, C)PyTorch 会把前两维当成 batch 维自动在每组 (N, C) 上做矩阵乘法。这一步把每个节点的一跳邻居特征聚合回来。Chebyshev 的递推直接照公式写2.0 是多项式系数K 越大感受野越广。class TemporalConv(nn.Module): def __init__(self, in_channels, out_channels, Kt): super().__init__() self.conv nn.Conv2d( in_channels, out_channels * 2, kernel_size(Kt, 1), # 时间维卷积节点维不动 padding(Kt // 2, 0), ) def forward(self, x): y self.conv(x) # (B, 2*out, T, N) p, q torch.chunk(y, 2, dim1) # 沿通道维拆成两半 return p * torch.sigmoid(q)GLU 的核心在两行通道翻倍是给门控预留空间torch.chunk把结果平均切成 P 和 Qp * sigmoid(q)是门控输出。Kt 建议用奇数比如 3 或 5配合paddingKt//2能保持时间长度不变。最后组装 STGCN 主体。class STConvBlock(nn.Module): def __init__(self, in_channels, hidden_channels, out_channels, K, Kt): super().__init__() self.tconv1 TemporalConv(in_channels, hidden_channels, Kt) self.cheb ChebConv(hidden_channels, hidden_channels, K) self.tconv2 TemporalConv(hidden_channels, out_channels, Kt) self.residual nn.Conv2d(in_channels, out_channels, 1) \ if in_channels ! out_channels else nn.Identity() def forward(self, x, lap): res self.residual(x) out self.tconv1(x) out self.cheb(out, lap) out self.tconv2(out) return out res # 残差连接 class STGCN(nn.Module): def __init__(self, num_nodes, in_channels, hidden_channels, out_channels, K, Kt, num_pred): super().__init__() self.block1 STConvBlock(in_channels, hidden_channels, hidden_channels, K, Kt) self.block2 STConvBlock(hidden_channels, hidden_channels, hidden_channels, K, Kt) self.output nn.Conv2d(hidden_channels, num_pred, 1) def forward(self, x, lap): x self.block1(x, lap) x self.block2(x, lap) x self.output(x) # (B, num_pred, T, N) return x[:, :, -1, :] # (B, num_pred, N)最后一个 1x1 卷积把通道扩成 num_pred只取最后一个时刻的输出直接映射到未来 num_pred 步。这种“一步到位”的多步预测叫 one-shot 策略STGCN 原文就是这么做的。跑前向验证b, c_in, t, n 64, 1, 12, 307 x torch.randn(b, c_in, t, n) adj (np.random.rand(n, n) 0.01).astype(float) # 模拟稀疏邻接 lap build_scaled_laplacian(adj) model STGCN(num_nodesn, in_channels1, hidden_channels32, out_channels32, K3, Kt3, num_pred12) y model(x, lap) print(y.shape) # torch.Size([64, 12, 307])输出形状是 (64, 12, 307)含义是 64 个样本、每个节点未来 12 步的预测值。到这里模型已经通了剩下的是训练和评估这一步最容易让人产生“我在调参其实在碰运气”的错觉。4. 训练与评估STGCN 的几个必调参数和“看着收敛其实过拟合”的时刻模型能跑和模型训练出有效结果之间隔着一条数据处理的河。图卷积类的模型对数据顺序和归一化极其敏感训练阶段最常见的错误是Loss 下降得漂亮换到测试集立刻崩掉。4.1 训练集 / 验证集 / 测试集怎么切归一化参数在哪里拟合时间序列数据绝对不能用随机切分。我见过有人用train_test_split(random_state42)切交通数据把 3 月某天的样本放进训练集4 月的同班车样本放进测试集评估结果虚高得离谱。正确做法是严格按时间顺序切。total raw.shape[0] train_end int(total * 0.7) val_end int(total * 0.8) train_raw raw[:train_end] mean train_raw.mean(axis(0, 1), keepdimsTrue) std train_raw.std(axis(0, 1), keepdimsTrue) 1e-6 def normalize(raw): return (raw - mean) / std train_X, train_Y to_samples(normalize(raw[:train_end]), 12, 12) val_X, val_Y to_samples(normalize(raw[train_end:val_end]), 12, 12) test_X, test_Y to_samples(normalize(raw[val_end:]), 12, 12)注意mean和std只从训练段计算验证和测试段复用。这要求训练段数据分布能代表整体。如果流量有很强的星期周期性训练集最好覆盖完整的周一到周日否则节假日的预测会整体漂移。4.2 损失函数与评估指标MAE、RMSE、MAPE 的实现细节交通流数据大量存在缺失值和零值。很多公开数据集用 0 表示无车流这和真正的“零流量”语义重叠。直接用nn.L1Loss()会让缺失位置参与梯度更新把模型往平均值方向拉。我习惯先写一个带 mask 的损失函数。def masked_mae(pred, true, null_val0.0): mask (true ! null_val).float() m 1e-4 mask.sum() return (torch.abs(pred - true) * mask).sum() / m def masked_rmse(pred, true, null_val0.0): mask (true ! null_val).float() m 1e-4 mask.sum() return torch.sqrt(((pred - true) ** 2 * mask).sum() / m) def masked_mape(pred, true, null_val0.0): mask (true ! null_val).float() denom torch.abs(true) 1e-4 # 防止分母为 0 return ((torch.abs(pred - true) / denom) * mask).sum() / (1e-4 mask.sum())MAPE 分母里的1e-4是关键。零流量时刻算出的百分比可能是百分之几千直接把整个指标带歪。加常数会让 MAPE 不再“纯净”但换来了可比性。报告中同时给出 MAE 和 RMSE让读者能从绝对误差和粗差异两个角度判断。训练循环里我一般用 Adam 初始学习率 0.001配上ReduceLROnPlateaupatience 设 5。批量大小在显存允许范围内尽量取大图卷积的einsum对 batch 扩张很敏感同样数据量跑 32 batch 和 128 batch 的收敛速度差距明显。optimizer torch.optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, patience5, factor0.5, verboseTrue) for epoch in range(100): model.train() total_loss 0.0 for xb, yb in train_loader: # xb: (B, F, T, N) yb: (B, F, pred_steps, N) pred model(xb, lap) loss masked_mae(pred, yb[:, 0]) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() val_loss evaluate(model, val_loader, lap) scheduler.step(val_loss)4.3 参数调整优先级K 阶、隐层通道、时间核宽度的顺序STGCN 的可调参数不少新手最容易乱调。我调参时的固定顺序是第一K 阶。K 从 2 提到 3观察验证集 MAE 变化。如果没变化说明路网空间依赖只有一跳有效继续加大 K 只会增加过平滑风险。如果明显下降可以试 K4但不要超过 5。第二Kt 时间核宽度。Kt3 是起步如果历史窗口长比如 24 步Kt5 甚至 7 有效。这里有个容易忽略的交互Kt 变大后padding 也要跟着变否则时间维长度缩水后面的残差连接形状对不上。第三隐层通道数。从 32 起步显存允许就试 64 和 128。这个参数与数据集节点数强相关N307 的 METR-LA 和 N1700 的 PEMS-BAY 最优值差别很大。节点越多单个通道能表达的模式越有限通常需要更大通道数。最后调学习率和 dropout。STGCN 这类模型对 dropout 不敏感0 到 0.3 之间影响不大建议先不动。5. STGCN 避坑指南PyTorch 实现里最常见的 5 个翻车现场下面的坑我都踩过按“现象 → 原因 → 解决”写每条都能对应到具体报错或异常指标。5.1 Loss 从一开始就不降甚至第一个 epoch 就 NaN现象训练第一步打印的 loss 是 inf 或 nan或者前几个 epoch 稳定下降后突然爆掉。原因拉普拉斯矩阵没有做缩放特征值超出 Chebyshev 多项式预期的 [-1, 1] 范围递推项数值指数膨胀。解决用build_scaled_laplacian里的lam_max做缩放或者放弃 Chebyshev改用 GCN 风格的归一化D^{-1/2} \tilde A D^{-1/2}让最大特征值天然落在 2 附近稳定性更好。5.2 在验证集上 MAPE 突然变成 300%现象MAE 看着正常MAPE 却高得离谱。原因真实流量中有很多接近 0 的值MAPE 的分母是真实值微小分母会把误差放大到无法解读。解决评估时单独统计“非零流量区间”的 MAPE或者把分母改成max(true, threshold)threshold 按数据集分布取 10 或 20。报告里别只给一个 MAPE把 MAE 放旁边一起看。5.3 节点数一上千训练 OOM现象N1700 的数据集在 2080Ti 上 16G 显存跑不满一个 batch。原因einsum(btnc,co-btno, ...)生成中间张量 (B, T, N, C_out)N 翻一倍显存就翻一倍而且 Chebyshev 的每一阶都会保留一份。解决先把 batch 降到 16再用混合精度如果还不够把 lap 转成torch.sparse_coo_tensormatmul在稀疏张量上会省很多内存和计算。5.4 换一个数据集预测结果变成一条平线现象训练 loss 正常但所有节点的预测值都收敛到整体均值附近。原因节点特征顺序和邻接矩阵顺序没有对齐。公开交通数据集的检测器编号和矩阵行列顺序经常不一致直接读进来等于随机打乱了图结构。解决在数据处理脚本里加一个assert adj.shape[0] node_order.shape[0]并且打印前 5 个节点的邻居编号做人工核验。这个检查 30 秒能做完能省一整天排错时间。5.5 PyTorch 环境装了好几天代码始终跑不到 GPU 上现象torch.cuda.is_available()返回 False模型一直吃 CPU。原因PyTorch 和 CUDA 版本不匹配或者驱动太旧。解决装 GPU 版 PyTorch 前先查nvidia-smi支持的 CUDA 版本再按对应版本安装。装完后用一段小矩阵乘法验证不要一上来就跑完整模型。环境稳定是 STGCN 复现的前提环境问题不解决后面所有调参都是白费。6. 把 STGCN 从“能跑”做到“可信”多步预测、消融验证与结果可视化模型跑通只是第一步。我会再补两个验证动作第一个是可视化单节点预测对比第二个是消融。可视化代码很简单但能暴露很多指标掩盖的问题import matplotlib.pyplot as plt with torch.no_grad(): pred model(test_X, lap).cpu().numpy() # (B, num_pred, N) true test_Y.numpy()[:, 0] # (B, num_pred, N) node_idx 42 plt.figure(figsize(8, 4)) plt.plot(true[0, :, node_idx], labelground truth) plt.plot(pred[0, :, node_idx], labelSTGCN) plt.legend() plt.savefig(stgcn_pred_vs_truth.png, dpi150)看这张图不是看曲线贴得有多紧而是看预测是不是滞后一拍——如果真实曲线在拐点处总是提前或滞后一个时间步说明模型没有学到趋势只是在做平滑复制。这种情况下一百行调参代码都救不回来需要回看 Kt 和门控机制。消融实验我一般做三组去掉图卷积只用时间卷积把 GLU 的 sigmoid 固定成 1退化成普通卷积把 K 从 3 降到 1。三组对照能直接回答“空间信息到底贡献了多少”。这种验证比调参更值得投入因为它决定了 STGCN 在你这个数据集上是不是正确的选择——如果去掉图卷积性能不变那说明路网结构对预测没有帮助你需要的只是一个时序模型。我现在拿到新的时空预测任务第一件事不是调 K而是先画一个节点 24 小时的真实流量曲线再拿验证集预测叠上去看形状。形状对不上参数调得再漂亮也是自欺欺人。这个习惯帮我少掉过很多自己骗自己的实验。希望帮到你。本文还有配套的精品资源点击获取