简介这份资源面向希望深入理解并动手复现时空预测模型的PyTorch学习者与研究者聚焦于将论文中的公式、图示及原作者实现转化为可运行的模块化代码适合具备一定深度学习基础、想从零搭建或对照学习时空预测结构的中高级开发者。压缩包共65个文件以53个Python源码为主辅以少量XML配置、说明文档与许可证文件整体约57KB体量轻便。代码按models目录组织每个文件夹对应一个完整模型结构统一假定输入张量为(batch, sequence, channel, height, width)五维格式util目录提供针对大尺寸数据的patch分割方法、可继承重写的训练与测试模板类以及生成目录树的工具便于按需改造。已有580人学习读者可借此获得一套结构清晰、内聚性强的复现参考理解模型拆分与组合思路并在此基础上快速搭建自己的实验流程。1. 时空预测模型复现从论文公式到能跑通的 PyTorch 工程第一次拿到时空预测的论文源码时我盯着model.py里那堆einsum和邻接矩阵乘法看了整整一个下午——公式在纸上明明只有三行落到代码里却牵扯到时间维卷积、空间图卷积、残差连接三套逻辑交织。这类项目最坑的地方不在于模型本身多复杂而在于数据预处理、滑动窗口切分、多步预测的评估口径这三处任何一处对不上loss 曲线就会给你表演玄学。这份 PyTorch 复现工程把训练、验证、测试的完整链路都摊开了适合两类人一是想跑通时空预测 baseline 做对比实验的研究生二是需要把交通流量、气象序列这类时空数据接进自己业务管线的工程师。下面按「能干什么 → 怎么跑 → 怎么改」的顺序拆。2. 工程骨架与数据管线时空张量是怎么被喂进模型的2.1 目录结构与模块职责复现类项目的目录设计直接决定了你二次开发时改哪里、不动哪里。这套工程的组织方式我比较认可职责边界清晰stgnn-repro/ ├── configs/ # yaml 配置数据集/模型/训练超参分离 │ ├── base.yaml │ └── metr_la.yaml ├── data/ │ ├── raw/ # 原始 npz / csv │ └── processed/ # 切好的滑窗样本缓存 ├── datasets/ │ ├── __init__.py │ └── spatiotemporal.py # Dataset 滑窗切分 归一化 ├── models/ │ ├── blocks.py # 时间卷积块、图卷积块 │ └── stgcn.py # 主干网络组装 ├── engine/ │ ├── trainer.py # 训练循环、早停、checkpoint │ └── metrics.py # MAE/RMSE/MAPE按 horizon 分别统计 ├── scripts/ │ ├── preprocess.py │ └── train.py └── requirements.txtdatasets/和models/是两条独立的改动线换数据集只动前者换模型结构只动后者engine/里的训练逻辑基本不用碰。这种切法在复现多个 baseline 做消融时特别省事我一般会把models/下的每个变体单独建文件共用blocks.py里的基础算子。2.2 滑动窗口切分时空预测的第一个翻车点时空预测的输入不是单条序列而是「历史 T 步 × N 个节点 × C 个特征」的三维张量输出是「未来 H 步 × N 节点」。切分逻辑写错后面全白搭。核心代码大致长这样import numpy as np import torch from torch.utils.data import Dataset class STWindowDataset(Dataset): def __init__(self, data, input_len12, pred_len12, stride1): # data: (T_total, N, C) 已归一化 self.data data self.input_len input_len self.pred_len pred_len # 样本数 可切出的完整窗口数量 self.num_samples (len(data) - input_len - pred_len) // stride 1 self.stride stride def __len__(self): return self.num_samples def __getitem__(self, idx): start idx * self.stride x self.data[start : start self.input_len] # (T_in, N, C) y self.data[start self.input_len : start self.input_len self.pred_len, :, 0] # (T_out, N) # 转成 (N, C, T) 方便后续时间维卷积 x np.transpose(x, (1, 2, 0)) return torch.as_tensor(x, dtypetorch.float32), \ torch.as_tensor(y, dtypetorch.float32)这里有两个参数必须和论文对齐input_len和pred_len。交通流量类数据集常见配置是 12 步历史预测 12 步未来按 5 分钟采样即 1 小时但气象序列可能用 24 步。stride控制样本重叠度训练集用 1 保证样本量验证测试集也建议用 1否则评估指标会因窗口选择产生偏差。注意y只取第 0 个特征通道——多变量输入单变量预测是这类任务的默认设定如果你的业务要预测全部通道把[:, 0]去掉并相应改 loss 的维度处理。2.3 归一化与反归一化别在评估时丢掉尺度z-score 归一化必须在训练集上统计均值方差然后应用到验证测试集这是铁律。但很多人栽在反归一化上模型输出的是归一化空间的预测值算 MAE 前必须乘回标准差加回均值否则指标看起来漂亮得离谱。class StandardScaler: def __init__(self): self.mean None self.std None def fit(self, train_data): # train_data: (T, N, C) self.mean train_data.mean(axis(0, 1), keepdimsTrue) self.std train_data.std(axis(0, 1), keepdimsTrue) 1e-8 def transform(self, data): return (data - self.mean) / self.std def inverse_transform(self, data): return data * self.std self.meankeepdimsTrue不能省否则广播维度对不上。1e-8是防除零的常规操作某些节点流量长期为零时标准差会趋近 0。评估阶段对预测值和真值同时做inverse_transform再算指标这一步我在engine/metrics.py里封装成统一入口避免训练脚本和评估脚本各写一套导致口径不一致。3. 模型主干拆解时间卷积与图卷积如何拼成 STGCN3.1 时间卷积块用一维卷积抓局部时序模式时空预测模型的时间维处理通常用门控一维卷积或因果卷积。这套复现用的是 GLU门控线性单元结构比 LSTM 训练快且并行度高import torch.nn as nn class TemporalConvBlock(nn.Module): def __init__(self, in_ch, out_ch, kernel_size3): super().__init__() self.conv nn.Conv2d(in_ch, out_ch * 2, kernel_size(1, kernel_size), padding(0, kernel_size // 2)) self.sigmoid nn.Sigmoid() def forward(self, x): # x: (B, C_in, N, T) h self.conv(x) # (B, 2*C_out, N, T) a, b h.chunk(2, dim1) # 各 (B, C_out, N, T) return a * self.sigmoid(b) # GLU 门控kernel_size3对应感受野覆盖前后各 1 步堆叠多层可扩大时序感受野。padding用kernel_size // 2保持时间维长度不变方便残差连接。注意输入布局是(B, C, N, T)节点维 N 在时间维 T 前面这是为了后续图卷积在 N 维上操作时内存连续。如果你从别的代码库搬模块过来第一件事就是核对张量布局(B, T, N, C)和(B, C, N, T)混用是血泪教训级别的 bug 来源。3.2 图卷积块邻接矩阵决定空间信息流向空间维建模依赖预定义的邻接矩阵基于路网距离或节点相似度。图卷积的核心是「聚合邻居特征」class GraphConvBlock(nn.Module): def __init__(self, in_ch, out_ch, adj): super().__init__() # 对称归一化: D^{-1/2} A D^{-1/2} d adj.sum(dim1) d_inv_sqrt torch.pow(d, -0.5) d_inv_sqrt[torch.isinf(d_inv_sqrt)] 0. self.adj_norm (d_inv_sqrt.unsqueeze(1) * adj * d_inv_sqrt.unsqueeze(0)) self.weight nn.Parameter(torch.empty(in_ch, out_ch)) nn.init.xavier_uniform_(self.weight) def forward(self, x): # x: (B, C_in, N, T) x x.permute(0, 3, 2, 1) # (B, T, N, C_in) out torch.matmul(self.adj_norm, x) # 邻居聚合 out torch.matmul(out, self.weight) # 特征变换 return out.permute(0, 3, 2, 1) # 回到 (B, C_out, N, T)归一化那几行是图卷积的标准预处理d_inv_sqrt里把 inf 置零是为了处理孤立节点度为 0。adj作为 buffer 注册进模型而不是普通属性这样model.to(device)时会自动跟着搬。邻接矩阵的构建在scripts/preprocess.py里常见做法是基于节点间距离阈值或高斯核阈值取多少直接影响模型效果我一般会先画一下度分布再定。3.3 主干组装残差连接与输出层把时间块和图块交替堆叠每个 ST 块内部用残差连接class STBlock(nn.Module): def __init__(self, ch, adj, kernel_size3): super().__init__() self.tconv TemporalConvBlock(ch, ch, kernel_size) self.gconv GraphConvBlock(ch, ch, adj) self.norm nn.LayerNorm(ch) def forward(self, x): residual x h self.tconv(x) h self.gconv(h) h h residual # 残差 return self.norm(h.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)LayerNorm放在通道维上permute 两次是为了让 norm 作用在最后一维。堆 23 个 STBlock 后接一个Conv2d把通道压到 1再取最后pred_len个时间步作为输出。残差连接要求输入输出通道一致所以每个 block 的ch固定通道扩展只在第一层做。4. 训练与评估让 loss 曲线不再玄学4.1 训练循环的关键配置import torch from torch.optim import Adam from torch.optim.lr_scheduler import ReduceLROnPlateau def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss 0. for x, y in loader: x, y x.to(device), y.to(device) optimizer.zero_grad() pred model(x) # (B, T_out, N) loss criterion(pred, y) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0) optimizer.step() total_loss loss.item() * x.size(0) return total_loss / len(loader.dataset)clip_grad_norm_的max_norm5.0是时空模型训练的常规保险图卷积层梯度爆炸不算罕见。优化器用 Adam、学习率 1e-3 起步配合ReduceLROnPlateaupatience5factor0.5在验证 loss 停滞时降学习率。早停 patience 一般设 1520 个 epochcheckpoint 只保存验证集 MAE 最优的那份。4.2 分 horizon 评估别只看平均指标多步预测的评估必须按预测步长分别统计因为误差随 horizon 增长是常态平均指标会掩盖长步预测的崩坏def evaluate(model, loader, scaler, device, horizons(3, 6, 12)): model.eval() preds, trues [], [] with torch.no_grad(): for x, y in loader: pred model(x.to(device)).cpu().numpy() preds.append(pred) trues.append(y.numpy()) preds scaler.inverse_transform(np.concatenate(preds)) trues scaler.inverse_transform(np.concatenate(trues)) results {} for h in horizons: mae np.abs(preds[:, h-1] - trues[:, h-1]).mean() rmse np.sqrt(((preds[:, h-1] - trues[:, h-1]) ** 2).mean()) results[fMAE{h}] mae results[fRMSE{h}] rmse return resultshorizons(3, 6, 12)对应 15 分钟、30 分钟、60 分钟预测这是交通预测论文的标准报告口径。MAPE 要小心处理零值常见做法是加一个小的 epsilon 或只对真值大于阈值的样本计算。5. 避坑与排查复现路上最容易踩的五个坑5.1 现象loss 不降反升几轮后 NaN原因通常是学习率过大叠加图卷积的度归一化除零。检查adj是否有全零行d_inv_sqrt的 inf 是否被正确置零。解决把max_norm降到 1.0学习率降到 5e-4并在归一化后打印torch.isnan(adj_norm).any()确认。5.2 现象训练集指标很好测试集 MAE 高出一截八成是归一化用了全量数据的统计量造成信息泄漏。确认scaler.fit只传训练集切片验证测试集只调transform。另一个可能是滑窗切分时训练集和测试集有重叠窗口stride和切分边界要按时间顺序硬切不能随机打乱。5.3 现象换了数据集后模型完全不收敛先核对邻接矩阵的节点数是否和数据特征维的 N 一致。不同数据集的节点编号顺序可能不同邻接矩阵必须按同一套节点索引构建。再检查特征通道数单特征数据集要把in_ch改成 1别沿用多特征的配置。5.4 现象GPU 显存爆掉batch_size 调到 1 也不行时空张量(B, C, N, T)在 N 较大时比如 800 节点显存占用是平方级增长的。解决把图卷积的matmul换成稀疏矩阵乘法torch.sparse.mm邻接矩阵用to_sparse()转换或者把时间块和图块的计算拆开做梯度累积。5.5 现象复现指标和论文差 23 个点先确认评估口径论文报的是不是归一化空间的指标、horizon 定义是否一致、是否用了相同的训练/验证划分比例。再检查随机种子时空模型对初始化敏感固定torch.manual_seed和np.random.seed后跑三次取均值再对比。差 1 个点以内属于正常波动差 3 个点以上大概率是数据管线某处不一致。6. 二次开发换数据集与改结构的实操路径改造这套工程我习惯从数据侧和模型侧两条线分别验证。数据侧换新数据集时先只跑scripts/preprocess.py把切好的样本 shape 打印出来核对(num_samples, C, N, T_in)和(num_samples, T_out, N)确认无误再进训练。邻接矩阵如果新数据集没有现成的用节点坐标算高斯核距离矩阵阈值从 0.1 开始试观察度分布是否合理——全连接或大量孤立节点都会让图卷积失效。模型侧做消融时最省事的改法是在configs/里加开关而不是改代码model: num_blocks: 2 channels: 64 kernel_size: 3 use_graph: true # 关掉即退化为纯时间卷积 baseline use_residual: true然后在stgcn.py里根据use_graph决定是否插入GraphConvBlock。这样一套代码能同时跑出「有时间无空间」「有空间无时间」「完整 STGCN」三组对比写论文或做技术选型时直接出表。验证改造是否成功我的习惯是跑一个「过拟合小样本」测试取 200 个训练样本关掉正则和 dropout训练 50 个 epoch如果训练 loss 降不到接近零说明模型结构或数据管线有硬伤别急着上全量数据。这个测试帮我省过好几次在错误方向上调参的时间。最后说个我踩过的坑有次改完图卷积的归一化方式指标涨了两个点兴冲冲准备写进报告结果发现是评估脚本里inverse_transform的 std 用了测试集的统计量——等于偷看了答案。从那以后我每次改完评估相关代码都强制用一个手工构造的小张量走一遍前向和反归一化确认数值对得上再跑全量。希望这些能帮你在复现路上少绕几个弯。本文还有配套的精品资源点击获取