简介面向交通流预测与时空图网络研究的 Python 深度学习资源基于注意力时空图卷积网络AST-GCN实现聚焦 PEMS04 与 PEMS08 两个公开交通数据集上的流量预测任务。适合具备一定 PyTorch 基础、希望复现注意力时空图卷积网络或开展交通预测实验的研究者与开发者。 资源共 25 个文件以 rar 压缩包发布体积约 20.71MB。核心为 13 个 Python 源码文件覆盖模型定义、训练测试流程、数据预处理与评估工具另附 2 个模型配置文件、2 份论文资料会议原文与配套演示文档、说明文档以及 Docker 环境配置可帮助读者快速搭建可复现的实验环境。整体目录按源码、配置、论文和工具模块组织查找与引用方便。 目前已有 1982 人浏览学习。通过这套代码可以完成从交通数据加载、模型训练到指标评估的完整流程尤其适合缺少完整工程代码的论文复现场景参考价值明确。1. ASTGCN 是什么为什么 PEMS04/PEMS08 上的交通预测总拿它当基线做交通流量预测的人迟早会撞上 ASTGCN 这个名字。它是 2020 年提出的 Attention-based Spatio-Temporal Graph Convolutional Network把图卷积和时间卷积组合起来专门处理路网传感器数据PEMS04 有 307 个传感器PEMS08 有 170 个每 5 分钟记一条流量、速度、占用率你要做的是用过去 1 小时预测未来 1 小时。你可能以为这类任务用 LSTM 就能解决但 LSTM 看不到路网拓扑而 ASTGCN 用图卷积把“传感器之间的空间关系”直接建模进去在公开数据集上把 MAE 往下压了一截。之后几年几乎所有交通预测论文都拿它当对比基线。这篇文章我按自己复现时的实际路径来写PEMS 数据怎么读、邻接矩阵怎么用 python 构建、模型结构怎么拆、训练时哪些参数值得调、最容易踩的坑在哪。适合刚接触时空图神经网络的研究生也适合想拿公开数据练手图模型的工程师。2. 从数据到训练样本PEMS08 数据格式与 python 构建邻接矩阵2.1 PEMS04 和 PEMS08 的 npz 文件里到底有什么从公开渠道下载的 PEMS04、PEMS08 数据集解压后核心是一个.npz文件。用np.load打开后常见的有data、label、adj三个键。data和label的维度一般写作(样本数, 节点数, 特征数)PEMS04 节点数是 307PEMS08 是 170特征数固定为 3流量、平均速度、占用率。时间粒度为 5 分钟一天 288 个时间片所以一个月的样本量就是节点数乘 288 乘天数数量级在几十万条。.npz里的adj键在多数复现中并不直接拿来用因为它的组织和论文描述的邻接矩阵不一致我一般只把它当作参考真正用来构图的是传感器坐标文件。这里有一个容易误导新手的点很多人以为拿到data和label就可以直接丢进DataLoader实际上还差两步。第一步data是按原始时间顺序存储的矩阵切片你需要自己切滑动窗口第二步图卷积需要的邻接矩阵不会自己出现必须根据传感器经纬度用距离公式构造。两步做完数据才算准备好。import numpy as np raw np.load(PEMS08.npz) data raw[data] # 假设 shape 为 (样本数, 170, 3) label raw[label] # shape 同 data print(data.shape, label.shape) # 输出类似 (34272, 170, 3) (34272, 170, 3)代码说明raw[data]取得的是整个时间序列的观测label是与data对齐的未来序列。不同来源的数据集在样本数上会有差别但你只需要关注后两维节点数和特征数。如果打印出来的特征维度不是 3先检查是不是读错了文件PEMS04 和 PEMS08 的标准特征就是流量、速度、占用率这三项。这里不直接打印adj是因为它的组织方式因版本而异直接打印往往是一堆看不出含义的索引对。2.2 滑动窗口怎么切用过去 12 步预测未来 12 步模型输入输出对齐是复现 ASTGCN 时第一个容易出错的环节。它的标准设定是用过去 12 个时间片也就是 1 小时的数据预测未来 12 个时间片的流量。为什么是 12 不是 6因为交通流量有明显的日周期性1 小时窗口能覆盖“早高峰前 1 小时”这类完整的上坡段而 15 分钟、30 分钟预测虽然也能做但窗口太短时模型很容易退化成“拿上一时刻外推”。ASTGCN 的论文里分别预测了 3 步、6 步和 12 步复现时我建议先只做 12 步这是最常被对比的设定。滑动窗口的构造逻辑很直接从时间轴第 0 个位置开始取[i, i12)作为输入取[i12, i24)作为目标步长为 1。这样相邻样本之间有大量重叠数据量会被放大到和原始时间片数量几乎一样多。需要注意归一化要在这之后做而且是先用训练集拟合均值方差再应用到验证集和测试集顺序反了会造成数据泄漏后面避坑章节会细说。def make_windows(seq, history12, horizon12): X, Y [], [] for i in range(len(seq) - history - horizon 1): X.append(seq[i : i history]) Y.append(seq[i history : i history horizon]) return np.stack(X), np.stack(Y) # data 形状: (样本数, 节点数, 特征数)先转成 (样本数, 特征数, 节点数) 再切 data_t data.transpose(0, 2, 1) # (S, 3, 170) X, Y make_windows(data_t) Y Y[..., 0:1] # 只保留流量作为预测目标代码说明make_windows里的len(seq)是总的原始时间片数量history和horizon分别控制输入输出长度。切完后X的形状是(样本数, 12, 3, 170)Y的形状是(样本数, 12, 1, 170)。注意我在data上先做了transpose把特征维度放到第二维这是为了后续输入 PyTorch 卷积层时直接以(B, C, T, N)使用省去在模型内部反复换轴。预测目标只取流量这个特征因为 PEMS04/08 的标准评价指标 MAE、RMSE、MAPE 都是针对流量计算的。2.3 用 python 构建邻接矩阵高斯核、距离阈值、自环一个都不能少PEMS 数据集没有直接给你“邻接矩阵”只有传感器坐标。ASTGCN 论文里用的是高斯核加权两个传感器之间的距离越近边的权重越大距离超过一定阈值就断开。这一步直接决定图卷积能“看”到多远的邻居阈值设得太大整张图全连通图卷积退化成全局平均设得太小每个节点只和自己相连模型学不到空间信息。我一般用距离标准差作为高斯核的分母阈值设为 0.8这样保留下来的边大约是距离最近的 10% 到 20%。import numpy as np def build_adj(coords, sigmaNone, threshold0.8): # coords: (N, 2)每一行是 (lat, lon) n len(coords) dist np.zeros((n, n)) for i in range(n): # 实际项目中用 haversine 公式算经纬度距离这里示意欧氏距离 dist[i] np.sqrt(((coords - coords[i]) ** 2).sum(axis1)) if sigma is None: sigma dist.std() adj np.exp(-(dist ** 2) / (sigma ** 2)) adj[adj threshold] 0 np.fill_diagonal(adj, 1.0) # 自环保留避免图卷积丢当前节点 return adj # 使用adj build_adj(coords) # shape (170, 170) 或 (307, 307)代码说明sigma不传时用距离矩阵的标准差这样阈值对 PEMS04 和 PEMS08 都适用不用因为节点密度不同而反复调。threshold的作用是把权重小于 0.8 的边直接砍掉这个值在 PEMS08 上表现稳定但换到别的城市数据集时最好画一下权重分布再定别照抄。fill_diagonal把对角线设成 1这一步很多人会漏漏掉之后切比雪夫图卷积做矩阵乘法时当前节点自己的特征贡献会被邻居淹没模型训练初期 loss 会明显偏高。到这里训练样本和邻接矩阵都准备好了。下一步是把模型结构拆开看ASTGCN 到底怎么把时间维和空间维揉在一起的以及每一步张量的形状变化这决定了你写代码时维度对不上该去哪里排查。3. 模型结构拆解时空注意力、图卷积与张量形状的演进3.1 三个时空模块级联的骨架从 (B, C, T, N) 到预测输出ASTGCN 的整体结构可以理解成“三个时空模块首尾相接再接一个全连接输出层”。每个时空模块内部都做同样四件事先算时间注意力再算空间注意力接着做图卷积最后做时间卷积。输入张量的形状全程保持(B, C, T, N)不变其中B是 batch sizeC是特征通道数T是时间步长 12N是节点数 170 或 307。很多初学者抄代码失败就是没搞懂这个形状约定在某个permute之后把T和N弄反了。我按论文思路写一个简化版时空模块骨架方便你看清楚数据流向。import torch import torch.nn as nn class STBlock(nn.Module): def __init__(self, in_c, out_c, K3, heads3): super().__init__() self.spatial_attn SpatialAttention(in_c, heads) self.cheb_conv ChebConv(in_c, out_c, K) self.temporal_conv nn.Conv2d(out_c, out_c, (3, 1), padding(1, 0)) self.glu nn.Conv2d(out_c, out_c, 1) self.norm nn.LayerNorm([out_c, 12, 170]) self.residual nn.Conv2d(in_c, out_c, 1) if in_c ! out_c else nn.Identity() def forward(self, x): # x: (B, C, T, N) attn self.spatial_attn(x) # (B, N, N) h torch.einsum(bctn,bnm-bctm, x, attn) h self.cheb_conv(h) # (B, out_c, T, N) h self.temporal_conv(h) # (B, out_c, T, N) h self.glu(h) * torch.sigmoid(h) # 门控线性单元 return self.norm(h self.residual(x))代码说明einsum这一行把空间注意力矩阵应用到每个时间片的特征上等价于对每个时刻的节点特征做一次图注意力加权。temporal_conv用的卷积核是(3, 1)只在时间维上滑动不跨节点保证空间信息不被卷积核混掉。glu之后的sigmoid乘法是 GLU 门控它的作用是让模型自己决定每一时刻的信息保留多少这在交通流量这种有早晚高峰突变的序列上比纯 ReLU 稳定。LayerNorm对[out_c, T, N]做归一化注意N用了 170这是写死 PEMS08 的写法换成 PEMS04 要改成 307或者干脆把N从输入形状里动态取出来。3.2 空间注意力和时间注意力为什么传感器之间要动态加权图卷积本身用的是固定的邻接矩阵但 ASTGCN 觉得这不够早晨 8 点的拥堵可能从收费站 A 传到收费站 B晚上 6 点的拥堵却可能从 C 传到 A固定边权无法表达这种随时间变化的影响。所以它在图卷积之前加了一个空间注意力模块输出一个N×N的权重矩阵对每个时刻重新分配邻居的重要性。这个矩阵不是从邻接矩阵复制来的而是根据当前输入的特征实时计算出来的相当于给图卷积装了一个动态路由器。时间注意力的逻辑类似它作用在T×T维上捕捉的是时间片之间的关系。比如今天是周三早上 8 点那今天 7:45 的流量和昨天 7:45 的流量高度相关模型通过时间注意力给这两个时刻分配更高的权重。这两类注意力都采用了多头设计论文里默认头数是 3。头数不是越大越好我试过 8 个头训练时间变长但指标几乎没有变化PEMS 这类单变量预测任务 3 个头足够。实现时有一个显存上的坑空间注意力矩阵形状是(B, N, N)PEMS04 的N307一个 batch 64 张卡会直接吃掉大半显存后面避坑章节会给出具体解法。3.3 切比雪夫图卷积为什么 K3 是默认值拉普拉斯要先标准化图卷积是 ASTGCN 的空间建模核心它不直接用邻接矩阵做一层 GCN而是用切比雪夫多项式展开到 K 阶相当于让每个节点聚合 1 到 K 跳邻居的信息。K3是论文里的默认值含义是每个传感器最多能看到三跳以内的路网关系。K 设得更大理论感受野更广但超过 5 之后参数增多、训练变慢而且多跳信息高度重叠边际收益很小。实现切比雪夫卷积前邻接矩阵必须先做拉普拉斯标准化再把特征值缩放到[-1, 1]区间否则矩阵连乘会导致数值不稳定训练到一半 loss 变成 nan。import torch import torch.nn as nn def rescale_laplacian(adj): d adj.sum(axis1) 1e-6 d_inv_sqrt np.diag(d ** -0.5) lap np.eye(adj.shape[0]) - d_inv_sqrt adj d_inv_sqrt # 特征值缩放: L 2*L/lambda_max - I lam np.linalg.eigvalsh(lap).max() return 2 * lap / lam - np.eye(adj.shape[0]) class ChebConv(nn.Module): def __init__(self, in_c, out_c, K): super().__init__() self.weights nn.ParameterList([ nn.Parameter(torch.randn(in_c, out_c) * 0.05) for _ in range(K) ]) def forward(self, x, lap): # x: (B, C, T, N), lap: (N, N) B, C, T, N x.shape x x.permute(0, 2, 1, 3).reshape(B * T, C, N) # (BT, C, N) x0 x x1 torch.bmm(lap.unsqueeze(0).expand(x.shape[0], -1, -1), x) out torch.bmm(x0.transpose(1, 2), self.weights[0]) # 0 阶 out out torch.bmm(x1.transpose(1, 2), self.weights[1]) # 1 阶 for k in range(2, len(self.weights)): xk 2 * torch.bmm(lap.unsqueeze(0).expand(x.shape[0], -1, -1), x1) - x0 out out torch.bmm(xk.transpose(1, 2), self.weights[k]) x0, x1 x1, xk return out.reshape(B, T, N, -1).permute(0, 3, 1, 2)代码说明rescale_laplacian函数只需要在训练前调用一次把邻接矩阵转成标准化拉普拉斯之后所有 batch 共用。ChebConv的 forward 里先算 0 阶和 1 阶再用递推公式x_k 2 * L * x_{k-1} - x_{k-2}算高阶项这是切比雪夫多项式的基本递推关系。需要注意的是torch.bmm乘法里lap的unsqueeze(0).expand只是为了适配 batch 维度实际不会把矩阵复制进显存。这里的权重形状是(in_c, out_c)意思是对每个节点做一次线性变换如果输入输出通道不同还得在这个模块外面套一个通道对齐的卷积。这版代码是教学用的最简实现真正跑实验时建议在这个基础上确认梯度回传路径重点看x0、x1的复用有没有破坏计算图这是我复现时排查过最久的地方。4. 训练与超参让 ASTGCN 在 PEMS 数据上收敛的最小配置4.1 数据加载与训练循环早停比调更多超参更值钱ASTGCN 的训练循环和普通 PyTorch 模型没有本质区别但有几个细节直接影响结果优化器用 Adam初始学习率 0.001损失函数用 MSE指标报告用 MAE、RMSE、MAPE训练集、验证集、测试集按 6:2:2 划分注意按时间顺序切不要随机打乱。我习惯先把训练集跑 100 个 epoch每 10 个 epoch 保存一次 checkpoint同时记录验证集 MAE如果连续 10 个 epoch 没有下降就早停。from torch.utils.data import DataLoader, TensorDataset import torch.optim as optim dataset TensorDataset( torch.FloatTensor(X_train), torch.FloatTensor(Y_train) ) loader DataLoader(dataset, batch_size64, shuffleTrue) model ASTGCN(num_nodes170, in_channels3, out_channels64, K3, heads3) optimizer optim.Adam(model.parameters(), lr0.001) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) best_mae float(inf) patience 0 for epoch in range(100): model.train() for x, y in loader: pred model(x) # (B, 1, 12, 170) loss torch.nn.functional.mse_loss(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() val_mae evaluate(model, X_val, Y_val) # 自定义验证函数 if val_mae best_mae: best_mae val_mae torch.save(model.state_dict(), best_astgcn.pth) patience 0 else: patience 1 if patience 10: break scheduler.step()代码说明ASTGCN是 3.1 和 3.3 中模块的组装体完整代码较长核心是输入(B, 3, 12, 170)、输出(B, 1, 12, 170)。StepLR每 10 个 epoch 把学习率减半到第 30 个 epoch 之后学习率已经很小模型开始精细调整这一步不加也能收敛但加上之后 MAPE 通常能再降 0.3 到 0.5 个百分点。早停的patience我设 10PEMS 这类数据噪声大验证指标经常连续 5 个 epoch 不动然后突然下降太早停会错过最优解。evaluate函数用 MAE 做判断但保存 checkpoint 时也会把 RMSE、MAPE 一起打印出来方便和论文对比。4.2 六个必调超参参考值、调参顺序与为什么别动它们ASTGCN 能改的超参不少但真正值得动手的只有六个切比雪夫阶数 K、注意力头数 heads、输出通道数 out_channels、batch size、学习率、dropout。下表是我在 PEMS08 上跑过多次后觉得比较稳的起点新手按这个配置先跑通再逐项调整。参数参考值调参方向注意事项K切比雪夫阶数3先试 2 和 4K4 提升很有限训练时间明显变长heads注意力头数3通常不用动8 头会显存爆炸收益微弱out_channels64可以试 12864 到 128 指标提升约 0.1翻倍显存batch_size64显存小就 32PEMS04 用 64 需要 10G 以上显存learning rate0.001降到 0.0003 再看0.01 直接发散dropout0.3过拟合就加到 0.5数据量接近百万一般不会严重过拟合调参顺序上我一般建议先固定 K3、heads3、out_channels64只调 batch size 和学习率保证模型先收敛然后再试 K4最后才动 dropout。优先调数据预处理而不是模型结构很多项目把时间注意力或门控单元删掉指标掉一截又加了回来这是结构层面的改动成本很高。先确认数据没问题再怀疑模型。4.3 指标怎么算PEMS08 上的合理水平是多少MAE、RMSE、MAPE 是交通预测的标准三件套。MAE 是绝对误差均值RMSE 放大了大误差的惩罚MAPE 看相对误差百分比。三者都要算但判断模型好坏时我主要看 MAE 和 MAPE。流量特征在原始数据里是“每 5 分钟通过的车辆数”PEMS08 的流量均值在几十的量级所以 ASTGCN 论文报告的水平大致是 1 小时预测时 MAE 在 20 上下MAPE 在 10% 上下你复现时不要死盯绝对数字因为数据预处理、样本划分方式不同数字波动很大。def evaluate_metrics(pred, true): pred pred.cpu().detach().numpy() true true.cpu().detach().numpy() mae np.abs(pred - true).mean() rmse np.sqrt(((pred - true) ** 2).mean()) mape np.mean(np.abs((pred - true) / (true 1e-6))) * 100 return mae, rmse, mape代码说明MAPE 计算时在分母加了1e-6这是为了防止流量为 0 的时段除零。PEMS 数据里有夜间低流量时段某些传感器可能连续 1 小时流量为 0不加这个 epsilon 会导致 MAPE 出现几百的异常值。验证时建议按预测步长分别打印指标比如第 1 步的 MAE 和第 12 步的 MAE前者通常只有后者的一半这能帮你判断模型是不是在“短时记忆”上偷懒。如果你的模型第 1 步和第 12 步的 MAE 差距小于 10%多半是标签构造出了问题预测值在时间轴上错位了。5. 避坑清单复现 ASTGCN 时最容易翻车的 5 个位置5.1 归一化在切分之前做验证集泄漏进训练集现象训练集 MAE 一路下降验证集指标却在高位震荡而且换随机种子结果波动很大。原因对整个data矩阵先做了(x - mean) / std均值和方差包含了验证集、测试集的信息模型相当于提前看到了未来数据的分布特征。解决先按时间顺序切出 train、val、test再在 train 上调用scaler.fit用同一组参数 transform 另外两个集合。5.2 邻接矩阵对角线为 0图卷积丢了自己现象loss 能收敛但预测曲线比真实曲线平滑很多峰值全被削平。原因构建邻接矩阵时没有把对角线设成 1切比雪夫图卷积的 0 阶项被邻接矩阵的一阶项稀释当前节点自己的特征贡献变得很小。解决在build_adj里加一行np.fill_diagonal(adj, 1.0)或者在模型里对输入的邻接矩阵单位化。这个坑很隐蔽因为 loss 不会报错指标也只差 0.5 左右但画图时一眼就能看出来。5.3 标签错位预测序列比真实序列晚了一个时间片现象预测曲线和真实曲线形态很像但整体向右平移了 5 分钟MAE 却意外地低。原因滑动窗口切错了索引把[i12-1, i1211)当成了目标模型学到的是“把上一时刻的流量复制过来”这个廉价解法。解决切完窗口后随机抽 10 条样本打印X[i, -1]和Y[i, 0]确认预测起点和输入终点至少相隔 5 分钟更严格的验证是计算预测序列与真实序列的滞后相关性峰值出现在 lag0 才对。5.4 显存溢出PEMS04 的 307 节点把空间注意力撑爆现象batch_size 设 64forward 到空间注意力时报 CUDA out of memory。原因空间注意力矩阵是(B, N, N)PEMS04 的 N307、B64 时光这个矩阵就是 64×307×307×4 字节约 96MB再乘多头和梯度显存直接翻车。解决先把 batch_size 降到 32如果还不行就把训练切换成混合精度torch.cuda.amp可以把显存占用砍半最后兜底的手段是把注意力头数从 3 降到 1牺牲少量精度换训练稳定性。5.5 MAPE 异常大罪魁祸首是夜间零流量现象MAE 正常RMSE 正常但 MAPE 高达几百甚至上千。原因夜间部分传感器流量为 0真实值为 0 时预测值只要不是 0相对误差就是无穷大。解决在 MAPE 计算时对真实流量加一个小 epsilon或者直接只统计一天中 6 点到 22 点的 MAPE。业界论文里通常不会公开这些细节所以别人报 10% 的 MAPE 时先问清楚他有没有过滤低流量时段。6. 从复现到可用迁移到新数据集与验证模型是否学对的 3 个技巧6.1 从 PEMS04 迁到 PEMS08只需要改五个地方换数据集时不必重写模型但要逐项检查五个变量节点数N307 变 170、邻接矩阵坐标文件、距离阈值、拉普拉斯预处理、标签维度。最容易漏的是坐标距离单位PEMS04 和 PEMS08 的坐标文件单位可能不一致直接套同一个sigma会得到密度完全不同的图。我的检查顺序是先打印邻接矩阵的稀疏度确认每行平均有 5 到 15 个非零元素再开始训练。几个配置对比时把命令行入口封装成 argparse 参数用 codex 跑深度学习实验脚本批量扫描会更省事。6.2 验证模型有没有学对先画高峰时段曲线模型训练完别急着看指标先画三张图早高峰 7 到 9 点、晚高峰 17 到 19 点、夜间 0 到 2 点的预测对比。如果高峰时段预测曲线比真实曲线平滑很多说明模型偏向平均预测空间注意力没起作用如果低谷时段误差和高峰时段一样大说明模型没学到日周期时间注意力大概率失效。我会再跑一个最简单的 baseline拿上一时刻流量作为预测值ASTGCN 的 MAE 必须显著低于这个 baseline否则整个模型没有比朴素外推更聪明。6.3 我的收尾习惯先跑 1/4 数据验证再跑全量现在接手新数据集我的固定动作是先抽出 1/4 数据、把N改成对应节点数、batch 降到 16跑通一个 epoch 确认 shape 和梯度没问题再上全量数据。这一步能省下半天排错时间。最后提醒一句ASTGCN 是 2020 年的模型后来出现的 DSTAGNN、PDFormer 等模型在 PEMS 上有更好的指标但 ASTGCN 作为基线的价值仍然存在——你新提出的模块到底有没有用得有它垫底才知道。希望这些踩坑经验能帮你少走点弯路。本文还有配套的精品资源点击获取