空间-时间图卷积网络

关于

这是论文《空间-时间图卷积网络：一种用于交通预测的深度学习框架》中STGCN的PyTorch实现。

论文

https://arxiv.org/abs/1709.04875

引用

@inproceedings{10.5555/3304222.3304273,
author = {Yu, Bing and Yin, Haoteng and Zhu, Zhanxing},
title = {Spatio-Temporal Graph Convolutional Networks: A Deep Learning Framework for Traffic Forecasting},
year = {2018},
isbn = {9780999241127},
publisher = {AAAI Press},
booktitle = {Proceedings of the 27th International Joint Conference on Artificial Intelligence},
pages = {3634–3640},
numpages = {7},
series = {IJCAI'18}
}

相关工作

1. TCN：序列建模中通用卷积和循环网络的实证评估
2. GLU和GTU：基于门控卷积网络的语言建模
3. ChebNet：具有快速局部化谱滤波的图卷积神经网络
4. GCN：基于图卷积网络的半监督分类

相关代码

1. TCN：https://github.com/locuslab/TCN
2. ChebNet：https://github.com/mdeff/cnn_graph
3. GCN：https://github.com/tkipf/pygcn

数据集

来源

1. METR-LA：DCRNN作者的Google Drive
2. PEMS-BAY：DCRNN作者的Google Drive
3. PeMSD7(M)：STGCN作者的GitHub仓库

预处理

使用来自ChebNet的公式：

模型结构

我的代码与作者代码的不同之处

1. 修复了错误
2. 添加了早停法
3. 添加了Dropout方法
4. 提供了一组不同的超参数
5. 提供了两种不同类型的图卷积配置文件（ChebyGraphConv和GraphConv）
6. 增加了METR-LA和PEMS-BAY数据集
7. 采用了不同的数据预处理方法

需求

安装依赖项：

pip3 install -r requirements.txt

STGCN 快速上手指南

本指南基于 PyTorch 实现的时空图卷积网络（STGCN），主要用于交通流量预测任务。

环境准备

• 操作系统：Linux / macOS / Windows
• Python 版本：建议 Python 3.6+
• 核心框架：PyTorch
• 其他依赖：详见项目根目录下的 requirements.txt

提示：国内用户建议使用清华或阿里镜像源加速 pip 安装，例如在后续安装命令中添加 -i https://pypi.tuna.tsinghua.edu.cn/simple。

安装步骤

1. 克隆项目代码到本地：

   git clone https://github.com/hazdzz/STGCN.git
   cd STGCN
   

2. 安装所需依赖包：

   pip3 install -r requirements.txt
   

   （国内加速版）

   pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 数据准备： 本项目支持 METR-LA、PEMS-BAY 和 PeMSD7(M) 数据集。需手动下载数据并放置于指定目录（参考 README 中的 Dataset 部分）。

   数据预处理采用 ChebNet 公式计算加权邻接矩阵，请确保数据格式符合模型输入要求。

基本使用

以下是最简单的训练与评估流程示例（具体脚本名称请以仓库实际文件为准，通常为 main.py 或类似入口）：

1. 配置参数： 修改项目提供的配置文件（如 config_peaks.py 或 .yaml 文件），设置数据集路径、图卷积类型（ChebyGraphConv 或 GraphConv）及超参数。

2. 启动训练：

   python main.py --config config/metrla_config.py --mode train
   

3. 模型评估： 训练完成后，使用保存的模型权重进行预测评估：

   python main.py --config config/metrla_config.py --mode eval --weights logs/best_model.pth
   

注意：该实现相比原论文代码增加了早停（Early Stopping）、Dropout 机制，并修复了部分 Bug，支持更多数据集和图卷积变体。