时空变换网络实战：交通流预测的Transformer方案与代码解析

简介：面向交通流预测研究及毕业设计、课程设计场景，这份Python实现的时空变换网络（ST-Transformer）源码包，整合了模型训练、验证、图卷积计算与单热编码等完整模块，并配有PEMSD7路网速度数据集，代码附有注释，逻辑清晰，新手也能快速看懂。压缩包共8个文件，包含6个Python脚本与2个CSV数据文件，脚本覆盖模型搭建、图层运算、数据预处理、训练与验证全过程，数据集提供邻接矩阵与速度流量序列，整体约484KB，轻量易部署。目前已有172人学习下载，适合需要快速复现时空预测基线、完成课程设计或毕业设计的高年级本科生与研究生。借助这份源码包可以掌握交通流数据清洗、时空依赖建模、模型评估的完整链路，省去自行整理数据和搭建框架的时间，直接聚焦算法调优与结果分析。

1. 时空变换网络和它的高分项目形态：这不是又一个LSTM

交通流预测里有个反直觉的场景：一个路口堵住了，相邻路口的流量往往要过二十分钟才出现连锁变化，而且堵点还会顺着路网往远处传。只靠LSTM管时间、只靠图卷积管空间，都没法把这个“跨时段、跨路段”的传播关系一次抓准，于是有了专门把自注意力搬进路网的时空变换网络。这种用Python实现、自带数据集的“高分项目”形态，正是这份方案要拆的东西——它建模什么、怎么跑通、指标怎么调、坑在哪里。时空变换网络把路网当作一张图，把每个时间步当作序列，用Transformer的自注意力同时在空间和时间两个维度上建模流量依赖。适合第一次接触时空图预测的读者：你拿到一个zip，想快速跑通并读懂每一行代码，这份笔记就是你的跟练地图。

2. 拆开时空变换网络：空间注意力和时间注意力两条主线

时空变换网络的输入不是一张普通的二维表格，而是一个五维张量X ∈ R^{B,N,T,C}：B是batch大小，N是路测点（传感器/卡口）数量，T是历史时间步，C是特征通道（流量、速度、占有率）。常见的实现把节点和时间拆成两条注意力主线，再串成一个“时空块”。理解这两条主线，后面调参和排错都有地方下手。

2.1 空间注意力：邻接矩阵是注意力里的“先验闸门”

空间维度要回答的问题是：某个节点的流量，受哪些邻居影响？最直接的做法是对每个时间步，把N个节点当作N个token，做一次标准的自注意力。但路网不是全连接图，两个相距十公里的节点本来就不该有强相关，所以注意力里必须加一个“先验闸门”——邻接矩阵。

常见做法是先把邻接矩阵做拉普拉斯归一化，再把归一化结果变成attention bias：有边相连的位置bias为0，不相连的位置填入一个很大的负数，让softmax算出来权重接近0。下面是构造bias的标准代码：

import numpy as np import torch def build_spatial_bias(adj, nodes): # adj: (N, N)，带权重或0/1邻接矩阵 # 拉普拉斯归一化：D^-1/2 @ A @ D^-1/2，让不同出度的节点量纲一致 deg = adj.sum(axis=1) + 1e-6 diag_inv_sqrt = np.diag(1.0 / np.sqrt(deg)) norm_adj = diag_inv_sqrt @ adj @ diag_inv_sqrt # 不相邻的位置塞一个很大的负数，softmax之后权重趋近于0 bias = torch.zeros(nodes, nodes) mask = torch.from_numpy((norm_adj > 0).astype(np.float32)) bias[mask == 0] = -1e9 return bias

这段代码有两个参数值得较真。第一，deg为什么要加1e-6：交通流邻接矩阵里有些传感器只有一条入边甚至没有边，直接归一化会除零，加一个小常数只是保数值稳定，不会改变拓扑结构。第二，-1e9这个负数有多大才够：softmax里只要比正常logit小一个数量级以上就能把权重压到接近0，实际工程里-1e9足够，但要注意如果后面接了float16训练，-1e9在fp16下会溢出成-inf，反而让梯度消失，这时改-1e4更稳妥。

