ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

ASTGCN交通流预测源码复现:PEMS04/PEMS08实战指南

ASTGCN交通流预测源码复现:PEMS04/PEMS08实战指南

简介:面向交通流预测与图神经网络研究的深度学习源码包,基于AST-GCN(注意力时空图卷积网络)实现,适配PEMS04与PEMS08公开数据集。内含完整训练与测试流程,提供模型配置、数据预处理、评估指标等模块,适合有一定Python基础的研究者、研究生复现经典论文或开展交通预测方向的实验对比。包体共25个文件,核心为13个Python脚本,覆盖astgcn/mstgcn模型定义、train/test主程序、lib工具库(数据准备、utils、metrics);另含PEMS04/PEMS08配置文件、依赖清单与Dockerfile,以及两篇PDF文献(AAAI论文与PPT讲解)。压缩包20.71MB,目录结构清晰,按读README即可快速运行。目前已有1982人学习下载。通过源码可掌握时空图卷积的构建方法、注意力机制的应用以及交通数据集的标准化处理流程,便于后续扩展与改进模型。

1. ASTGCN 交通流预测:这份 PEMS04/PEMS08 源码包到底能跑出什么

做交通流预测的人应该都听过 ASTGCN,Attention Based Spatial-Temporal Graph Convolutional Networks,2019 年 AAAI 的文章。它把注意力机制和图卷积结合到一起,专门处理路网传感器的时间序列预测问题。这篇笔记要拆的这份资源,是 ASTGCN 的完整源码包,带着 PEMS04 和 PEMS08 两个数据集的配置,加上论文 PDF 和作者 PPT,基本就是照着论文复现的完整工具箱。适合已经在做时空序列预测、想对比 baseline 的研究生,也适合刚入门图神经网络、想找一个能跑通的交通预测项目下手的工程师。项目里模型代码、数据预处理、训练测试脚本、Docker 部署文件都齐了,不用自己东拼西凑。

2. ASTGCN 的模型结构:注意力加图卷积,为什么能抓交通流的时空依赖

2.1 时空注意力机制在交通预测里解决什么问题

交通流数据跟普通时间序列不一样,它有两个维度的依赖。空间上,一条路的车流量受周边路网影响,高峰期一环堵车,二环三环跟着慢,这是图结构上的空间依赖;时间上,早高峰的拥堵模式跟昨天、上周同一天高度相关,这是长周期的时间依赖。普通的 RNN 或者 CNN 处理不了这种非欧几里得空间的关系,所以 ASTGCN 的思路是用图卷积处理空间维度,用注意力机制自动学习不同时间段和不同传感器节点之间的影响权重。

这套源码里的核心模型文件是lib/model/astgcn.py,继承自mstgcn.py里的基础类。模型分成三个分支,分别处理最近一小时、过去一天、过去一周三个时间片。每个分支都走「时间注意力 + 空间注意力 + 图卷积 + 标准卷积」这条流水线,最后把三个分支的输出融合起来做预测。源码里model_config.py里定义了num_for_predict、num_of_vertices这些参数,改数据集的时候主要就是动这个文件。

2.2 配置文件和控制参数怎么理解

这份资源里配置文件的格式是.conf,用 ConfigParser 解析。PEMS04 和 PEMS08 各有一份,打开configurations/PEMS08.conf能看到核心参数:

[data] data_file = PEMS08.npz model_file = ASTGCN adj_file = PEMS08.csv num_of_vertices = 170 [training] batch_size = 32 epochs = 200 learning_rate = 0.001 num_of_weeks = 1 num_of_days = 1 num_of_hours = 1

参数含义拆开说。num_of_vertices = 170是 PEMS08 路网中的传感器节点数,PEMS04 是 307,这个值必须和邻接矩阵 CSV 的行列数对上,否则图卷积的拉普拉斯矩阵维度会直接报错。num_of_weeks、num_of_days、num_of_hours对应三个时间分支的输入长度,每个分支的输入形状是(batch_size, C, num_of_hours * steps_per_hour, num_of_vertices),steps_per_hour在data_preparation.py里默认取 12,也就是五分钟一个采样点。改这三个参数就是在调整模型看多长的历史窗口。

