STGCN时空图卷积网络PyTorch实现与交通流量预测指南

简介：这是一份基于 PyTorch 的 STGCN 时空图卷积网络实现代码，源自 IJCAI 2018 会议论文的官方工程，面向希望复现时空序列建模算法的开发者与研究者。其核心价值在于将图卷积与时间卷积相结合，分别建模节点间的空间依赖与序列随时间的动态演变，可用于交通流预测、人体行为识别等典型场景。压缩包共 12 个文件，大小约 28.71MB，主要包含 Python 源码、Markdown 说明文档、LICENSE 及备份文件；其中 stgcn.py 定义模型主体结构，utils.py 和 main.py 构成数据预处理与训练评估流程，README 则提供环境配置与使用说明，目录组织清晰。目前已有 123 人学习。通过阅读和运行这套代码，可以系统掌握 STGCN 的模块化搭建方法、标准化坐标系处理、损失函数与 Adam 优化器配置，并基于验证集准确率完成模型调优；还可在此基础上尝试引入注意力机制或长短期记忆网络进行扩展，非常适合个人算法复现与学习研究。

1. STGCN时空图卷积网络：一段能直接跑的PyTorch预测代码

做交通流量预测或者气象站点温度外推这类任务时，最难受的不是模型选型，而是“数据既有时间顺序又有空间结构”这件事本身——纯 LSTM 丢掉了邻居关系，纯 GCN 又忽略了时间演化。STGCN 就是为这种局面设计的：它把图卷积和时间卷积串进同一个模块，输入一段历史序列，输出未来多步的预测，而且代码结构非常规整，很适合个人学习时逐行读。这份基于 PyTorch 的实现代码，覆盖数据预处理、邻接矩阵构建、模型训练到评估脚本的完整流程，适合想复现时空图卷积算法或把它接到自己数据集上的开发者。

2. STGCN模型结构拆解：ST-Conv Block与切比雪夫图卷积的落点

2.1 为什么图卷积用切比雪夫近似：避免特征分解的高成本

原始的频谱图卷积需要对角化图拉普拉斯矩阵，计算量是 O(N^3)，节点上千时根本没法放到训练循环里。STGCN 用的折中方案是切比雪夫多项式近似：把卷积核写成拉普拉斯算子的多项式，通过递推逐阶聚合邻居信息，完全绕开了特征分解。

K 阶多项式等价于感受野覆盖到 K-1 跳邻居。K=3 时，每个节点能看到自身加上 1 跳、2 跳邻居的信息。实际工程里，K 取 3 就足够覆盖大多数路网或传感器网络的局部结构，再大参数量和过拟合风险都会明显上升。这也是模型里唯一一个既影响精度又影响速度的超参数，后面避坑章节会单独说。

2.2 切比雪夫多项式递推与归一化拉普拉斯

空间卷积的输入是归一化拉普拉斯矩阵，先构建它，再递推生成每阶多项式。常见做法是在邻接矩阵 W 上做对称归一化，再加一个自环防止孤立节点。

def normalized_laplacian(w, alpha=0.8): # w: (N, N) 对称邻接矩阵，已经加过自环 d = w.sum(dim=1) + alpha # 度向量，alpha 防止除零 d_hat = torch.diag(d ** -0.5) L = torch.eye(w.size(0), device=w.device) - d_hat @ w @ d_hat return L def cheb_polynomials(L, K=3): # 返回 [T_0, T_1, ..., T_{K-1}]，每个都是 (N, N) N = L.size(0) T_0 = torch.eye(N, device=L.device) T_1 = L.clone() polys = [T_0, T_1] for i in range(2, K): T_i = 2 * L @ polys[-1] - polys[-2] polys.append(T_i) return polys