用这一段bias，空间注意力的Q、K、V仍然来自输入特征，只是注意力分数在softmax之前被邻接矩阵“闸门”筛了一遍。这样模型学到的是“在路网约束下，哪些邻居的流量变化对我影响最大”，而不是NLP里那种完全由数据自由发挥的相关性。

2.2 时间注意力：位置编码和预测边界决定模型看到哪一步

时间维度要回答的问题是：历史哪几个时刻对预测未来最有用？对每个节点独立做时间步注意力，和Transformer在NLP里的用法基本一致，但有两个关键差异。

第一是位置编码。交通流不是纯文本，它有强周期性：早高峰、晚高峰、周末效应。常见的做法是用可学习位置嵌入，而不是NLP里默认的sin/cos。原因是一段30天的流量序列，周期长度不是固定常数，sin/cos的固定频率很难匹配采样不均匀的交通数据，可学习嵌入把“第几步”交给模型自己去归纳周期：

import torch import torch.nn asnn class LearnableTimeEmbedding(nn.Module): def __init__(self, seq_len, d_model): super().__init__() # 每个时间步一个独立的可学习向量，形状 (1, T, d_model) self.pos = nn.Parameter(torch.randn(1, seq_len, d_model)) def forward(self, x): # x: (B, T, N)，先映射到 (B, T, d_model) 再加位置 return x + self.pos

第二是预测边界。如果模型是自回归多步预测，每预测下一步要把上一步输出拼回输入，那注意力必须加因果掩码，让当前位置看不到未来；如果模型是单次输入T步、一次输出P步，那解码端每个输出位置能看到完整的历史输入，不需要因果掩码。拿到的源码是哪种结构，直接决定了你能不能看到“未来信息泄漏”问题：很多翻车现场就是自回归结构里漏了mask，训练时指标漂亮，测试时一塌糊涂。

时间注意力的代价是O(T²)的矩阵乘法。seq_len=12时无感，但有人把seq_len拉到48甚至96，显存就开始吃紧，这跟后面避坑章的OOM问题直接相关。

2.3 为什么选Transformer而不是纯GCN或纯LSTM：三个可验证理由

选型不是追新，而是有可验证的依据。第一是长程依赖：拥堵传播往往滞后半小时以上，图卷积堆两层感受野只有两跳邻居，LSTM在长序列上容易遗忘早期状态，注意力则允许任意两个时间步直接建边，理论上能把“半小时前的上游拥堵”和“现在的下游流量”连起来。第二是可解释性：多头注意力的权重可以直接画成热图，答辩或汇报时能说明“模型关注了哪条邻居路径”，纯GCN的黑匣子属性更强。第三是benchmark表现：在METR-LA、PEMS04这类公开交通流数据集上，时空注意力结构近几年基本占据前列，作为高分项目，评审对这种结构的认可度远高于堆一个双层的LSTM。

但这套结构的代价也很现实：数据量小的时候容易过拟合，显存消耗比GCN大一个量级，而且对数据预处理非常敏感。后面第三、四、五章分别解决“怎么跑通”“怎么调优”“翻车了怎么排查”。

3. 跑通源码：从解压zip到看到第一条训练曲线的四步

拿到一个“源码+数据集”的zip，先别急着读模型文件，按四步走：确认目录和数据文件完整性，做滑动窗口预处理，跑一个最小训练脚本，最后用一套默认参数验证模型真的在学。

3.1 目录、数据文件与入口函数：先建立文件地图

这类项目的目录结构大同小异，常见做法是这样：

project/ ├── main.py # 训练/测试统一入口 ├── config.py # 全部超参数 ├── models/ │ ├── sttn.py # 时空变换网络主体 │ └── layers.py # 注意力、MLP、位置编码 ├── utils/ │ ├── data_loader.py # 加载 npz、切滑动窗口 │ └── metrics.py # MAE / RMSE / MAPE └── data/ ├── metr-la.npz # 流量/速度数据 └── adj.npz # 邻接矩阵

main.py是入口，config.py是全部旋钮，模型文件和数据处理都在各自目录里。拿到zip先不要跳进sttn.py，先确认数据文件是好的。交通流公开数据集最常见的格式是npz，里面有两个key：data形状是(N, T, C)，adj形状是(N, N)。用几行命令验证：