2.3 邻接矩阵和数据的组织方式

这里有个新手容易绕进去的点:数据文件是.npz,邻接矩阵是.csv,两者怎么配合。data_preparation.py里的load_dataset函数负责读PEMS08.npz,里面的data键形状是(num_samples, num_of_vertices, num_of_features),特征一般是流量、速度、占用率三列。邻接矩阵 CSV 是num_of_vertices x num_of_vertices的方阵,元素表示两个传感器节点之间的空间距离或连通性,ASTGCN 原文用的是高斯核加权邻接矩阵,源码里utils.py的calculate_laplacian函数把它转成拉普拉斯矩阵再喂给图卷积层。

数据划分在utils.py的data_split里,默认按 6:2:2 分成训练、验证、测试。这个比例和交通预测论文里的常见做法一致,刻意不打乱时间顺序,因为时间序列预测必须用连续时间段来验证,打乱了就变成纯回归问题,失去时序意义。

3. 把源码跑起来:环境配置、数据准备和训练全流程

3.1 环境搭建和依赖安装

资源包里有 Dockerfile 和 requirements.txt,两条路可以走。有 Docker 环境就直接构建镜像,没有就用 pip 装依赖。我一般是先看 requirements.txt 里锁了哪些版本再决定:

pip install -r requirements.txt python train.py --config configurations/PEMS08.conf

requirements.txt 里主要是 TensorFlow 1.x 系列、Keras、NumPy、SciPy、pandas。这里有个大坑:这份源码是 2019 年的,用的是 TensorFlow 1.14 或者 1.15,用 TensorFlow 2.x 直接跑会报错,因为tf.Session、tf.get_default_graph这些 API 全被移除了。官方给的 Dockerfile 里锁的是tensorflow/tensorflow:1.14.0-gpu-py3镜像,建议直接用这个,省得在依赖地狱里折腾。

3.2 数据预处理流程拆解

tools/data_preparation.py负责把原始数据切分成模型能吃的格式。核心函数是data_transform,它做的事是把原始的时间序列按照num_of_weeks、num_of_days、num_of_hours三个窗口滑窗采样。每次滑窗生成三个输入张量:week_data形状是(N, C, num_of_weeks * steps_per_hour, num_of_vertices),day_data和hour_data类似,只是时间步数不同。对应的标签是当前时间点之后num_for_predict步的流量值。

执行数据预处理需要先跑make_data.py,它会读取配置文件里的路径参数,把 PEMS08.npz 转换成训练、验证、测试三个 npz 文件放到data/PEMS08/目录下。这一步的输出直接决定训练时train.py里load_data能不能正确读到数据,顺序反了会看到Data file not found的错误。

3.3 训练入口和参数调优的落点

train.py是训练入口,它做的事情可以用下面这段逻辑概括:

config = ConfigParser() config.read(config_file) adj = load_adj(config["data"]["adj_file"]) data = load_data(config["data"]["data_file"]) model = ASTGCN(adj, config["model"]) model.train(data["train_x"], data["train_y"], data["val_x"], data["val_y"])

load_adj从 CSV 读邻接矩阵并计算拉普拉斯,load_data返回的train_x是一个字典,键是week、day、hour,值是对应的滑动窗口输入张量。ASTGCN类的train方法里,batch_size控制每轮迭代的样本数,learning_rate初始是 0.001,优化器用的 Adam,中间还有一个lr_schedule在训练后期自动衰减,通常到第 120 个 epoch 左右学习率会降到原来的十分之一。

3.4 自定义模型结构时的修改点

如果你不是为了复现,而是想在这个模型上改网络结构,重点看几个文件。lib/model/astgcn.py里的_build_graph方法定义了三个分支的完整计算图,想加注意力头数、改卷积核大小都在这里动。lib/model/mstgcn.py的MSTGCN类实现了基础的单分支时空卷积块,ASTGCN继承它再叠加注意力。lib/model/model_config.py里的ModelConfig类初始化了num_filter和kernel_size这些超参,我把num_filter从 64 改到 128 试过,PEMS08 上的验证损失能降一点,但训练时间几乎翻倍,值不值要自己权衡。

