简介一套基于Python的时空变换网络ST-Transformer交通流预测完整实现面向智慧交通、深度学习和时空数据挖掘方向的学习者与开发者可用于解决城市道路车流量短时预测与拥堵趋势分析问题。压缩包共9个文件以6个Python脚本为核心分别对应模型定义ST_Transformer.py、图卷积与注意力层实现GCN_models.py、layers.py、训练与验证流程train.py、validation.py及数据预处理One_hot_encoder.py另含2个PEMSD7高速公路交通流数据集文件CSV格式和1份README说明文档便于快速上手整体大小约451KB。目前已有387人学习下载。借助该资源读者可完整体验从原始交通数据加载、时空特征编码到模型训练、效果验证的完整管线并理解时空卷积与注意力机制如何协同提升预测精度同时基于公开真实数据集可直接运行对比实验或替换自有数据开展二次开发为相关课题研究和工程应用提供可复用的参考。1. 交通流预测里的时空变换网络解决的不只是“下一个时刻堵不堵”你在做城市路网短时交通流预测时最深的体会通常是单点时序模型在早晚高峰延续、事故疏通、潮汐路段上总是慢半拍。LSTM 把每个检测器当成独立序列图卷积能把路网结构用起来但它们各自只在“时间”或“空间”一个侧面上建模。时空变换网络把两者放在同一套特征空间里做变换每一层输入的是一张【时刻 × 路口】的交通流矩阵输出既带上相邻路段的分布信息也保留整段时间上的变化模式。放在路口级流量预测、区域拥堵趋势判断里它能回答“上游排队会不会在 15 分钟后传到我这条路段”这类纯时序模型回答不了的问题。这篇笔记用 Python 源码加真实路网的数据组织方式把一套可跑的时空变换网络从数据到训练全程讲一遍新手可以照着复现熟手可以直接拿走调参数。2. 把时空变换网络拆成三块空间图、时间窗与变换层2.1 交通流预测为什么不是单点时序问题常规检测器每 5 分钟上报一次车流量。A 路口堵了影响要隔 510 分钟才会传导到相邻的 B 路口再进一步扩散到更远的 C 路口。这种空间级联效应是道路拓扑结构带来的而且有时效性距离越远特征到达得越晚、越衰减。纯 LSTM 或 LightGBM 这类按单点历史值建模的方案天然看不到“邻居路口此刻发生了什么”所以它们对突发拥堵的响应总是慢。反过来只做空间建模也不行——流量还有日周期性、早晚高峰形态和事件后恢复过程这些全落在时间维度上。所以这类问题的标准建模对象是一个三维张量X [N, T, F]。N 是路网节点数一般取道路交叉口或检测器点常见的公开规模是 100300 个T 是历史时间步数F 是每个节点的特征数通常是流量、速度、占有率中的一种或几种。预测目标是未来 H 个时间步的流量矩阵 Y [N, H]。时空变换网络要做的就是在这个三维张量上同时完成“邻居信息怎么汇总”和“历史信息怎么沿时间轴演化”两件事。2.2 空间模块选型GCN、图注意力或邻接矩阵直乘法空间模块的选择通常有三种按实现成本从低到高排列第一种是图卷积 GCN把邻接矩阵 A 和输入特征做一次归一化聚合。最常见的公式是 Z σ(D^{-1/2} Â D^{-1/2} X W_s b)其中 Â A II 是自环矩阵D 是度矩阵。这一步等价于对每个节点做邻居特征的加权平均权重由道路连接关系决定。它的好处是参数少、训练稳定缺点是只能表达静态拓扑。第二种是图注意力 GAT把邻居权重变成可学习的注意力系数。适合邻居重要性动态变化的场景但实现时要小心多头注意力的拼接维度和 mask 处理代码量多一截。第三种是自适应邻接矩阵。先用可学习参数做空间依赖的隐式度量再与真实邻接矩阵加权融合。这个方法在数据量充足时效果好但对小数据集容易过拟合。我一般建议第一版先上 GCN。交通流预测里一个节点的流量受其邻居影响有明确的拓扑先验GCN 的归纳偏置正好够用GAT 和自适应矩阵都属于锦上添花等 GCN 跑通、确认网络本身没问题后再换不迟。2.3 时间模块的设计边界为什么 Transformer 要处理时间维度时间维度建模有 LSTM、一维卷积、Transformer 编码器三条路。LSTM 的问题不是精度而是慢——批量训练几百个节点时双向 LSTM 的显存占用和推理延迟都比较高。一维卷积快但它对长距离时序依赖的建模能力弱感受野要靠堆层数来扩。Transformer 的注意力机制天然覆盖整个时间窗口而且能同时看到前后文比较适合“早高峰前 30 分钟的特征一直影响到当前时刻”这类场景。注意这里说的“时间变换”和 NLP 里的原生 Transformer 有一个关键差异交通流数据不是离散 token 序列每个时间步都对应一个连续数值向量而且时间步之间的相关性是单调衰减的。所以直接套原版 Transformer 并不理想更常见的做法是把多头注意力作为变换算子叠加残差和层归一化并把位置编码换成可学习的时间嵌入。这个结构在实现上仍然叫 Transformer 编码器层只是把输入从 token 序列换成时间切片序列。2.4 三种输入形状和输出形状先把张量维度写清楚很多刚上手的同学在写网络时第一个被绊倒的就是张量形状。设定一批输入 x 的原始形状是 [B, T, N, F]B 是 batch sizeT 是历史窗口N 是节点数F 是特征数。按模块不同需要做 reshape空间 GCN 层工作时它只关心节点维度和特征维度所以把 x 变成 [B*T, N, F]做完空间聚合再变回 [B, T, N, F]。时间变换层工作时它只关心时间维度和特征维度所以把 x 变成 [B*N, T, F]做完时间变换再变回原形。输出层做一个线性映射把 [B, T, N, F] 在时间维度上聚合最后得到 [B, H, N, 1]H 是预测步数。阶段张量形状说明输入样本[B, T, N, F]一批原始时空窗口空间 GCN 输入[B*T, N, F]把时间并进 batch时间注意力输入[B*N, T, F]把节点并进 batch预测输出[B, H, N, 1]未来 H 步每个路口的流量这个形状切换是整个源码里最容易埋 bug 的地方后面避坑章节会单独展开。3. 数据集与源码落位把路网检测数据切成交叉时空窗口3.1 原始交通流数据结构和邻接矩阵构建拿到一套交通流数据集第一件事不是写模型而是确认数据字典里有没有三张表节点表、流量时间序列、邻接关系。流量时间序列最常见的存储格式是 CSV 或 NPY行代表时间点列代表检测器 ID也就是一个 [T, N] 的矩阵。邻接关系有两种表达一种是直接给出 [N, N] 的邻接矩阵另一种是给出两两检测器之间的路网距离按距离阈值或 K 最近邻生成。如果只有经纬度或距离常见的邻接矩阵构建方式是这样import numpy as np def build_adjacency(distances, k8, sigma210.0, epsilon0.5): # distances: [N, N]两两检测器之间的物理距离矩阵 n distances.shape[0] adj np.zeros((n, n)) for i in range(n): dist distances[i] idx np.argsort(dist)[1:k1] # 去掉自身取前 k 个最近邻居 weights np.exp(-dist[idx] ** 2 / sigma2) adj[i, idx] weights adj[idx, i] weights # 保证对称 adj[adj epsilon] 0.0 return adj参数说明k 控制每个节点连接多少个邻居经验范围是 510城市路网通常取 8sigma2 是高斯核的带宽决定权重随距离衰减的速度太大会让所有邻居权重相近太小则只有极近节点有贡献。epsilon 是弱连接剪枝阈值低于它的边直接清零目的是减少噪声。距离矩阵的度量单位建议统一成公里或米否则 sigma2 完全没有物理含义纯靠试。3.2 用滑动窗口生成训练样本的 Python 代码流量时间序列本身是二维的 [T, N]。生成监督学习样本要做的是滑动窗口切分每个样本取连续 Tp 步作为输入紧接着的 H 步作为标签。注意窗口滑动步长通常设为 1 或 3步长为 1 时样本量最大但相邻样本高度重叠会增加训练耗时步长为 3 时样本量约减为三分之一训练更快精度损失通常不大。def make_samples(flow, his_len12, horizon3, step1): # flow: [T, N] 的流量矩阵 # 返回 x: [样本数, his_len, N] 和 y: [样本数, horizon, N] samples_x, samples_y [], [] total flow.shape[0] for start in range(0, total - his_len - horizon 1, step): end start his_len x flow[start: end] # [his_len, N] y flow[end: end horizon] # [horizon, N] samples_x.append(x) samples_y.append(y) return np.stack(samples_x), np.stack(samples_y)逻辑说明start 从 0 开始步长取 1 时数据利用率最高his_len 和 horizon 是超参数而不是网络参数前者决定模型看多长的历史后者决定预测多远。默认 his_len12 对应 5 分钟粒度下的一小时horizon3 对应预测未来 15 分钟。如果是预测早高峰特征可以酌情把 his_len 加到 24让模型看到 2 小时前的状态。3.3 均值方差归一化与反归一化的正确顺序交通流数据有很明显的日内波动凌晨流量接近 0早高峰接近峰值直接喂给网络会让损失函数被大数值样本主导。常见做法是用训练集统计量做 Z-Score 归一化然后网络输出必须反归一化回真实流量值再算指标。from sklearn.preprocessing import StandardScaler scaler StandardScaler() # 注意只在训练集上 fit验证集和测试集用同一套参数 transform flow_train flow[:train_len] scaler.fit(flow_train) flow_scaled scaler.transform(flow) # 全量统一按训练集均值/方差变换 def inverse_scale(y_scaled): return scaler.inverse_transform(y_scaled)提示归一化要按“路口”这列来做还是按整体矩阵来做取决于原始数据里各检测器是否有明显的量纲差异。如果所有检测器都是同一类设备、量纲一致直接对整个矩阵 fit 即可如果不同路段流量量级差别很大建议对每列分别 fit。4. 时空变换网络的最小实现PyTorch 源码逐块拆讲4.1 空间图卷积层邻接矩阵只做一次归一化空间 GCN 层的源码可以写成下面这样。关键是邻接矩阵在初始化时做一次对称归一化并转为稀疏张量前向过程里不要重复计算。import torch import torch.nn as nn import torch.nn.functional as F class SpatialGCN(nn.Module): def __init__(self, in_dim, out_dim, adj): super().__init__() self.adj adj # [N, N] 对称归一化后的邻接矩阵 self.fc nn.Linear(in_dim, out_dim) self.bn nn.BatchNorm1d(out_dim) def forward(self, x): # x: [B*T, N, in_dim] out torch.einsum(bnf,nm-bmf, x, self.adj.to(x.device)) out self.fc(out) # 线性变换到 out_dim out self.bn(out.permute(0, 2, 1)).permute(0, 2, 1) return F.relu(out)逻辑说明einsum 里的 bnf,nm-bmf 就是在做邻居聚合等价于 x 和邻接矩阵的矩阵乘法fc 层完成特征维度的映射BatchNorm 在节点维上归一化缓解图卷积中“度较大的节点特征被稀释”问题。邻接矩阵必须提前做好对称归一化再传进来不要在 forward 里重复操作否则训练会很慢。4.2 时间变换层在时间维度上做多头注意力时间变换层负责捕捉历史窗口内不同时刻之间的依赖。这里用简化版多头注意力采用可学习的相对时间偏移作为位置编码。class TemporalTransformer(nn.Module): def __init__(self, feature_dim, num_heads, dropout0.3): super().__init__() self.attn nn.MultiheadAttention( feature_dim, num_heads, dropoutdropout, batch_firstTrue ) self.norm1 nn.LayerNorm(feature_dim) self.norm2 nn.LayerNorm(feature_dim) self.ffn nn.Sequential( nn.Linear(feature_dim, feature_dim * 2), nn.ReLU(), nn.Dropout(dropout), nn.Linear(feature_dim * 2, feature_dim) ) def forward(self, x): # x: [B*N, T, feature_dim] shortcut x out, _ self.attn(x, x, x) out self.norm1(out shortcut) shortcut out out self.ffn(out) out self.norm2(out shortcut) return out逻辑说明这里的注意力计算是标准的 self-attentionquery、key、value 都是同一个序列LayerNorm 放在残差连接之后收敛更稳定。feature_dim 一般取 64 或 128num_heads 取 2 或 4因为交通流时间序列长度通常只有 1224头数太大会把特征切得太碎。4.3 主干模块和训练循环源码与核心参数的对应关系把空间 GCN 和时间 Transformer 交替堆叠就构成了一个最小的时空变换网络。class STTransformer(nn.Module): def __init__(self, num_nodes, his_len, feature_dim64, num_heads2, num_layers2, adjNone, horizon3): super().__init__() self.num_nodes num_nodes self.his_len his_len self.horizon horizon self.input_proj nn.Linear(1, feature_dim) self.blocks nn.ModuleList() for _ in range(num_layers): self.blocks.append(SpatialGCN(feature_dim, feature_dim, adj)) self.blocks.append(TemporalTransformer(feature_dim, num_heads)) self.output_proj nn.Linear(feature_dim * his_len, self.horizon) def forward(self, x): # x: [B, his_len, num_nodes, 1] B, T, N, F x.shape x self.input_proj(x) # [B, T, N, feature_dim] for gcn, tfm in zip(self.blocks[::2], self.blocks[1::2]): # 空间 xs x.permute(0, 1, 2, 3).reshape(B * T, N, -1) xs gcn(xs) xs xs.reshape(B, T, N, -1) x xs # 时间 xt x.permute(0, 2, 1, 3).reshape(B * N, T, -1) xt tfm(xt) xt xt.reshape(B, N, T, -1).permute(0, 2, 1, 3) x xt # 输出: 展平每个节点的时间维映射到 horizon 步 x x.reshape(B, N, -1) out self.output_proj(x) # [B, N, horizon] return out.permute(0, 2, 1) # [B, horizon, N]训练循环里损失函数和优化器的选择比网络本身更影响结果model STTransformer(num_nodesn_nodes, his_len12, feature_dim64, num_heads2, num_layers2, adjadj_norm, horizon3) optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemin, factor0.5, patience5 ) criterion nn.L1Loss() for epoch in range(200): model.train() for xb, yb in train_loader: xb xb.unsqueeze(-1).float() # [B, T, N, 1] yb yb.float() # [B, H, N] optimizer.zero_grad() pred model(xb) # [B, H, N] loss criterion(pred, yb) loss.backward() nn.utils.clip_grad_norm_(model.parameters(), 5.0) optimizer.step() # 每个 epoch 结束后在验证集上看 MAE并调整学习率参数说明学习率 1e-3 是这类网络的常见起点200 个 epoch 足够让模型收敛。ReduceLROnPlateau 会在验证指标连续 5 个 epoch 不下降时把学习率减半。L1Loss 代替 MAE 计算对流量尖峰不那么敏感。梯度裁剪值设 5.0防止图卷积聚合度差异引发梯度爆炸。5. 交通流预测的踩坑与排查五个血泪经验5.1 随机切分数据集导致验证集信息泄漏现象训练时 MAE 很好验证集上 MAE 也很漂亮所有指标都像打了兴奋剂一放到线上实时预测就翻车误差比训练时大 30% 以上。原因交通流是强时序关联数据随机切分训练集和验证集会让同一段时间窗口的样本一部分进训练集、一部分进验证集。模型的记忆里已经包含验证集时刻的交通状态相当于开卷考试。解决只用时间序切分。按时间先后顺序取前 70% 做训练、后 10% 做验证、最后 20% 做测试。如果流量数据跨年最好按“训练在去年测试在今年”的原则切这样才贴近真实场景。5.2 全量归一化把测试集统计混进训练现象训练损失下降很快验证损失也正常但预测曲线整体平滑得过分真实数据里的陡峭峰值全部被抹平。原因有些人图省事用全量数据的均值方差执行归一化包括未来数据。这等于把测试集的数值分布提前泄露给了训练过程。模型学到的输出范围被压缩到训练集和测试集的平均值附近预测值天然偏向中庸。解决严格区分 fit 和 transform。StandardScaler 只在训练集上 fit验证集和测试集只调用 transform不参与统计量计算。反归一化时也同样使用训练集拟合的 scaler。5.3 图卷积训练到一半 loss 变 NaN现象loss 一路正常下降到某个 epoch 后突然变成 NaN重启训练又能撑几个 epoch反复在同一个节点附近炸掉。原因邻接矩阵没有加自环或对称归一化时度矩阵里有 0 值。部分孤立节点没有连接任何邻居D 中对角元为 0D^{-1/2} 出现无穷inf 乘 zero 之后就是 NaN。另一种常见原因是邻接矩阵不是对称的GCN 的特征聚合方向不一致数值震荡加剧。解决在原始邻接矩阵上加单位矩阵 I然后再计算度矩阵和归一化矩阵对度矩阵的下标加上一个极小的 epsilon 如 1e-8 防止除零。不要只在训练前打印一次 adj训练中途也检查一次 adj 是否有 NaN 漂移。5.4 预测曲线整体被压低峰值全部消失现象MAE 和 RMSE 都在下降但你画出预测折线图后会发现早高峰的预测值始终比真实值低 20%夜间低流量时段反而略高预测曲线像被“拉平”了。原因L1 或 L2 损失对误差取平均之后模型为了减小整体误差会倾向于输出条件均值。交通流分布本身有长尾峰值条件均值天然低于尖峰时刻的真实值。有些实现里还用了“将多个时间步展平后再卷积”的做法等效于时间维度上的平均池化进一步压低峰值。解决在训练损失中加入峰值时刻的加权项比如 WMAE对高峰时段样本乘以 1.5 或 2 的权重。另一个实用技巧是预测误差较大的路段单独放大权重避免网络只优化易学习的低速路段。最后确认反归一化步骤没有把缩放因子用错这类问题在形状对的情况下更容易被归因到模型。5.5 维度变形算得对就是报负数维度现象forward 跑到 reshape 步骤时报错提示 “shape 包含 -1” 且计算出的维度是负数一看代码没写错矩阵乘法也算得通但就是崩。原因时间变换层把 [B, N, T, F] 转成 [BN, T, F] 时BN 是动态值batch size 不整除就会让后续 reshape 到 [B, T, N, F] 时对不上维度。另一处是 last batch 的尺寸和其他 batch 不一致模型在最后一个 batch 上生成错误形状。解决训练 DataLoader 设置 drop_lastTrue放弃不足一个 batch 的数据对 reshape 前的张量记录下原始 B、T、N、F 值避免用 -1 推导。前向传播里尽量显式传递 B、T、N 变量而不是靠 shape[-1] 倒推。调试阶段在模型 forward 中插入打印每个关键张量的 shape配一个 batch size 1 的随机输入跑通再上数据。6. 验证与进阶用消融和基线衡量模型到底学到什么6.1 三个基线对比持久性模型、纯时序模型、纯空间模型时空变换网络最怕的不是分数差是“显得好但毫无因果”。我一般会准备三个基线来定位它到底赢在哪。基线做法代表模型它能验证什么持久性模型用最近一个时间步的流量作为未来预测把 x[:, -1, :] 直接当预测数据本身有没有下降趋势纯时序模型只用时间维度信息不用邻接矩阵LSTM、TCN时间建模的收益占比纯空间模型只用当前时刻的空间邻居特征GCN-only空间建模的收益占比评估指标统一用 MAE、RMSE、MAPE 三个MAPE 在凌晨低流量时段会虚高所以同时记录白天 7 点到 22 点的分时段 MAPE。对比时如果时空变换网络比纯时序模型只提升 3%但比持久性模型提升 40%说明主要收益来自时序建模空间结构还没被充分利用反过来说明邻接矩阵构造需要重做。6.2 按高峰时段拆开看误差而不只看平均指标平均指标会掩盖模型在关键时段的失败。真实路网中管理者最关心的是早晚高峰和突发事件的预测质量。评估时把预测误差按小时维度拆开peak_mask (test_timestamps_hour 7) (test_timestamps_hour 9) nonpeak_mask ~peak_mask def report(y_true, y_pred, mask): mae np.mean(np.abs(y_true[mask] - y_pred[mask])) rmse np.sqrt(np.mean((y_true[mask] - y_pred[mask]) ** 2)) return mae, rmse我见过一个模型整体 MAPE 只有 12%看起来不错但拆开一看早高峰 7 点到 8 点 MAPE 高达 28%等于在最需要准确预测的场景里几乎不可用。这类问题只有按时段拆误差才会暴露建议把“全天指标 分时段指标”写进固定的评估流程里不要只看一张测试报告。6.3 从实验到上线的三个执行习惯第一固定种子。时空变换网络涉及多头注意力和随机丢弃不固定种子的话同一份代码跑两次结果差异可能达到 5% 以上没法做消融对比。在训练脚本开头设置 torch.manual_seed(42) 和 numpy.random.seed(42)必要时在 DataLoader 里同样设置 worker 的种子。第二保留最朴素的 baseline 进监控。上线后把持久性模型作为实时监控的参照线如果某一天时空变换网络的实时误差和持久性模型持平大概率是数据源出了问题而不是模型变笨了。这个习惯能帮你快速定位到传感器掉线或时间戳错位。第三每周滚动重训。交通流存在节假日效应、季节变化和新开通道路带来的拓扑改变长期不更新会让模型悄悄退化。我一般用“每天增量更新特征 每周全量重训”的节奏重训时保留前几周的测试集作为回归测试集确保新一轮模型不比上一轮差。我的习惯是把“能否在验证集上稳定复现提升”当作第一验收标准如果一次实验的提升换个随机种子就消失那这个优化不值得上线。时空变换网络的真正优势在于把空间和时间揉进了同一个特征空间里这个收益是结构性的不依赖某一次随机初始化。希望帮到你。本文还有配套的精品资源点击获取