import numpy as np d = np.load("data/metr-la.npz") print(list(d.keys())) # 期望 ['data', 'adj'] print(d["data"].shape) # 期望 (207, 34272, 1) 左右 print(d["adj"].shape) # 期望 (207, 207)

这一步如果shape对不上，后面所有窗口切分都是白做。另外注意np.load默认allow_pickle=False，有些老数据集打包时混入了Python对象，如果报错可以加allow_pickle=True，但这只是兜底，不代表数据没问题。

3.2 数据预处理：滑动窗口、归一化与防数据泄漏

交通流项目的预处理核心是滑动窗口：用历史T步预测未来P步。窗口会重叠，样本数相当大，METR-LA这种规模的数据切完能有几万条样本。下面这段是这类项目里最常见的数据处理代码：

import numpy as np def load_data(data_path, seq_len, pred_len, stride=1): raw = np.load(data_path) flow = raw["data"][..., 0].T # 取流量通道，转成 (T, N) adj = raw["adj"] # (N, N) # 按时间顺序切分，train 段占前 70% split_t = int(flow.shape[0] * 0.7) mean = flow[:split_t].mean(axis=0, keepdims=True) std = flow[:split_t].std(axis=0, keepdims=True) + 1e-6 norm_flow = (flow - mean) / std X, Y = [], [] for i in range(0, len(norm_flow) - seq_len - pred_len + 1, stride): X.append(norm_flow[i: i + seq_len]) # 历史 seq_len 步 Y.append(norm_flow[i + seq_len: i + seq_len + pred_len]) # 未来 pred_len 步 return np.stack(X), np.stack(Y), adj, mean, std

归一化这里有个新手最容易踩的坑：统计量必须只用训练段计算。如果对整段数据做Z-score，测试段的分布信息已经偷偷混进了归一化参数里，算出来的指标偏乐观，复现论文数据对不上往往就是这里出问题。还有std加1e-6是防止某节点流量全天为0时除零。

stride=1时窗口重叠严重，样本量、训练时间都翻倍；stride=5能显著提速，但会跳过一部分时序组合，指标通常略差一点。如果你只是验证代码能跑，用stride=5先跑通再换回1。

3.3 最小训练脚本：模型要能过拟合再谈泛化

很多新手拿到模型文件第一件事是直接把整个训练脚本扔上去跑几千个epoch，结果loss不下来就怀疑人生。我一般会先写一个最小训练脚本，目标只有一个：让模型在一个batch上过拟合。能过拟合，说明前向反向、loss回传的链路是通的。

import torch import torch.nn.functional as F def train_minimal(model, x, y, steps=200): """x: (B, T, N), y: (B, P, N)，只喂同一个batch""" opt = torch.optim.Adam(model.parameters(), lr=1e-3) for step in range(steps): opt.zero_grad() pred = model(x) # (B, P, N) loss = F.l1_loss(pred, y) # L1 即 MAE，和评价指标对齐 loss.backward() opt.step() if step % 50 == 0: print(f"step {step}, loss {loss.item():.4f}")

注意这里的loss直接用L1 loss，因为最后评价指标是MAE，训练目标跟评价指标对齐，比用MSE省去一次“换了loss重新调参”的往返。正常情况下一两百步之内loss应该明显下降，如果一直卡在某个常数不动，先查两个地方：一是输入有没有归一化，二是模型输出层有没有接反了维度，输出应该是(B, P, N)，很多实现容易写成(B, N, P)。

跑通最小脚本之后再切回完整训练流程，这时候加梯度裁剪和混合精度是顺理成章的事：

torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)

交通流数据偶尔会有尖峰，梯度裁剪能防止一个异常的batch把权重顶飞。显存紧张就加torch.cuda.amp自动混精，注意前面提到的-1e9bias在fp16下的溢出问题。

3.4 默认参数表：一套能起步的配置

过拟合测试通过后，用下面这套参数起跑：

这套参数在PEMS04（307个节点）上通常能跑到MAE 17~22，在METR-LA（207个节点）上大约2.5~3.5，是公开benchmark的常见水平。每个zip里的模型细节不同，数值不完全一样，但这不重要——重要的是先有一个“能跑出数”的起点，然后按第四章的思路去调。如果连这个数量级都达不到，先回头检查数据预处理，而不是急着换模型结构。