逻辑说明：normalized_laplacian 返回的是标准的对称归一化拉普拉斯 L = I - D^-1/2 A D^-1/2。其中 D 是度矩阵，alpha 是我习惯加的稳定常数，理论上加自环之后度不会为零，但某些稀疏图里数值计算仍可能不稳定，加一个小常数能避免 d ** -0.5 出现 inf。cheb_polynomials 的递推式是标准的切比雪夫递推：T_2 = 2L T_1 - T_0，逐阶往上算。K=3 时返回三张矩阵，0 阶是单位阵，1 阶是拉普拉斯本身，2 阶是两跳邻居的聚合权重。注意这里 L 必须是 (N, N) 且与后续输入的节点数一致，否则 einsum 阶段会直接报维度不匹配。

2.3 图卷积层的 PyTorch 实现

图卷积层的输入我统一约定为 (B, N, C)，批量、节点、特征各占一维，输出是 (B, N, C')。空间维 N 是真正的“图维度”，图卷积只负责在节点之间传播信息。

class ChebGraphConv(nn.Module): def __init__(self, K, in_ch, out_ch): super().__init__() self.K = K # 权重维度: 把 K 阶多项式结果在特征维上拼接后做线性变换 self.weight = nn.Parameter(torch.randn(K * in_ch, out_ch) * 0.05) self.bias = nn.Parameter(torch.zeros(out_ch)) def forward(self, x, polys): # x: (B, N, C) b, n, c = x.shape x_list = [x] for k in range(1, self.K): # 逐阶聚合: T_k @ x x_g = torch.einsum('ij,bjc->bic', polys[k], x_list[-1]) x_list.append(x_g) x_cat = torch.cat(x_list, dim=-1) # (B, N, K*C) out = torch.einsum('bnc,co->bno', x_cat, self.weight) + self.bias return out

逻辑说明：先让每一阶多项式分别作用在输入上，得到 K 份 (B, N, C) 的张量，然后在特征维拼接成 (B, N, KC)，最后乘一个 (KC, C') 的权重矩阵，把不同阶邻居的信息融合起来。这个写法比“每阶各设一组权重”参数更少、矩阵乘法更快，而且实现简洁。einsum 的 'ij,bjc->bic' 表示用 T_k 的第 i 行去加权所有节点 j 的特征，聚合到节点 i，相当于做一次标准的图卷积消息传递。

2.4 时间卷积层：GLU 门控与整体模块拼装

时间卷积处理序列的 T 维。PyTorch 的 Conv2d 输入是 (B, C, H, W)，我把空间节点 N 放在 H 位、时间 T 放在 W 位，这样卷积核 (1, Kt) 就只在时间方向滑动，空间维度完全不动。

class TemporalConv(nn.Module): def __init__(self, in_ch, out_ch, Kt=3): super().__init__() # 输出 2*out_ch，一半走门控，一半走线性 self.conv = nn.Conv2d(in_ch, 2 * out_ch, kernel_size=(1, Kt), padding=(0, Kt // 2)) def forward(self, x): # x: (B, C, N, T) y = self.conv(x) # (B, 2*out_ch, N, T) p, q = torch.chunk(y, 2, dim=1) # 沿通道维拆成两半 return p * torch.sigmoid(q)

说明：GLU（门控线性单元）把卷积输出从通道维切成 P、Q 两份。P 作为主线信息直接参与相乘，Q 经过 sigmoid 后变成 0～1 的门控值，控制主线信息的保留程度。这个设计和 LSTM 的门控思路一致，但完全由卷积实现，训练并行度更高。Kt=3、padding=1 时时间维长度不变，后续模块可以直接堆叠，不用反复对齐序列长度。Kt 取奇数是个隐含要求，偶数核配 padding 会导致输出长度比输入多 1，多模块堆叠时误差会累积。

接着把时间卷积和图卷积拼成 ST-Conv Block，结构是「时间卷积 → 图卷积 → 时间卷积」。图卷积期望节点维在倒数第二维，所以中间要临时做一次维度转换。

class STConvBlock(nn.Module): def __init__(self, K, Kt, in_ch, hidden_ch, out_ch): super().__init__() self.tconv1 = TemporalConv(in_ch, hidden_ch, Kt) self.gconv = ChebGraphConv(K, hidden_ch, hidden_ch) self.tconv2 = TemporalConv(hidden_ch, out_ch, Kt) def forward(self, x, polys): # x: (B, C, N, T) out = self.tconv1(x) # (B, hidden_ch, N, T) b, c, n, t = out.shape out = out.permute(0, 3, 2, 1) # (B, T, N, hidden_ch) out = self.gconv(out, polys) # 图卷积，维度不变 out = out.permute(0, 3, 2, 1) # 换回 (B, C, N, T) out = self.tconv2(out) # (B, out_ch, N, T) return out

最后拼出完整模型。我习惯在出口加一个 1x1 卷积把通道压到 1，让输出直接是 (B, 1, N, T)，这样训练循环里少处理一次维度。

class STGCN(nn.Module): def __init__(self, K=3, Kt=3, in_ch=1, hidden_ch=32, out_ch=32, n_block=2): super().__init__() self.blocks = nn.ModuleList() for i in range(n_block): c_in = in_ch if i == 0 else hidden_ch c_out = out_ch if i == n_block - 1 else hidden_ch self.blocks.append(STConvBlock(K, Kt, c_in, hidden_ch, c_out)) self.final = nn.Conv2d(out_ch, 1, kernel_size=(1, 1)) def forward(self, x, polys): # x: (B, 1, N, T) for blk in self.blocks: x = blk(x, polys) return self.final(x) # (B, 1, N, T)

两个 STConvBlock 中间不需要刻意加 ReLU，时间卷积内部的 GLU 已经带了非线性。这个模型的参数量主要由 hidden_ch 决定，32 起步在几百个节点的数据上跑起来很轻量。

3. 数据准备与邻接矩阵：决定结果上限的一步

3.1 原始数据到模型输入的张量组织

原始传感器数据一般是 (T, N) 的矩阵，T 是时间帧数，N 是传感器数量。要喂给 STGCN，得先变成 (B, T_in, N)，每个样本包含一段连续历史、覆盖所有节点。

def make_samples(feats, tw=12, th=3): # feats: (T, N) x_list, y_list = [], [] for i in range(feats.shape[0] - tw - th + 1): x_list.append(feats[i:i + tw]) # 历史 12 帧 y_list.append(feats[i + tw:i + tw + th]) # 未来 3 帧 return np.array(x_list), np.array(y_list)

为什么 tw 取 12、th 取 3：12 帧相当于 5 分钟粒度数据下的一小时，短时预测是 STGCN 最擅长的区间。tw 再长会稀释近邻时间步的权重，训练时间和显存同步上涨。th 取 3 是因为多步预测误差会累积，超过 6 步预测曲线会明显上翘，所以先把 3 步跑通，再考虑拉长预测步长。

如果原始数据每个节点有多个特征，比如流量、速度、占有率，输入变成 (T, N, D)。处理方式是把每个时间步的特征堆到通道维，模型输入从 (B, 1, N, T) 变成 (B, D, N, T)，时间卷积仍然只沿 T 维滑动，空间卷积对每个特征通道分别聚合。D 太大时可以先做一次全连接降维，避免通道数直接拉爆参数量。

3.2 用经纬度距离构造 KNN 邻接矩阵

邻接矩阵是图卷积的空间骨架。最常见做法是用传感器坐标算欧氏距离，每个节点只连接最近的 K 个邻居，边权重用高斯核衰减。

def knn_adjacency(coords, k=8, sigma2=0.1): # coords: (N, 2) 平面坐标或经纬度 n = len(coords) diff = coords[:, None, :] - coords[None, :, :] dist = np.sqrt((diff ** 2).sum(-1)) # (N, N) adj = np.zeros((n, n), dtype=np.float32) for i in range(n): nbrs = np.argsort(dist[i])[1:k + 1] # 跳过自己 for j in nbrs: w = np.exp(-dist[i, j] ** 2 / sigma2) adj[i, j] = w adj[j, i] = w # 对称化 return adj

参数说明：k 是邻居数，几百个节点时取 8 比较均衡，上千节点可以加到 10～12。sigma2 控制距离衰减速度，0.1 表示距离约 0.316 个单位的边上权重仍有 0.4。如果坐标直接用的经纬度，sigma2 通常要放大到几十甚至上百，否则距离一远权重全趋近于零，图卷积等于白做。判断 sigma2 是否合适有个土办法：把邻接矩阵打印出来，看每行的非零元素占全图节点的比例，一般在 5%～10% 之间比较合理。

固定阈值图（距离小于某个值就连边）在均匀分布的传感器上能用，但遇到边缘地区的稀疏节点会产生大量孤立点。KNN 保证每个节点都有固定数量的邻居，比固定阈值稳健。对称化步骤不能省，切比雪夫递推要求图是无向的，非对称矩阵会导致多项式结果畸形，模型训练时 loss 会莫名震荡。

拿到邻接矩阵后，处理顺序是固定的：先对称化、再加自环、最后算归一化拉普拉斯。反过来会让自环在归一化中被削弱，孤立节点的自身信息传不出去。

3.3 归一化与训练验证划分：防止数据泄漏的关键操作

时序数据不能像图像那样随机打乱后切分。随机划分会引入未来信息，验证集里混入了训练集时间段的相邻样本，验证指标虚高，放到线上会直接翻车。

def split_by_time(x, y, train_ratio=0.7): n = len(x) idx = int(n * train_ratio) return (x[:idx], y[:idx]), (x[idx:], y[idx:])

min-max 归一化要在划分之后、只在训练集上统计极值，再用同一套 min、max 处理验证集。这个细节常被漏掉，直接对整个数组归一化再切分，验证集的极值信息就泄漏到了训练分布里，结果是验证集 MAE 永远比真实测试好一截，而且很难察觉。

def min_max_normalize(data, data_min, data_max): return (data - data_min) / (data_max - data_min + 1e-8) def inverse_min_max(data, data_min, data_max): return data * (data_max - data_min) + data_min

提示：min 和 max 必须在训练集上计算，验证集和测试集只能应用同一组极值参数。极值建议用 np.save 存成独立文件，方便推理阶段还原量纲。

归一化为什么选 min-max 而不是 z-score：STGCN 输出的是回归值，min-max 让模型输出天然落在有界区间，逆变换回来不容易产生离谱的外推。z-score 在训练时更稳，但推理时一旦预测值超出训练分布的范围，逆变换后的绝对值可能为负——交通流量和温度这类物理量不允许这种结果。归一化对象是全部节点共同统计，不是按节点单独算。单独算会把不同节点的量纲各自拉平，模型学的是“每个节点的波动形态”，跨节点的空间对比信息就丢失了。

4. 训练与推理：从参数配置到多步预测的输出处理

4.1 训练的默认超参数参考

下面这组参数是我在多个公开路网数据上跑通的起点，不能算最优，但能保证网络稳定收敛、不炸 loss。

早停实现很简单：每轮结束看验证 loss，超过 patience 轮没有刷新最小值，就加载历史最优权重并结束循环。保存的模型应该是验证集最优的那份，而不是最后一轮。最后一轮往往已经过拟合。

4.2 多步预测的输入输出对齐

模型输出的维度是 (B, 1, N, T)。预测目标是未来 th 帧，所以取输出时间维的最后 th 帧，转成 (B, th, N) 后和标签对齐。

def forward_pass(model, x_batch, polys, th=3): # x_batch: (B, tw, N) x = x_batch.permute(0, 2, 1).unsqueeze(1) # (B, 1, N, tw) out = model(x, polys) # (B, 1, N, T) out = out[:, :, :, -th:] # 取最后 th 帧 return out.squeeze(1).permute(0, 2, 1) # (B, th, N)

这里的两次维度操作容易看晕：permute(0, 2, 1) 把 (B, tw, N) 换成 (B, N, tw)，unsqueeze 在通道维插入 C=1。模型内部的时间卷积用核宽 1 处理 N 维，用 Kt 处理 T 维，所以 N 和 T 谁在前谁在后不能搞混。取 -th: 帧是因为模型输出的是一个完整序列，末尾几帧对应最远期的预测，和标签的 th 步对齐。

4.3 训练循环与评估脚本

训练主体是一个标准的 PyTorch 循环。损失函数先用 MSE 跑通，如果数据里离群点很多再换 Huber。

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=5, factor=0.5) loss_fn = nn.MSELoss() for epoch in range(epochs): model.train() train_loss = 0 for xb, yb in train_loader: # xb: (B, tw, N), yb: (B, th, N) optimizer.zero_grad() pred = forward_pass(model, xb, polys) # (B, th, N) loss = loss_fn(pred, yb) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 5) optimizer.step() train_loss += loss.item() * len(xb) train_loss /= len(train_dataset) model.eval() with torch.no_grad(): val_loss = compute_val_loss(model, val_loader, polys) scheduler.step(val_loss) if val_loss < best_val: best_val = val_loss torch.save(model.state_dict(), "best_stgcn.pt")

clip_grad_norm_ 设 max_norm=5 是我必写的，STGCN 在图卷积梯度过大时会梯度爆炸，尤其训练头几轮。scheduler 用 ReduceLROnPlateau 比固定衰减更省心，验证 loss 卡住时自动砍半学习率。

推理和评估前，记得把预测结果还原回真实量纲再算指标：

def evaluate(pred, true, data_min, data_max): pred = inverse_min_max(pred, data_min, data_max) true = inverse_min_max(true, data_min, data_max) mae = np.mean(np.abs(pred - true)) rmse = np.sqrt(np.mean((pred - true) ** 2)) mape = np.mean(np.abs((pred - true) / (true + 1e-4))) * 100 return {"MAE": mae, "RMSE": rmse, "MAPE": mape}

MAPE 在速度这类非零数据上很好用，但真实值接近零时 1e-4 的防护可能不够，分母太小会把 MAPE 撑到几千。如果数据里大量零值，干脆不报 MAPE，只报 MAE 和 RMSE。

5. STGCN代码复现的避坑清单：五条踩过的坑

5.1 归一化后忘记恢复真实量纲

现象：验证集 MAPE 只有 8%，看起来相当好，但把预测曲线画出来一看，全在 0 到 1 之间浮动，完全不是真实的流量数值。

原因：训练时输入做了 min-max 归一化，模型输出的预测值也是 0～1 尺度，但评估脚本里没有做逆变换，直接拿归一化尺度的预测值和真实量纲的标签算误差。这个问题隐蔽在“只用 loss 评估”的工作流里，训练和验证的 loss 都在归一化尺度上，曲线正常，只有画到真实尺度才暴露。

解决：把逆变换写进 evaluate 函数的第一行，并且要求标签也做同样的逆变换，避免一个在原始尺度、一个在归一化尺度的错位。data_min 和 data_max 单独用 np.save 存档，别从日志里手抄，手抄容易抄错精度。

5.2 KNN 邻接矩阵的 k 值设成全局平均

现象：把 k 从 8 改到全部节点数之后，验证 loss 不降反升，模型输出的预测曲线变得特别平，节点之间的差异几乎消失。

原因：k 太大导致邻接矩阵变稠密，图卷积把多跳邻居的信息反复混在一起，空间特征被平均化，每个节点的输出趋同。极端情况下图卷积等价于在全图上做了一次平滑滤波，空间维度形同虚设。

解决：k 控制在节点总数的 5%～10% 之间。节点数 100 时取 8 左右，节点数 1000 时加到的 10～15。调完把邻接矩阵的稀疏率打印出来，非零元素占比超过 30% 就说明图太稠密了，需要回退。

5.3 切比雪夫多项式的 K 不等于越大越好

现象：K 从 3 调到 6 后，训练 loss 顺利下降，但验证 loss 前几轮正常、后面突然飙高，过拟合特征非常明显。

原因：K 控制感受野的同时直接拉高了参数量。K=6 时每层图卷积的权重维度是 6 倍输入通道，两个 STConvBlock 叠加参数翻倍，数据量不够就过拟合。

解决：K=3 起步，先确认流程能跑通，再考虑 K=4 并配合缩小 hidden_ch。如果一定要用大 K，必须加 weight_decay 和早停阈值收紧，并且盯紧验证集。

5.4 对整段时序做随机 shuffle 造成数据泄漏

现象：训练收敛很快，验证集指标很漂亮，但同一份模型换到另一段连续时间的测试数据上明显变差，差到怀疑代码写错了。

原因：构建数据集时直接对整个样本集做了随机打乱。滑窗样本之间天然有时间重叠，验证集里混入了训练集时间段的相邻样本，模型等于提前看到了未来的答案。

解决：按时间顺序切分，前 70% 训练、后 30% 验证。训练 loader 内部的 shuffle 没有问题，它只是打乱了训练集内部顺序，没有引入未来信息。我还会额外检查切分点两侧的滑窗样本没有重叠，否则一次样本同时出现在两边，泄漏更隐晦。

5.5 PyTorch 维度 layout 的混乱

现象：训练时 einsum 报size mismatch，或者 Conv2d 报Expected 4D input，排查半天发现是输入张量的维度顺序不对。

原因：STGCN 内部各层对维度顺序的假设不一样。Conv2d 要求 (B, C, N, T)，图卷积习惯 (B, T, N, C)，早期代码里来回 permute，中间一次没换对就报错。

解决：固定一个主 layout：(B, C, N, T)，只在 STConvBlock 内部做临时转换。模型入口统一转一次，出口统一转一次，中间不再碰维度。每个模块开头加一行注释标注期望的输入形状，这是我在这份代码上花时间最多的地方，后来写任何图神经网络都沿用这套约定。

提示：维度问题最容易发生在“模型能跑但结果全错”的场景，loss 正常下降，但预测与标签形状错位。检查方法是跑一个 batch，打印模型输出和标签的 shape，逐维对比。

6. 进阶：把训练好的STGCN用到自己的数据上

跑通默认的交通流 Demo 之后，最值得做的事是把它迁移到自己手头的数据。这个阶段的核心不是改网络，而是改数据接口和验证习惯。

先做一张误差热力图：把测试集上每个节点的预测误差按节点位置画出来，用颜色映射到图结构上。这张图比任何指标都直观，能一眼看出误差是集中在某个区域还是随机分布。如果误差集聚在某个方向，通常是邻接矩阵构造有问题——KNN 只连了距离近的节点，而实际流量是沿着道路走向传播的，需要换成真实路网连接图或者基于导通时间的距离。

迁移到新数据时，重新审视两件事。第一是坐标单位，原始 Demo 里坐标可能已经换算成平面米制单位，直接换成经纬度而不同步调整 sigma2，高斯核的衰减曲线完全失效，邻接矩阵退化成几乎全是零值，图卷积退化成本地 MLP。第二是时间粒度，5 分钟粒度的数据用 Kt=3、tw=12 没问题；换成小时粒度后，序列相关性变弱，tw 要缩短到 6 以下，否则模型会把过多的历史加权进当前时刻的预测。

另一个实用技巧是预测误差的自检脚本。每轮训练结束，除了验证 loss，额外算一个“最后一帧预测误差的移动平均”，也就是误差 EMA。如果这个值在 10 轮内变化超过两倍，说明数据里有明显的分布漂移，大概率是归一化参数和训练时的邻接多项式不是同一套，或者数据切分出了重叠。

从那以后，我每次拿到一份新的图时序数据，都强制走一遍“邻接矩阵可视化 → 归一化参数存档 → 三段式验证”的流程，再也没有出现验证集好、测试集翻车的情况。这套 PyTorch 版 STGCN 实现把上述流程都封装成了可以直接改参数运行的模块，当成个人学习的起点很合适。希望帮到你。

本文还有配套的精品资源，点击获取