基于PyTorch的时空预测模型复现：图卷积与时间卷积工程实践

简介：这是一份面向时空预测模型研究与PyTorch进阶学习者的PyTorch复现代码包。针对交通流量、气象等序列预测场景，资源将多个论文模型拆解为内聚的小Module再组装，便于读者对照公式与图示理解结构，并适配(batch, sequence, channel, height, width)五维输入，适合希望从零搭建时空预测网络的开发者。压缩包共65个文件，以53个py源码为核心，同时包含6个xml配置、Markdown说明、gitignore等辅助文件，整体仅57KB，轻量易下载。目前已有579人学习使用。util目录下提供了针对大尺寸数据的patch分割方法，以及自写的训练流程模板类TrainingTemplate和TestingTemplate，继承后只需重写关键方法即可快速完成训练与测试；content_tree工具可自动生成目录树，方便梳理项目结构。models目录下每类结构单独成文件夹，代码组织清晰，既可用于项目二次开发，也可作为阅读论文实现的辅助笔记，兼顾学习与实战需求。

1. 时空预测模型复现：一套能直接改的PyTorch工程

如果你拿到一份城市路网的拥堵指数、一排站点未来几小时的空气质量，或者一栋楼未来一天的用电负荷，你遇到的问题都属于时空预测：每个观测点既随时间变化，又受相邻观测点影响。把每个点当成独立序列用 LSTM 硬训，等于扔掉了空间关系；用纯图神经网络又容易忽略时间依赖。这套工程复现的是把图卷积和时间卷积拼在一起做时空建模的经典思路，用 PyTorch 从零手写，不依赖专用 GNN 框架。数据管道、训练循环、评测脚本都齐，适合正在啃论文又不想从零造轮子的开发者，拿来就能跑，改一改就能接自己的数据。

2. 项目架构与目录设计：数据管道、模型层和训练循环怎么组织

2.1 目录结构与模块职责

整个工程的组织方式很直白，主目录下按输入输出拆分，源码集中放在src里。

我一般会把utils.py当成一个纯工具模块，里面不放任何模型相关代码，这样在调试 notebook 里也可以单独import utils去算指标或者处理邻接矩阵。model.py也只依赖 PyTorch 基础模块，不引用数据相关的类，保证模型纯函数式，方便单测。

config.yaml 的内容大概长这样：

data: data_path: data/raw/traffic.csv adj_path: data/raw/adjacency.npy history: 12 horizon: 12 model: hidden_channels: 64 num_layers: 2 kernel_size: 3 train: epochs: 60 batch_size: 64 lr: 0.001 seed: 42 device: cuda

参数说明：history是回看窗口长度，horizon是预测未来多少个步长；hidden_channels控制图卷积和时间卷积的特征维度，模型宽度基本由它决定；kernel_size是时间卷积核大小，更大的核能看更长的局部时间模式，但会显著增加显存占用和过拟合风险。

2.2 数据加载：归一化、滑窗与批次化

时空预测的数据一般是一个二维矩阵：行是时间戳，列是节点编号。比如交通数据就是[时段数, 路段数]，空气质量就是[小时数, 站点数]。这类数据量通常不大，完全可以一次性读入内存，再在内存里做滑窗，不需要像 NLP 项目那样搞复杂的流式加载。

数据加载最核心的部分是滑窗构造。下面这份代码是简化后的 Dataset 实现：

import torch import numpy as np import pandas as pd from torch.utils.data import Dataset class SpatioTemporalDataset(Dataset): def __init__(self, data_path, history=12, horizon=12): # 读原始数据，shape: [T, N]，T 是时间步数，N 是节点数 df = pd.read_csv(data_path, index_col=0) self.data = df.values.astype(np.float32) self.T, self.N = self.data.shape self.history = history self.horizon = horizon def __len__(self): # 最后一个可用样本要留出 history + horizon 的长度 return self.T - self.history - self.horizon + 1 def __getitem__(self, idx): # 输入是过去 history 步，标签是未来 horizon 步 x = self.data[idx : idx + self.history, :] y = self.data[idx + self.history : idx + self.history + self.horizon, :] return torch.from_numpy(x), torch.from_numpy(y)