4. 指标与调参：把MAE从能跑压到能用的三层手段

训练跑通只是第一步，高分项目需要指标能打。这一章讲透三个评价指标怎么选，训练曲线怎么看，以及从三个层面各动一个参数的具体顺序。

4.1 三个评价指标：MAE、RMSE、MAPE各自的坑

MAPE是翻车重灾区。夜间流量接近0时，预测值只差1，MAPE却能贡献百分之几百的误差。常见的处理办法是加epsilon，或者只统计流量高于某个阈值的样本，并在报告里写上口径，否则答辩时评委一句“MAPE怎么这么高”就能让你卡住。RMSE则相反，它放大峰值误差，如果你的模型总是低估早高峰流量，RMSE会比MAE更早暴露问题，所以调参时两个指标要一起看。

4.2 训练曲线往哪看：loss降指标涨才是真的在学

训练日志要同时记录train loss和val MAE（或val loss），两个曲线对照才有意义。正常的信号是：前二三十个epoch两者同步下降，之后train loss继续降但val指标开始反弹，这就是过拟合的拐点。如果一开始train loss就降得很慢，val指标几乎不动，先别调参，先确认模型是否在训练集上能过拟合，这一步见第三章的检验方法。

一个特别有效的baseline是“均值回归”：用历史窗口的均值当作预测值，算一下它的MAE。如果时空变换网络连这个baseline都打不过，问题大概率不在模型容量，而在数据预处理或实现bug。很多调参调到最后发现是归一化写错，白白浪费几个晚上，这种血泪经验就是要靠这个baseline来兜底。

4.3 三层可调参数：结构层、训练层、数据层各动一个

调参建议按数据层、结构层、训练层的顺序来，因为数据层决定信号上限，结构层决定拟合能力，训练层只是让优化过程不崩。

数据层第一个动的是seq_len。从12拉到24，观察val MAE是否有明显收益。交通流的自相关窗口一般在一小时以上，但更长历史不总是有用，如果收益小于0.05就回退，因为seq_len翻倍意味着时间注意力的显存开销翻四倍。pred_len同理，从12拉到36时误差会明显上涨，这是数据本身的预测极限，不是模型问题，别在这一步上死磕。

结构层优先动num_layers和d_model。PEMS04这种中小规模数据集，层数超过4层几乎必然过拟合，d_model从64加到128则要盯显存。有个常见误区是一味加层数，实际上交通流数据的有效感受野有限，层数收益边际递减很快。

训练层最后动。学习率用warmup+cosine退火是这类项目的标准配置，前几个epoch让权重稳定再给满学习率；batch_size如果从64降到16能明显稳住梯度，说明训练本身不稳定，但更大的batch反而可能提前过拟合。梯度裁剪clip_grad_norm=5.0直接写死在训练脚本里当保险丝用，不碍事。

最后提醒一句：这些参数只在公开数据集上有意义。如果你要换成自己采集的卡口数据，输入不再是(N,T,C)的标准npz，还需要先把数据整理成节点流量矩阵，邻接矩阵也要按路网拓扑生成，而不能用经纬度距离简单截断，这一步的影响远大于后面所有的调参。

5. 避坑指南：时空变换网络跑不动的六个常见问题

这一章把最常翻车的六个问题按“现象→原因→解决”写清楚。多数问题不是模型难，而是数据、归一化、掩码这些外围环节在捣鬼。

5.1 数据与指标：泄漏、除零与注意力失效

坑一：复现指标对不上论文或README的数值。

现象：别人写MAE 17.5，自己跑出来21+，调了好几天压不下去。原因排在第一位的是数据泄漏：归一化统计量用了全样本，或训练集验证集随机切分，导致验证阶段混进了未来信息。解决方法是严格按时间顺序切分，train占前70%、val占中间10%、test占最后20%，归一化统计量只从train段计算，再固定随机种子复跑一遍。如果差得还是大，就把第五章其他坑逐个过一遍，八成还在外围。

坑二：MAPE打印出来是一个三位数，比如588%。