3.5 metrics.py 的评估指标怎么读

lib/metrics.py实现了三个交通预测标准指标,MAE、RMSE、MAPE,训练完的每个 epoch 都会输出验证集上的表现。最终测试时test.py会加载最优权重并算测试集上的三个指标。读结果的时候注意看 MAPE 的表现,ASTGCN 在 PEMS04 上是 8% 左右、PEMS08 上是 9% 左右的水平,如果你的实验环境跑出来的 MAPE 超过 15%,大概率是数据预处理阶段出了问题,不是模型问题。

4. 避坑指南:ASTGCN 源码包运行中的五个常见翻车现场

4.1 TensorFlow 2.x 跑出AttributeError: module 'tensorflow' has no attribute 'Session'

现象:按照 README 装了最新版 TensorFlow,python train.py启动后在第一轮训练就直接抛AttributeError。

原因:源码里大量使用tf.Session()、tf.placeholder、tf.get_default_graph,这些 API 在 TensorFlow 2.0 里被彻底移除。这份源码的基准环境是 TF 1.14,官方 Dockerfile 也锁在这个版本,用 2.x 跑是必挂的。

解决:两个办法,第一个是直接用仓库里的 Dockerfile 构建镜像,镜像名大概率是tensorflow/tensorflow:1.14.0-gpu-py3,docker build完docker run进去跑训练,省心。第二个是在 conda 里专门建一个 py3.6 + TF 1.14 的环境,pip install tensorflow==1.14.0,然后保证numpy版本不高于 1.16,因为 1.17 以上有些 C API 不兼容 TF 1.x。

4.2 GPU 显存 OOM 但代码逻辑没错

现象:模型开始训练,Epoch 1 刚跑完一个 batch 就报ResourceExhaustedError: OOM when allocating tensor。

原因:batch_size = 32加上三个时间分支的输入叠加,在 8GB 显存的卡上就是会爆。PEMS08 是 170 个节点,但num_filter是 64,每个分支的中间张量体积很大,我实测在 11GB 的 1080Ti 上勉强能跑。PEMS04 有 307 个节点,同样配置显存需求直接翻倍。

解决:先把batch_size降到 8 或者 16,我一般先降到 8 确认能跑通,再逐步往上加。如果还不行,把model_config.py里的num_filter从 64 改成 32,显存占用大概能省一半,代价是精度有轻微下降。另外看train.py里有没有设config.gpu_options.allow_growth = True,没设的话 TF 会默认把显存全占满。

4.3 换成 PEMS04 数据集后维度对不上

现象:把配置文件从PEMS08.conf换成PEMS04.conf,训练时报InvalidArgumentError: Incompatible shapes。

原因:PEMS04 的节点数是 307,PEMS08 是 170。配置文件里的num_of_vertices没改,或者改了但邻接矩阵 CSV 还是 PEMS08 的,拉普拉斯矩阵的维度跟输入数据的第二维对不上,矩阵乘法直接崩。

解决:换数据集时三个地方必须同步检查。第一,configurations/PEMS04.conf里的num_of_vertices是不是 307。第二,data/PEMS04/下是不是既有.npz数据文件也有对应的邻接矩阵 CSV。第三,模型结构如果改过,lib/model/model_config.py里的num_of_vertices初始化值也要跟着改,这地方不读配置文件而是直接读参数,是个容易漏的小坑。同样的道理,以后换任何一个新数据集,这三处都要一起改。

4.4 复现结果跟论文对不上,差了老远

现象:跑完测试集,MAE 在 20 以上,论文里是 17 左右,怎么调都降不下来。

