Decoupled Spatial-Temporal Graph Neural Networks

Code for our paper: Decoupled Dynamic Spatial-Temporal Graph Neural Network for Traffic Forecasting.

Traffic forecasting is an indispensable part of building intelligent transportation systems and has remained an enduring research topic in academia and industry. Recently, spatial-temporal (ST) graph neural networks have been proposed to model complex temporal and spatial dependencies in traffic data, and have made significant progress. However, existing models simply connect the spatial and temporal models in series, which ignores the special characteristics of spatial and temporal information. Moreover, the serial connection structure may cause error accumulation, leading to worse model performance.

To address the problem, we propose a novel spatial-temporal framework consisting of a unique spatial gate and a residual decomposition mechanism, which is capable of facilitating the sufficient learning process of downstream modules via decoupling spatial and temporal signals. With the decoupled ST framework, we also propose Decoupled Dynamic Spatial-Temporal Graph Neural Network (D$^2$STGNN in short), which aptly captures spatial-temporal dependencies and is enhanced by a dynamic graph learning module, for learning the dynamic characteristics of traffic networks. Extensive experiments on four real-world traffic datasets demonstrate the effectiveness of the proposed method.

1. Run the model and reproduce the result?

1.1 Data Preparation

For convenience, we package these datasets used in our model in Google Drive or BaiduYun.

They should be downloaded to the code root dir and replace the raw_data and sensor_graph folder in the datasets folder by:

cd /path/to/project
unzip raw_data.zip -d ./datasets/
unzip sensor_graph.zip -d ./datasets/
rm {sensor_graph.zip,raw_data.zip}
mkdir log output

Alterbatively, the datasets can be found as follows:

• METR-LA and PEMS-BAY: These datasets were released by DCRNN[1]. Data can be found in its GitHub repository, where the sensor graphs are also provided.

• PEMS03 and PEMS04: These datasets were released by ASTGCN[2] and ASTGNN[3]. Data can also be found in its GitHub repository.

1.2 Data Process

python datasets/raw_data/$DATASET_NAME/generate_training_data.py

Replace $DATASET_NAME with one of METR-LA, PEMS-BAY, PEMS04, PEMS08.

The processed data is placed in datasets/$DATASET_NAME.

1.3 Training the Model

python main.py --dataset=$DATASET_NAME

E.g., python main.py --dataset=METR-LA.

1.4 Load a Pretrained Model

Check the config files of the dataset in configs/$DATASET_NAME, and set the startup args to test mode.

Download the pre-trained model files into the output folder and run the command line in 1.3.

1.5 Results and Visualization

2. More QA?

Any issues are welcome.

3. To Do

• Add results and visualization in this readme.
• Add BaiduYun links.
• Add pretrained model.
• 添加中文README

References

[1] Atwood J, Towsley D. Diffusion-convolutional neural networks[J]. Advances in neural information processing systems, 2016, 29: 1993-2001.

[2] Guo S, Lin Y, Feng N, et al. Attention based spatial-temporal graph convolutional networks for traffic flow forecasting[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2019, 33(01): 922-929.

[3] Guo S, Lin Y, Wan H, et al. Learning dynamics and heterogeneity of spatial-temporal graph data for traffic forecasting[J]. IEEE Transactions on Knowledge and Data Engineering, 2021.