现象：MAE还算正常，一汇报MAPE就非常难看。原因是夜间流量接近0，相对误差被分母放大。解决方法是计算MAPE时只统计流量高于某个阈值（比如每小时5辆）的样本，分母再加一个epsilon，同时在报告里注明这个口径，否则这个数字没法跟别人对比。

坑三：把注意力热图画出来，所有节点权重几乎一样。

现象：attention map一片均匀，模型退化成“全路网平均流量”预测器。原因要么是空间bias没生效（比如邻接矩阵本身就是全0），要么是bias加的位置不对，被后面的归一化层洗掉了。解决方法是先检查build_spatial_bias输出里对角线和非对角线的数值差异，再看模型forward里bias有没有真正参与attention score的计算；更直接的办法是跑一个“把邻接矩阵当唯一输入、不学任何可学习参数”的对照组，如果这个对照组输出都是一团浆糊，说明数据或代码链路有问题。

5.2 训练与显存：OOM、损坏文件与预测曲线变平

坑四：batch_size=64一上来就CUDA out of memory。

现象：显存直接爆掉，报错在QK^T矩阵乘法附近。原因是多头注意力里Q @ K.T的形状是B×heads×T×N×N（或B×heads×N×T×T），N=307时这一项就能吃满一张12G卡。解决方法有三个：把batch降到8~16；用梯度累积模拟大batch；优先确认空间注意力和时间注意力是分两个子层实现的，而不是在一个矩阵里一次性算完，后者会同时爆两个维度的平方开销。

坑五：zip解压后np.load抛EOFError或中途停止输出。

现象：数据集文件解压后load到一半就报错，shape打印不出来。原因是zip包在传输或解压过程中不完整，分卷包没合并全。解决方法是先对比压缩包在zip内的原始大小和解压后文件的大小，少了字节就是缺包；重新完整解压一遍，然后用np.load(..., mmap_mode='r')读一次，能读完再走预处理流程，不要带着坏数据调半天模型。

坑六：预测曲线整体比真值小一圈，峰值全部被削平。

现象：MAE指标看着不错，一画图发现早高峰的预测值明显偏低。原因是L1 loss对中位数友好，对极端峰值惩罚不足，模型学会了“别猜太高，猜平了平均损失最小”。解决方法是把loss改成Huber loss，或者对流量高于75分位的样本加大权重，让模型对峰值更敏感。调完之后MAE可能略涨但RMSE会降，这个trade-off要看你汇报时更想强调哪个指标。

6. 进阶验证：花半小时做四个消融，把项目说服力拉满

模型跑通、指标调好之后，还差最后一步：证明“时空变换”这四个字里，空间和时间两个部件都真的在起作用。方法就是消融实验，半小时能出结果，但答辩和评审的说服力提升一大截。

最常见的消融组合是四组对比：完整STTN、去掉空间bias（等价于纯时间Transformer）、去掉时间位置编码（等价于时间上无序的注意力）、换成两层LSTM当baseline。前两组的差距说明空间注意力有没有利用路网结构，后两组说明时间建模的贡献。代码改动很小：

# 消融 2：把空间 bias 全置 0，等价于关掉空间先验 model.spatial_bias = torch.zeros_like(model.spatial_bias)

跑完把四个MAE列成一张表，整个项目的逻辑链就闭合了。如果消融结果表明去掉空间bias后指标反而更好，那说明你的数据里节点间的空间依赖本来就不强，这时候不要硬凹，坦白把这个发现写进报告里，同样是加分的。

有条件的话，再把测试集高峰时段（比如早8点到9点）的注意力权重导出来，画成以拥堵节点为中心的热图。你能直观看到模型把注意力投给了路网拓扑上的哪些邻居——这个证据比任何文字描述都有说服力。早年我跑这类项目只盯指标，后来发现模型几乎没用到邻接结构，等于把Transformer用成了带位置编码的均值回归，从那以后我习惯每次实验先加一个全零邻接矩阵的对照组，既防自己盲目乐观，也让评审看到你确实知道自己在做什么。希望你跑通的不只是一个能出数的zip，而是一套能讲清楚“为什么有效”的方案，希望帮到你。

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