STGCN实战：从零构建PyTorch时空图卷积网络预测城市交通流

关注

VIP免费原创最新推荐文章于 2026-09-30 09:00:11 发布 246 阅读 AI 写同款· GEO 优化 ›

别再只用CNN+LSTM了！用PyTorch复现STGCN搞定交通流量预测（附完整代码）本文详细介绍了如何使用PyTorch实现STGCN（时空图卷积网络）进行交通流量预测，包括环境配置、数据预处理、模型架构解析及训练优化技巧。STGCN通过结合空间和时间特征，显著提升了预测准确性，适用于智慧城市建设中的交通管理场景。阅读详情

1. 从交通数据到图结构：理解STGCN的输入

大家好，我是老张，一个在AI和智能硬件领域摸爬滚打了十多年的工程师。今天，我想和大家聊聊一个听起来很酷、用起来更酷的技术——时空图卷积网络（STGCN），并且手把手地带你用PyTorch从零开始，构建一个能预测城市交通流的模型。别被“时空”、“图卷积”这些词吓到，我会用最直白的方式，让你理解它到底在做什么，以及我们怎么一步步把它做出来。

想象一下，你每天开车上班，最头疼的是什么？十有八九是堵车。交通管理部门也想预测未来的车流速度，好提前疏导。传统的办法，比如看历史平均速度，在复杂的城市路网面前常常失灵，因为交通流在时间和空间上都是相互影响的。一个路口的拥堵，会像涟漪一样扩散到周边道路；而早高峰的规律，又和晚高峰截然不同。STGCN就是为了解决这种“时空耦合”问题而生的。

那么，我们怎么把一条条马路、一个个路口变成模型能“吃”下去的数据呢？这就是第一步，也是至关重要的一步：构建图（Graph）。在STGCN眼里，整个城市交通网络就是一张巨大的图。每个路口或者传感器检测点，就是图中的一个“节点”（Node）。连接两个路口的道路，就是图中的一条“边”（Edge）。这比把城市地图简单切成小方格（像处理图像那样）要合理得多，因为它忠实反映了道路之间真实的连接关系。

我们通常用一个叫邻接矩阵（Adjacency Matrix） 的数学工具来表示这张图。假设我们有3个路口A、B、C。如果A和B之间有路直接相连，我们就在矩阵的对应位置记上1（或者用道路距离、车流相关性计算一个权重），否则就是0。这样，一个N个路口的网络，就对应一个N×N的矩阵。这个矩阵，就是模型理解空间关系的“地图”。

接下来是每个节点的特征。在交通预测里，最常用的特征就是速度、流量、密度。以经典的PeMS数据集为例，它记录了加州高速公路上大量检测器每5分钟的速度和流量。所以，在某个时刻，每个节点（检测器）都有一个特征向量，比如[速度，流量]。我们把连续多个时间片（比如过去1小时，12个5分钟）的数据堆叠起来，就形成了一个三维的数据块：[节点数量，特征数量，历史时间步长度]。

为了让模型训练得更稳定，我们得先给数据“洗个澡”，做一下标准化。最常用的方法是Z-score标准化，也就是减去所有数据的平均值，再除以标准差。这样，数据就会乖乖地服从均值为0、标准差为1的分布。处理完的数据，我们再按时间顺序，用一个滑动窗口截取出一个个样本。例如，用过去12个时间步（1小时）的数据作为输入（X），去预测未来3个时间步（15分钟）的数据作为目标（Y）。这一步的代码逻辑很清晰：

def generate_dataset(X, num_timesteps_input, num_timesteps_output):
    """
    X: 原始数据，形状为 (num_nodes, num_features, total_timesteps)
    返回：样本特征和标签
    """
    indices = [(i, i + num_timesteps_input + num_timesteps_output)
               for i in range(X.shape[2] - (num_timesteps_input + num_timesteps_output) + 1)]

    features, target = [], []
    for i, j in indices:
        # 取出过去的时间步作为特征
        features.append(X[:, :, i: i + num_timesteps_input].transpose((0, 2, 1)))
        # 取出未来的时间步作为要预测的目标（这里以预测速度为例，取第一个特征通道）
        target.append(X[:, 0, i + num_timesteps_input: j])
    return torch.from_numpy(np.array(features)), torch.from_numpy(np.array(target))

经过这一系列操作，原始的、带有时空信息的表格数据，就被我们成功地转换成了图神经网络最喜欢的“图结构数据”。这相当于为模型准备好了食材，接下来就该看看我们的“主厨”STGCN是怎么烹饪这道大餐的了。

2. 庖丁解牛：STGCN的核心模块搭建

理解了数据长什么样，我们就能动手搭建模型的核心部件了。STGCN的巧妙之处在于，它没有使用复杂的循环神经网络（RNN）来处理时间序列，而是全部采用了卷积操作，这使得它的训练速度飞快，并且能更好地捕捉长期依赖。整个模型像搭积木，主要由两种基础模块构成：TimeBlock（时间块） 和 STGCNBlock（时空块）。