原因:最常见的是数据预处理时没有做标准化。ASTGCN 原文对流量数据做了 Z-Score 归一化,源码data_preparation.py里用StandardScaler处理后再进模型,但很多人单独提取数据的时候把这步跳过了。另一个可能是训练提前停了,train.py默认跑 200 个 epoch,但中间用 early stopping 机制保存最优权重,如果 patience 设太小,模型还没收敛就停了。

解决:先确认data/PEMS08/目录下的数据是已经标准化过的版本还是原始数据,如果是从外部下载的原始.npz,需要先跑make_data.py完整走一遍预处理流程。再看训练日志里val_loss的曲线,如果最后几十个 epoch 还在持续下降,就调大epochs或者调大 early stopping 的 patience。另外检查num_for_predict是不是 12,即预测未来一小时,原文的实验配置是 60 分钟预测窗口,如果改成 3 或者 6,指标数值会变化但不会跟论文一致。

4.5 Docker 里跑训练发现挂载路径不对

现象:docker run进去之后python train.py报FileNotFoundError: data/PEMS08.npz,但明明镜像构建时已经 COPY 了数据。

原因:常见做法是用docker run -v把宿主机的项目目录挂载进容器,但挂载会覆盖镜像里对应路径的数据文件,如果宿主机目录里没有数据,容器里自然也没有。

解决:检查启动命令的挂载点,确保宿主机项目根目录下有完整的data/PEMS08.npz和configurations/PEMS08.conf。我习惯用绝对路径挂载:docker run -v /home/user/ASTGCN:/workspace/ASTGCN -it astgcn-image bash,进去之后cd /workspace/ASTGCN再跑。如果直接靠镜像里的 COPY 不挂载,也可以跑,就是每次改代码都要重新 build 镜像,特别麻烦。

5. 进阶验证技巧:用一张图快速确认模型有没有真正学到时空依赖

训练结束后别只盯着 MAE、RMSE 看,还要看预测结果的分布对不对。ASTGCN 源码包里有figures/model_result.png可以参考,但没有直接提供绘图脚本。我一般会写一段独立脚本,把测试集某几天的真实值和预测值画成双线对比图:

import numpy as np import matplotlib.pyplot as plt from sklearn.metrics import mean_absolute_error pred = np.load("output/prediction.npy").squeeze() true = np.load("output/truth.npy").squeeze() node_idx = 25 start, end = 0, 96 plt.figure(figsize=(12, 4)) plt.plot(true[start:end, node_idx], label="ground truth", linewidth=1.2) plt.plot(pred[start:end, node_idx], label="astgcn prediction", linewidth=1.2, linestyle="--") plt.legend() plt.xlabel("time steps (5-min interval)") plt.ylabel("traffic flow") plt.title(f"PEMS08 node {node_idx} traffic flow forecasting") plt.tight_layout() plt.savefig("figures/check_node25.png", dpi=150)

这段脚本做的事是提取output目录下保存的模型预测结果和测试集真值,挑一个节点画 8 小时的对比曲线。判断模型学得好不好有三条标准。第一,预测曲线和真实曲线的波峰波谷相位是否对齐,如果整体滞后一两个时间步,说明模型对时间依赖的建模偏弱,需要增大num_of_hours或者让注意力机制更聚焦近期数据。第二,预测值是否平滑到失真,如果曲线几乎是一条直线,说明模型退化成只预测均值,大概率是图卷积部分失效了,检查一下邻接矩阵是不是全零。

第三点也是我最常看的,同一个节点在工作日和周末的表现差异。交通数据的工作日和周末模式完全不同,如果模型在两类数据上都跟得上节奏,说明时空依赖确实学进去了。如果只在工作日表现好、周末预测明显偏差,可以在配置里把num_of_weeks调大,让模型有更长的周期记忆。从那以后我每次跑完训练都会先画这一张图再写进实验记录,强制自己用可视化验证模型质量,光看指标数字真的会骗人,因为高频率时间点上的 MAE 被低估时,曲线图上会暴露得一清二楚。希望帮到你,按这套流程走下来,这个源码包应该能顺利跑通并成为你后面改模型的一个稳妥基线。

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

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进