这段代码有几个细节需要说明一下。self.data的 shape 是[T, N]，__getitem__返回的xshape 是[history, N]，yshape 是[horizon, N]，也就是用过去 12 个时间步的值预测未来 12 个时间步的值。这里故意没有在 Dataset 内部做归一化，因为归一化需要统计全体数据的均值和方差，如果每个样本独立归一化，训练和预测时的统计口径就不一致了，后面会引发数据泄露问题。

训练脚本里的归一化逻辑长这样：

def normalize(data): # 只对训练集统计均值和标准差 mean = data[:train_size].mean(axis=0) std = data[:train_size].std(axis=0) std[std == 0] = 1.0 return (data - mean) / std, mean, std

归一化的方式是经典 Z-Score，按列（每个节点单独）统计均值和方差，再对整个序列做标准化。std[std == 0] = 1.0是为了防止某些节点长时间没有变化导致除零，这种情况在实际传感器数据里经常出现，尤其是故障期间一直上报 0 值的节点。

DataLoader 端只需要设置正常的batch_size和shuffle，但要注意测试集不要打乱顺序，预测任务需要保持时间连续性，方便最后绘图对比。

2.3 训练循环与评测指标

训练循环本身不复杂，但有几个地方值得注意。时空预测任务里最常用的损失函数是 MAE（平均绝对误差），因为交通、气象这类数据的分布通常有长尾，MSE 会被少数极端值主导，MAE 对异常值更稳健，且与最终业务指标一致。

