Python实现时空变换网络：融合图注意力与Transformer的交通流预测

简介面向交通流预测的时空变换网络ST-TransformerPython实现适合毕业设计、期末大作业或课程设计阶段使用该项目以PEMSD7交通流数据为基础源码包含模型定义、图卷积网络模块、训练与验证流程代码注释较完整新手也能按注释逐步读懂。压缩包共8个文件包含6个Python脚本和2个CSV数据集大小仅484KB部署轻量Python脚本覆盖主模型、网络层、图卷积模型、训练、验证及数据预处理等模块CSV数据文件提供交通流预测所需的输入与标签数据便于完整复现实验。目前已有172人学习浏览。下载后解压即可运行目录结构清晰可快速定位关键代码既适合对照学习时空变换网络的完整实现流程也可直接作为毕设或课程设计的高分候选代码在此基础上进行改进与二次开发。1. 交通流预测的时空变换网络一个把路网结构和时间序列一起建模的 Python 项目交通流预测有个反直觉的事实凌晨三点的历史数据和早高峰的历史数据对当前时刻的参考价值完全不同而你当前路段的拥堵程度往往取决于五百米外上游道路十几分钟前的状态。时间依赖和空间依赖纠缠在一起这正是用 Python 实现时空变换网络Spatio-Temporal Transformer Network做交通流预测的核心动机。这个 zip 包里的项目把图注意力网络与 Transformer 的时间自注意力整合进同一个模型配好了可直接读取和训练的 METR-LA 数据集以及完整的训练、评估、可视化脚本。适合正在做智能交通方向课题、准备毕业设计或者想快速验证时空预测方案可行性的从业者。2. 为什么 LSTM 不够用把路网结构搬进注意力机制才算“时空”2.1 交通流数据的三个固有特性每一层都对应一个模型组件交通流数据不是普通的时间序列它有三个被反复验证的特性直接决定模型结构选型。第一是周期性。城市路网的车速、流量有明显的日周期和星期周期早高峰七点到九点、晚高峰五点到七点规律性强到可以用历史同期数据做基线预测。但周期本身会漂移节假日、天气、临时管制都会让“同期”失效。所以模型不能只记住周期还得有能力在周期基础上修正。第二是空间相关性。上下游道路的交通状态是强相关的——上游路口开始拥堵十几个采样间隔对应五到十分钟内下游路口的速度一定会跟着掉。这种相关性不局限于相邻路段快速路和辅路之间、匝道与主路之间也存在长距离弱相关。用欧氏距离衡量道路关系没有意义必须用路网拓扑。第三是突发扰动。交通事故、临时施工、信号灯配时调整会产生一个“冲击波”沿着路网传播并在时间轴上衰减。这种扰动没有周期性也不严格符合空间相邻关系但它恰恰是预测误差最大的来源。这三个特性叠加导致了两类经典模型的失效。ARIMA 只建模单点时间序列根本不看路网LSTM 虽然能抓长期依赖但每个节点的嵌入是独立训练的节点之间的特征共享只能靠样本里偶然出现的相关性没有任何结构先验。这里需要的是一个能把图结构显式写进计算图的模型而时空变换网络正好同时解决了“时间注意力”和“空间注意力”两件事。2.2 空间维度的骨架图注意力层怎么写空间维度的核心是把路网当成一张图N 个传感器是节点传感器之间的距离关系是边。常见的建图方式有两种距离阈值法和 KNN后面第 3 章会展开。拿到邻接矩阵之后空间聚合层用图注意力机制计算“每个节点应该从邻居那里聚合多少信息”。下面是最小可运行的图注意力层代码。import torch import torch.nn as nn import torch.nn.functional as F class GraphAttentionLayer(nn.Module): 单层图注意力输入 [B, N, F]输出 [B, N, F]。 adj 是 0/1 矩阵1 表示两条道路相邻。 def __init__(self, in_dim, out_dim, adj, dropout0.1): super().__init__() self.adj adj # [N, N]邻接矩阵 self.W nn.Linear(in_dim, out_dim, biasFalse) self.a nn.Parameter(torch.zeros(2 * out_dim, 1)) self.dropout nn.Dropout(dropout) self.leaky_relu nn.LeakyReLU(0.2) def forward(self, x): B, N, _ x.shape h self.W(x) # [B, N, out_dim] h_i h.unsqueeze(2).expand(B, N, N, -1) h_j h.unsqueeze(1).expand(B, N, N, -1) e torch.cat([h_i, h_j], dim-1) # 拼接起点和终点特征 e torch.matmul(e, self.a).squeeze(-1) e self.leaky_relu(e) # 非邻居的注意力分数直接置为负无穷softmax 后接近 0 mask self.adj.unsqueeze(0).expand(B, N, N) 0 e e.masked_fill(mask, -1e9) alpha F.softmax(e, dim-1) # 归一化注意力权重 out torch.matmul(alpha, self.dropout(h)) return out, alpha逻辑说明代码先把每个节点的特征经过线性层变换随后将起点特征和终点特征拼接通过向量 a 映射成标量注意力分数。关键操作是masked_fill把非邻居位置的分数压到负无穷这样 softmax 之后这些位置的权重趋近于零模型就只会在真实路网邻居之间分配注意力。返回的alpha可以拿来做可视化分析模型到底“看了”哪些道路。参数说明in_dim是输入特征维度实际项目里通常是 1只有速度或 3速度、流量、占用率out_dim建议设为输入维度的 2 到 4 倍dropout在注意力权重上生效能防止少数路段主导整个聚合结果。如果你想让图注意力支持多跳邻居把邻接矩阵换成 A 的 k 次幂即可但 k 超过 3 之后收益很小噪音反而变大。2.3 时间维度的骨架自注意力与残差怎么配合空间层处理完单个时间切片的图结构接下来需要让每个传感器节点“看”到自己的历史序列。这里用 Transformer 的MultiheadAttention做时间维度的自注意力。它和 LSTM 最大的区别是LSTM 按顺序逐步读入序列当前时刻的输出依赖前一个隐状态自注意力则一次性看到整个时间窗口用查询和键的相似度直接决定“哪几个历史时刻对当前最重要”。这对交通流很关键因为今天早高峰的形态可能与三天前更像而 LSTM 很难跨越几十步去精确对齐这种结构性相似。class STTNBlock(nn.Module): 时空变换块先做空间图注意力再做时间自注意力各带一个残差。 def __init__(self, feat_dim, adj, num_heads4): super().__init__() self.spatial GraphAttentionLayer(feat_dim, feat_dim, adj) self.temporal nn.MultiheadAttention( feat_dim, num_heads, batch_firstTrue ) self.norm1 nn.LayerNorm(feat_dim) self.norm2 nn.LayerNorm(feat_dim) def forward(self, x): # x: [B, T, N, F] B, T, N, F x.shape # 空间维度把时间当成样本共享图注意力权重 x_spatial, _ self.spatial(x.reshape(B * T, N, F)) x_spatial x_spatial.reshape(B, T, N, F) x self.norm1(x x_spatial) # 时间维度对每个节点独立做自注意力 x_temporal, _ self.temporal( x.permute(0, 2, 1, 3).reshape(B * N, T, F) ) x_temporal x_temporal.reshape(B, N, T, F).permute(0, 2, 1, 3) return self.norm2(x x_temporal)逻辑说明这段代码把整个 STTN 块拆成两个阶段。第一阶段把 [B, T, N, F] 重排成 [BT, N, F]让每个时间切片独立走一遍图注意力这样空间上的邻居聚合在每一时刻都会发生。第二阶段把维度重排成 [BN, T, F]把每个节点的时间序列独立送入自注意力。两个阶段都带残差连接和 LayerNorm目的是让深层堆叠时不发生梯度消失或表示塌缩。参数说明feat_dim在输入输出保持一致方便残差相加num_heads一般取 4 或 8头数太多会导致小数据集上注意力分布过散。实际项目里常见做法是堆叠两个 STTNBlock再跟一个两层 MLP 输出未来 horizon 个时间步的预测值。要注意的是这个简化版本没有写时间位置编码。严格来说Transformer 需要位置编码才能区分“三分钟前”和“三十五分钟前”但交通流数据本身带有强周期许多公开实现里不加位置编码仍然能收敛加了反而容易过拟合这一点在不同数据集上有玄学成分建议你加一个可开关的位置编码层做对照实验。3. 把 zip 解压到能跑通环境配置、目录规划与 METR-LA 数据预处理3.1 包内目录结构与 Python 环境准备拿到 zip 之后先不要急着装环境。第一步是解压并确认目录结构一个合格的时空预测项目至少应该有 models、data、utils、train.py、test.py 这几块。常见布局见下面的目录树。STTN/ ├── config.yaml ├── train.py ├── test.py ├── models/ │ ├── sttn.py │ └── attention.py ├── utils/ │ ├── data_loader.py │ ├── metrics.py │ └── visualization.py └── data/ ├── metr-la_adj.csv └── metr-la_flow.csvconfig.yaml负责集中管理所有超参数data_loader.py负责读取原始 CSV 并生成批量样本models目录放网络结构train.py和test.py分别是训练和评估入口。拿到陌生源码包时我一般先读config.yaml和train.py的启动参数再看data_loader.py的样本组织方式最后才进入模型细节。这个顺序能最快暴露“这个项目吃的是哪种数据格式”。环境建议用 conda 建独立环境Python 版本不低于 3.8。依赖主要是 torch、numpy、pandas、pyyaml如果你的 python 环境里还没有可以用下面两条命令一次性装好。# 建环境并激活 conda create -n sttn python3.8 -y conda activate sttn # 安装核心依赖torch 版本按自己的 CUDA 版本选 pip install torch numpy pandas pyyaml tqdm参数说明python3.8是保守选择代码里如果用了matmul和masked_fill3.8 以上版本都兼容CUDA 版本不确定时优先装 CPU 版 torch 跑通小 batch确认逻辑无误再换 GPU 版能省掉大量环境排查时间。切忌一上来就装最新版 torch有些旧源码里batch_first参数依赖特定版本 API出问题反而难定位。3.2 邻接矩阵构建距离阈值法和 KNN 的取舍METR-LA 数据集包含洛杉矶 207 个传感器节点每个节点每 5 分钟记录一次平均车速原始数据是两份 CSV一份是路由信息文件记录传感器经纬度和相互距离另一份是交通速度时间序列。构建图的第一步是把节点距离矩阵转成 0/1 邻接矩阵。业界最常用的两种方法分别是距离阈值法和 KNN下面是距离阈值法的实现。import pandas as pd import numpy as np def build_adjacency(dist_path, threshold0.03): 根据传感器距离矩阵生成邻接矩阵。 距离单位是经纬度坐标下的相对距离METR-LA 常用阈值 0.03。 dist pd.read_csv(dist_path, index_col0).values adj (dist 0) (dist threshold) # 自己到自己的边去掉避免自循环 np.fill_diagonal(adj, False) return adj.astype(np.float32)逻辑说明距离阈值法的核心假设是“只有离得够近的传感器才存在空间依赖”。dist 0排除对角线自环dist threshold保留短距离节点对。METR-LA 数据的坐标经过投影处理0.03 这个阈值大约对应三公里以内的相邻路段能在大约每个节点平均 6 到 8 个邻居的密度下运行。参数说明threshold是最需要调的参数。设得太小图会退化成若干孤立的子图空间信息传不出去设得太大全国节点全连在一起注意力矩阵变成稠密矩阵显存和计算量都会爆炸。KNN 方案则反过来——每个节点固定连接最近的 K 个邻居K 常在 8 到 16 之间。两者对比距离阈值法的优点是物理含义清晰缺点是不同路网密度下同一阈值的效果波动大KNN 的优点是保证每节点最少有一条边比较稳定缺点是远程强相关路段可能被截断。3.3 滑动窗口样本生成与归一化train、val、test 怎么切才不泄密拿到原始时间序列 [N, T_total] 之后下一个步骤是构造滑动窗口样本。假设要预测未来 12 个时间步60 分钟需要过去 12 个时间步60 分钟作为输入。样本构造和归一化必须一起做顺序错了模型基本白训。def create_dataset(flow, seq_len12, horizon12, train_ratio0.7, val_ratio0.1): flow: [N, T_total]按时间顺序切分训练/验证/测试集。 返回归一化后的 X/Y以及用于反归一化的 scaler。 N, T flow.shape scalers [] flow_norm np.zeros_like(flow) # 逐传感器做 MinMax 归一化统计量只来自训练区间 train_end int(T * (train_ratio val_ratio)) # 先留出测试区间 for i in range(N): train_slice flow[i, :train_end] mn, mx train_slice.min(), train_slice.max() scalers.append((mn, mx)) flow_norm[i] (flow[i] - mn) / (mx - mn 1e-6) X, Y [], [] for t in range(train_end - seq_len - horizon): X.append(flow_norm[:, t: t seq_len].T) # [seq_len, N] Y.append(flow_norm[:, t seq_len: t seq_len horizon].T) X np.stack(X, axis0) Y np.stack(Y, axis0) # 按时间顺序切分不 shuffle 测试集 train_size int(len(X) * train_ratio / (train_ratio val_ratio)) val_size int(len(X) * val_ratio / (train_ratio val_ratio)) return (X[:train_size], Y[:train_size], X[train_size:train_size val_size], Y[train_size:train_size val_size], X[train_size val_size:], Y[train_size val_size:], scalers)逻辑说明代码里最关键的一点是归一化统计量只从训练区间计算。很多新手图省事对整段序列做 MinMax结果测试区间的最大最小值泄漏进训练阶段模型在测试集上的 MAE 会虚低 10% 到 20%看起来“分数很高”拿到真实场景立刻翻车。另一个关键点是样本按时间顺序切分测试集永远在最末段模拟真实部署时“用历史预测未来”的状态。参数说明train_ratio0.7、val_ratio0.1是交通流基准实验的常见配置。seq_len12对应一小时历史窗口horizon12对应一小时预测长度。如果你想做多步预测对比把 horizon 改成 3、6、12 分别对应 15 分钟、30 分钟、60 分钟预测即可但要注意 horizon 增大后样本数量会减少模型难度也显著提升。4. 训练脚本调用与评估指标把模型跑到能出 MAE、RMSE、MAPE4.1 训练启动命令与关键超参数数据准备好之后训练入口一般封装成train.py通过命令行参数覆盖 config.yaml 里的默认值。以下是典型的启动命令。python train.py \ --dataset METR-LA \ --seq_len 12 \ --horizon 12 \ --batch_size 64 \ --lr 0.001 \ --epochs 100 \ --seed 2024参数说明batch_size64在 207 个节点、12 步历史的输入尺寸下显存占用大约 3 到 4 GB单卡 3080 级别如果你的显卡只有 6 GB 显存建议降为 32lr0.001配合 Adam 优化器是时空图模型的常用起点如果 loss 前 10 个 epoch 不下降优先把 lr 降到 0.0005而不是去调网络结构seed2024必须固定否则同一份代码两次训练的结果可能差出 0.3 个 MAE后面完全没法比较模型改动的好坏。训练循环内部还有两个细节值得注意梯度裁剪和早停。时空变换网络堆叠多层后梯度范数偶尔会冲到很大导致 loss 突然变成 NaN梯度裁剪能兜住这个风险。早停则避免过拟合——验证集 loss 连续 15 个 epoch 不下降就保存当前最优权重并退出。4.2 训练循环与带 mask 的损失函数交通传感器数据里存在缺失值METR-LA 原始 CSV 里缺失位置常以 0 填充。计算 loss 时必须把这些位置排除否则模型会被迫学习“把缺失位置预测成 0”整体指标被拉低。下面是带 mask 的 MAE 损失和训练循环的最小实现。import torch import torch.nn as nn import torch.nn.functional as F def masked_mae(pred, true, null_val0.0): 剔除缺失值后的 MAE。pred/true: [B, horizon, N] mask (true ! null_val).float() denom mask.sum().clamp(min1.0) return (torch.abs(pred - true) * mask).sum() / denom def train_one_epoch(model, loader, optimizer, device): model.train() total_loss 0.0 for xb, yb in loader: xb, yb xb.to(device), yb.to(device) out model(xb) # 输出形状 [B, horizon, N] loss masked_mae(out, yb, null_val0.0) optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() * len(xb) return total_loss / len(loader.dataset)逻辑说明masked_mae先把预测值和真实值绝对误差算出来乘上 mask只对有效位置求和最后除以有效位置总数。如果不对缺失值做 mask模型会把大量精力浪费在拟合缺失位置导致真实路段上的误差反而上升。clip_grad_norm_把梯度范数限制在 5 以内防止个别样本产生巨大梯度把整个参数空间冲乱。参数说明null_val0.0是 METR-LA 缺失值填充约定。换到 PEMS-BAY 数据集时同样是 0。如果数据源不同需要先看一眼 CSV 里缺失值的实际占位符是 NaN 就要先转成 0 或单独传 mask 矩阵。4.3 评估指标MAE、RMSE、MAPE 分别暴露什么问题交通流预测论文里几乎必出三个指标MAE平均绝对误差、RMSE均方根误差、MAPE平均绝对百分比误差。它们的含义完全不同只看一项会得出错误结论。import numpy as np def compute_metrics(pred, true): pred/true: [B, horizon, N]反归一化后的实际速度值单位 km/h。 mae np.mean(np.abs(pred - true)) rmse np.sqrt(np.mean((pred - true) ** 2)) # 真实速度接近 0 时除以小量避免 MAPE 异常 mape np.mean(np.abs((pred - true) / (true 1e-6))) * 100 return {MAE: mae, RMSE: rmse, MAPE: mape}逻辑说明反归一化这一步容易漏掉。模型输出的值是 [0, 1] 区间的归一化结果直接算指标会得到一个看起来漂亮的数字但根本不能和论文里的基准比较。正确做法是用第 3 章保存的每个传感器独立的 (min, max) 把预测值和真实值同时还原成 km/h再调用compute_metrics。指标解读MAE 告诉你平均每个时刻每个传感器偏差多少 km/h比如 MAE3.5 表示平均误差 3.5 公里每小时RMSE 因为先平方再开方对大误差更敏感如果 RMSE 明显大于 MAE说明存在少数预测崩掉的样本MAPE 是百分比方便不同路网之间横向比较但真实车速接近 0 时 MAPE 会爆表这就是加上 1e-6 小量的原因。实际项目里我一般以 MAE 作为主要调参依据RMSE 作为稳健性参考。5. 避坑指南归一化泄漏、显存溢出与 MAPE 爆表的 5 条排查记录5.1 现象训练 loss 震荡不收敛预测值全部接近历史均值原因数据里个别传感器的速度量纲差异大逐传感器归一化时用了全局统计量导致梯度在不同节点之间互相抵消。更常见的原因是学习率初始值太大Transformer 的注意力 softmax 在训练初期对输入扰动很敏感。解决先把学习率降到 0.0005 跑 20 个 epoch观察 loss 是否稳定下降同时确认data_loader里每一批样本都执行了归一化而不是只在预处理阶段归一化一次。如果仍不收敛检查位置编码是否数值不稳定必要时改用LayerNorm放在注意力之前。5.2 现象验证集 MAE 比训练集还低模型表现反常地好原因这是典型的数据泄漏。常见来源有两个——一是归一化统计量用了全量数据测试区间的 min/max 提前暴露二是滑动窗口切分后没有按时间顺序切 train/test而是随机 shuffle导致相邻时间窗口的样本同时出现在训练集和测试集测试时模型实际上“见过”几乎同样的数据。解决严格按时间切分数据集前 70% 训练、10% 验证、20% 测试归一化统计量只用训练段计算测试集任何情况下不做 shuffle。检查方法很简单打印出训练集最后一个样本的时间戳和测试集第一个样本的时间戳确认它们不相邻即可。5.3 现象batch_size 从 32 调到 64 后直接 CUDA Out of Memory原因图注意力层在计算注意力分数时生成了 [B, T, N, N] 的稠密矩阵。207 个节点时每个样本的注意力矩阵约 170 万元素batch 翻倍后显存占用翻倍加上 Transformer 的中间激活很容易超出显存。解决优先减小 batch_size如果一定要保持大 batch把邻接矩阵改为稀疏表示只对每个节点的前 K 个邻居计算注意力K8 时显存占用能降到原来的十分之一。还有个偏方是把整个模型切成 FP16 混合精度训练显存占用再降一半代价是训练稳定性变差需要配合梯度缩放。5.4 现象同一份代码、同一份数据两次训练结果差出 0.5 个 MAE原因训练脚本没有固定随机种子PyTorch 在 CUDA 上的某些算子默认使用非确定性算法尤其是scatter和atomicAdd这类归约操作每次前向传播的数值都有微小扰动。解决训练入口固定torch.manual_seed(seed)、np.random.seed(seed)并设置torch.backends.cudnn.deterministic True。同时建议在每次实验前把 config.yaml 连同代码版本一起记录否则你根本分不清结果差异是代码改动带来的还是随机性带来的。5.5 现象MAE 正常但 MAPE 出现几百甚至上千的异常值原因MAPE 的公式里分母是真实值。深夜某些传感器车速极低真实值可能是 5 km/h 甚至 1 km/h绝对值误差 2 km/h 对应百分之几十的误差数值一放大就不正常。更极端时真实值为 0 直接导致除零。解决计算 MAPE 时给分母加 epsilon比如true 1e-6同时设置一个阈值只统计真实速度大于阈值比如 10 km/h的样本。这是交通流预测评估里的惯例论文里都会注明“MAPE is computed on samples with actual speed 10 km/h”。如果你不做这个过滤报告里的 MAPE 会明显偏高和公开基准没法比。6. 把模型迁移到 PEMS-BAY 数据验证泛化能力与调参的一手经验METR-LA 跑通只是第一步。真正让这个时空变换网络有价值的事是在第二个数据集上验证泛化能力。PEMS-BAY 是加州湾区 325 个传感器的高速路网数据采样间隔同样是 5 分钟结构上和 METR-LA 几乎一致迁移成本很低。python train.py \ --dataset PEMS-BAY \ --num_nodes 325 \ --seq_len 12 \ --horizon 12 \ --batch_size 64 \ --lr 0.001 \ --epochs 100 \ --seed 2024迁移时需要改三处num_nodes从 207 改成 325data/目录下换成 PEMS-BAY 的邻接矩阵和流量 CSV。这两份数据同样公开可得预处理代码完全复用第 3 章的build_adjacency和create_dataset不需要改模型结构。我的经验是模型在 METR-LA 上能收敛在 PEMS-BAY 上几乎一定能收敛因为两者都是高速路网传感器数据分布接近。如果换到城市道路数据则需要重新调阈值。训练完之后记得做一件事把 15 分钟、30 分钟、60 分钟三个预测步长的指标同时画出来。交通流预测的规律是预测步长越长误差越大且增量不是线性的——从 15 分钟升到 30 分钟MAE 通常增加 0.5 到 1从 30 分钟升到 60 分钟MAE 再增加 1 到 1.5。如果你的模型出现 60 分钟误差急剧放大多半是训练时 horizon 步输出端的梯度传播不均可以考虑把输出端从单层 MLP 换成两层堆叠结构。最后记得用训练好的模型画一条 24 小时预测对比曲线横轴是时间纵轴是车速真实值和预测值画在同一个图上。你会直观看到早高峰和晚高峰附近的误差最大——这不是模型缺陷而是交通流的固有不确定性。我自己的习惯是先盯 MAE 调结构再盯峰值时段看曲线形态指标数字只是门槛曲线才决定模型能不能拿出手。如果峰值时段预测偏差能压住这份时空变换网络项目在大多数评审眼里就站得住脚。希望帮到你。本文还有配套的精品资源点击获取