import torch import torch.nn as nn def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0.0 for x, y in loader: x = x.to(device) y = y.to(device) pred = model(x) loss = criterion(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * x.size(0) return total_loss / len(loader.dataset)

loss.item()取出当前 loss 的数值用于累加，乘以x.size(0)是按 batch 大小加权，最后除以数据集总长度得到平均 loss。这里为什么不直接用total_loss += loss.item()？因为最后一个 batch 可能不足batch_size，直接累加平均会偏向小 batch，虽然是微小影响，但训练曲线会多一些毛刺。

评测指标定义在utils.py里，一般同时算三个：

def evaluate_metrics(y_true, y_pred): # y_true / y_pred: [B, H, N] mae = torch.mean(torch.abs(y_true - y_pred)).item() rmse = torch.sqrt(torch.mean((y_true - y_pred) ** 2)).item() mape = torch.mean(torch.abs((y_true - y_pred) / (y_true + 1e-8))).item() return mae, rmse, mape

MAPE 计算时在分母加了1e-8防除零，但真实值如果确实为 0，这个指标的意义会失真，所以看结果时最好同时关注真实值分布里有多少接近 0 的样本。验证阶段我通常每个 epoch 跑一次验证集，保存验证 MAE 最低的权重，而不是最后一个 epoch 的权重，防止过拟合后反而存了个更差的模型。

3. 核心代码拆解：图卷积与时间卷积的 PyTorch 实现

3.1 邻接矩阵的构造与归一化

图卷积的第一步是把节点关系表示成邻接矩阵。邻接矩阵的来源通常有两种：一种是数据本身带了物理连接关系（比如路段之间的上下游关系），另一种需要从节点坐标或距离矩阵手工计算。后者最常见的做法是用高斯核函数：

import numpy as np def build_adj_from_distance(dist_matrix, sigma=0.1): # dist_matrix: [N, N]，节点两两之间的距离 adj = np.exp(-(dist_matrix ** 2) / (sigma ** 2)) np.fill_diagonal(adj, 0) return adj

高斯核的距离越近，权重越大；距离越远，权重指数衰减。sigma是个超参数，控制着空间影响的衰减速度，一般需要根据实际数据尺度去试，我习惯先看距离矩阵的分位数再定。

拿到原始的邻接矩阵后，直接喂给图卷积会有问题。原始 A 的行和没有归一化，节点度数大的在聚合时会主导，迭代多层后容易梯度爆炸。常用做法是拉普拉斯对称归一化：

import scipy.sparse as sp def normalize_adjacency(adj): # 加自环：让节点聚合时保留自身信息 adj = adj + sp.eye(adj.shape[0]) deg = np.array(adj.sum(1)).flatten() d_inv_sqrt = np.power(deg, -0.5) d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0 d_mat_inv_sqrt = sp.diags(d_inv_sqrt) return d_mat_inv_sqrt.dot(adj).dot(d_mat_inv_sqrt).tocoo()

这里的归一化公式是 D^{-1/2} A D^{-1/2}，效果是让聚合时同时考虑自身节点度和邻居节点度，度数高的大节点不再天然占优。加自环的意思是A + I，让节点在做空间聚合时保留自己的上一时刻信息，不然第一层图卷积后节点只能看到邻居而看不到自己。归一化后的邻接矩阵会转成 COO 稀疏格式，方便后续和稀疏张量配合使用。

3.2 图卷积模块：聚合相邻节点信息

图卷积在 PyTorch 里最简单的实现，就是把邻接矩阵当成一个可微分的线性算子，直接在空间维度上做矩阵乘法。我见过有人引入 DGL 或 PyG 来做这一步，但复现项目完全没必要，稀疏矩阵乘法足够，而且能看清每一步在干什么。

import torch import torch.nn as nn class GraphConv(nn.Module): def __init__(self, in_channels, out_channels, adj): super().__init__() # adj 注册为 buffer，不参与梯度更新但会随模型一起迁移到 GPU self.register_buffer("adj", torch.from_numpy(adj.toarray()).float()) self.weight = nn.Parameter(torch.randn(in_channels, out_channels) * 0.1) self.bias = nn.Parameter(torch.zeros(out_channels)) def forward(self, x): # x: [B, T, N, C_in] B, T, N, C_in = x.shape # 空间维做邻接矩阵聚合 x = torch.einsum("btnc,nm->btmc", x, self.adj) # [B, T, N, N] 聚合邻居 x = torch.einsum("btmc,co->btmo", x, self.weight) # [B, T, N, C_out] return x + self.bias

einsum("btnc,nm->btmc")这一步是图卷积的灵魂：n和m都是节点维度，它做的事情等价于对每个节点m，把相邻节点n的输入特征按邻接矩阵权重求和。第二个einsum是普通的全连接变换，把聚合后的特征从C_in映射到C_out。整体效果就是「先聚合邻居，再变换特征」。

初始化时把权重按0.1缩放，是为了避免初始输出过大。如果这一步不做，模型刚开始训练时 loss 可能直接飞到几万，视觉上像是「模型坏了」，其实只是初始化范围太野。

3.3 时间卷积：因果卷积与门控线性单元

空间维度用图卷积建模后，时间维度还是需要单独处理的。常见选择是 Conv2d 配合门控线性单元（GLU）。这里用一维卷积在时间维上滑，配合 padding 保持序列长度不变，再用 GLU 提供非线性门控。

import torch import torch.nn as nn class TemporalConv(nn.Module): def __init__(self, in_channels, out_channels, kernel_size=3): super().__init__() self.conv = nn.Conv2d( in_channels, out_channels * 2, (1, kernel_size), padding=(0, (kernel_size - 1) // 2), ) self.glu = nn.GLU(dim=1) def forward(self, x): # x: [B, C, T, N] h = self.conv(x) return self.glu(h)

Conv2d 的卷积核是(1, kernel_size)，意思是空间维上卷积核大小为 1（不改变节点维度），只在时间维上滑动。输出通道数翻倍为out_channels * 2，这是为了 GLU 做准备：把输出沿通道维切成两半，一半通过 sigmoid 当门控信号，另一半直接当内容，两者逐元素相乘得到最终输出。

为什么要用 GLU 而不是直接 ReLU？门控让网络学会决定「哪些历史信息需要保留，哪些要丢弃」，对时间序列这种依赖关系复杂的场景效果会更好，而且梯度能直接穿过门控支路，训练也更稳定。这个设计在语义上相当于一个可学习的遗忘门。

3.4 完整模型前向传播：时空块如何串联

把图卷积和时间卷积拼在一起，就得到了一个经典的时空卷积块：先做时间卷积更新每个节点的时间特征，再做图卷积聚合空间信息，最后再做一次时间卷积。整体结构如下：

class STGCNModel(nn.Module): def __init__(self, num_nodes, history, horizon, in_channels, hidden_channels, out_channels, adj): super().__init__() self.history = history self.horizon = horizon self.num_nodes = num_nodes # 输入层：先把输入映射到 hidden_channels self.embed = nn.Conv2d(in_channels, hidden_channels, (1, 1)) # 两个时空卷积块 self.gconv1 = GraphConv(hidden_channels, hidden_channels, adj) self.tconv1 = TemporalConv(hidden_channels, hidden_channels) self.gconv2 = GraphConv(hidden_channels, hidden_channels, adj) self.tconv2 = TemporalConv(hidden_channels, hidden_channels) # 输出层：把最后一个时间步压缩到 horizon 维度 self.out = nn.Conv2d(hidden_channels, out_channels, (1, history)) def forward(self, x): # x: [B, T, N, F_in] # 转成 [B, F, T, N] 给 Conv2d x = x.permute(0, 3, 1, 2).contiguous() x = self.embed(x) # 第一个时空块：时间 -> 空间 -> 时间 x = self.tconv1(x) # [B, C, T, N] x = x.permute(0, 2, 3, 1).contiguous() # [B, T, N, C] 给图卷积 x = self.gconv1(x) # [B, T, N, C] x = x.permute(0, 3, 1, 2).contiguous() # 转回 [B, C, T, N] x = self.tconv2(x) # [B, C, T, N] # 第二个时空块：重复一次，加深非线性表示 x = self.tconv2(x) x = x.permute(0, 2, 3, 1).contiguous() x = self.gconv2(x) x = x.permute(0, 3, 1, 2).contiguous() x = self.tconv1(x) # 输出层：时间维从 history 压缩到 horizon x = self.out(x) # [B, C_out, H, N] return x.permute(0, 2, 3, 1) # [B, H, N, C_out]

这里最需要说明的是维度转换：Conv2d 期望输入是[B, C, T, N]，GraphConv 期望输入是[B, T, N, C]，两者之间必须反复permute。这一块是手写时空模型最容易翻车的地方，一不小心就是「维度不匹配」或者更隐蔽的「维度对但语义错」——比如把时间维和通道维搞混，模型能跑但效果极差。

模型最后用nn.Conv2d(hidden_channels, out_channels, (1, history))做输出投影，在时间维上用卷积核覆盖整个回看窗口，输出长度自动变成horizon。如果horizon和history不一致，需要调整卷积核大小或者后面再接一个线性层做维度映射。

4. 从零跑通：环境配置、训练命令与结果复现

4.1 环境依赖与数据准备

这个项目依赖不多，核心就 PyTorch、pandas、numpy、scipy 和 TensorBoard。PyTorch 建议使用 1.10 以上版本，因为后续如果要切混合精度训练，老版本支持不好。scipy 用于邻接矩阵的稀疏操作和归一化计算，如果你的邻接矩阵规模不大，用 numpy 直接算也完全可以。

数据准备这一步，我习惯先跑一段数据探查脚本，确认时间序列里没有大段空值和异常跳变。常见的脏数据是传感器掉线导致的整段 0 值，如果不处理，归一化时均值会被拉偏，模型会学到「输出 0」这种偷懒行为。

import pandas as pd def inspect_data(data_path): df = pd.read_csv(data_path, index_col=0) print(f"shape: {df.shape}") # [T, N] print(f"null count: {df.isnull().sum().sum()}") print(f"zero ratio: {(df.values == 0).mean():.4f}")

零值比例如果超过 5%，建议先做处理：要么把连续零值当成缺失填充，要么标记成一个特殊值。我处理交通数据时遇到过某路段连续三天全零的情况，后来查日志才知道是检测器坏了，这种数据留着只会干扰模型学习真实的空间关联。

4.2 训练命令与关键超参

环境准备好后，训练流程是直接跑src/train.py：

cd spatiotemporal_prediction python src/train.py --config config.yaml

训练脚本核心部分如下：

def main(): args = load_config("config.yaml") set_seed(args.train.seed) raw_data = pd.read_csv(args.data.data_path, index_col=0).values train_size = int(len(raw_data) * 0.7) val_size = int(len(raw_data) * 0.1) test_size = len(raw_data) - train_size - val_size # 按时间顺序划分，不打乱 train_data = raw_data[:train_size] val_data = raw_data[train_size : train_size + val_size] test_data = raw_data[train_size + val_size:] # 只用训练集统计归一化参数 mean, std = compute_stats(train_data) train_data = normalize(train_data, mean, std) val_data = normalize(val_data, mean, std) test_data = normalize(test_data, mean, std) train_set = SpatioTemporalDataset(train_data, args.data.history, args.data.horizon) val_set = SpatioTemporalDataset(val_data, args.data.history, args.data.horizon) test_set = SpatioTemporalDataset(test_data, args.data.history, args.data.horizon) ...

切分比例是 7:1:2，按时间先后切分而不是随机切分。时序预测如果随机打乱，会直接把未来信息泄漏进训练集，指标全线失真。归一化参数基于训练集单独统计，验证集和测试集直接复用同一组 mean、std，这是复现任何时序项目都必须遵守的纪律。

4.3 观察训练曲线与判断收敛

训练时每 5 个 epoch 记录一次训练 loss 和验证 MAE 到 TensorBoard：

from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter("logs") for epoch in range(epochs): train_loss = train_one_epoch(...) val_mae = evaluate(model, val_loader, device) writer.add_scalar("loss/train", train_loss, epoch) writer.add_scalar("metric/val_mae", val_mae, epoch)

我判断模型是否收敛从来不光看训练 loss，因为空间和时间模型往往出现过拟合先于 loss 下降停滞的情况。具体操作是：训练 loss 还在降但验证 MAE 连续 5 个 epoch 不再下降，就可以停了，再练下去就是纯浪费算力。如果验证 MAE 在一开始就比训练 loss 高一大截，先怀疑不是模型问题，而是数据划分和归一化口径出了问题。

5. 常见问题与避坑记录：时空模型复现路上的五个大坑

5.1 数据泄露：测试集的归一化参数算到了训练集头上

现象：验证集和测试集指标异常好，MAE 逼近 0，但把模型用到新数据上效果一塌糊涂。

原因：有人在数据加载里对全体数据做归一化后再切分训练测试集，或者把整个矩阵的均值和方差一次性算完。这样测试集的信息在归一化时已经透传给了模型，相当于考试前看到了答案。

解决：严格按「先切分，再分别统计」的顺序，只用训练集计算 mean 和 std，验证集、测试集复用这套参数。这是时序项目里最常见的翻车原因，每次跑新数据集我都先查一遍归一化代码。

5.2 邻接矩阵归一化不当导致梯度爆炸

现象：训练十几个 epoch 后 loss 突然变成 NaN，而且每次 NaN 出现的位置集中在某几个节点上。

原因：直接用原始邻接矩阵做消息传递，度数高的节点聚合了太多邻居信息，经过多层图卷积后数值指数增长，最后精度溢出。

解决：换用对称归一化D^{-1/2} A D^{-1/2}，给模型加梯度裁剪。图卷积层的权重初始化范围也尽量缩小一点，不要用 PyTorch 默认的 Kaiming 初始化范围。

5.3 滑窗数据集边界处理不当导致维度不一致

现象：最后一个 batch 报 shape mismatch，提示 batch 维度对不上。

原因：__len__计算时少减了窗口长度，或者__getitem__里切片越界返回到空数组。尤其是切分训练/验证集后，每个子集的样本数都在变，容易出现这种边界问题。

解决：__len__用T - history - horizon + 1计算，并且写一个简单测试，遍历一遍 Dataset，确保每个样本的 shape 都是[history, N]和[horizon, N]。这种问题跑一次全量遍历就能暴露，不要等到训练时才让框架报错。

5.4 显存超出：序列长度和图规模叠出来的开销

现象：跑 64 个节点的数据没问题，换到几百个节点的数据直接 CUDA Out Of Memory。

原因：图卷积的空间聚合复杂度是 O(B * T * N^2)，节点数一涨，显存消耗是平方级放大。

解决：常用手段有三个——减小 batch_size、减小 history 窗口、把邻接矩阵换成稀疏矩阵表示并配合 scatter 操作聚合。如果模型本身必须用全量图，我一般先把预测任务拆成「近邻子图」来做，限制每个节点只聚合距离最近的 K 个邻居，效果下降不多但显存省一大截。

5.5 用测试集反复调参

现象：模型在测试集上指标一路变好，看起来每一步改动都有提升，但换数据就崩。

原因：测试集的信息被用在了调参决策里，本质上成了第二个验证集，早停点也被污染了。

解决：调试阶段只用训练集和验证集，测试集只在全部确定后才碰一次。即使这样，真正上线前也建议在时间上再做一次「滚动验证」——用过去一段数据训练，预测未来一段，然后窗口滑动重复，这样才接近真实的使用方式。

6. 二次开发：换数据集、改图结构、加外部特征的改造路径

先聊最常做的改造：换成自己的数据。当前 Dataset 接收的是一个[T, N]矩阵，CSV 的格式是每行一个时间戳、每列一个节点。如果你的数据是长表格式，需要先pivot成宽表再喂进来。邻接矩阵如果不知道物理连接关系，可以用节点经纬度计算距离矩阵，然后用高斯核生成，配置里只需要换一下adj_path，utils.py里的build_adj_from_distance可以直接复用。

第二个值得改的点是图结构。当前用的是固定的静态邻接矩阵，一旦节点之间的影响关系本身就随时间变化，比如节假日交通流的方向模式完全改变，静态图就会成为瓶颈。常见做法是自适应图：用两个随机初始化的节点嵌入矩阵，通过矩阵乘法和 softmax 动态生成邻接矩阵：

node_emb1 = nn.Parameter(torch.randn(num_nodes, embed_dim)) node_emb2 = nn.Parameter(torch.randn(num_nodes, embed_dim)) adj = torch.softmax(torch.relu(node_emb1 @ node_emb2.T), dim=-1)

这样图结构不再是模型超参而是可学习的参数，能从数据里学出哪些节点应该相互影响。如果项目里想加外部特征（天气、节假日、事件），建议在时间卷积之后、图卷积之前拼进特征维度：temporal_out和external_embedding分别在特征维拼接，再喂给图卷积，而不是直接加在输入层。这样模型先提取完时间模式，外部信息再参与空间聚合，效果更可控。

从那以后我每次复现或改造时序模型，都会强制走一遍这套流程：先跑一个 1 个 epoch 的过拟合测试确认 pipeline 通，再检查归一化是否只基于训练集，最后才上完整训练。这套功夫救了不少次，希望帮到你。

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