This is the PyTorch implementation of AutoSTG in the following paper:
- Zheyi Pan, Songyu Ke, Xiaodu Yang, Yuxuan Liang, Yong Yu, Junbo Zhang and Yu Zheng. AutoSTG: Neural Architecture Search for Predictions of Spatio-Temporal Graphs. 2021. In Proceedings of the Web Conference.
- System: Ubuntu 16.04
- Language: Python 3.5
- Devices: a single GeForce GTX 1080 Ti CPU
- numpy == 1.19.1
- pandas == 1.1.1
- torch == 1.1.0
- torchvision == 0.3.0
- tables == 3.6.1
- ruamel.yaml == 0.16.12
Unzip dataset/dataset.zip with the following command:
cd dataset
unzip ./dataset.zip
The description please refers to the repository of DCRNN.
src/train_on_gpu0.sh gives an example to search and train the model on the two datasets:
cd src/
.- The settings of the models are in the folder model, saved as yaml format.
- For METR-LA: METR_LA_AutoSTG.yaml
- For PEMS-BAY: METR_LA_AutoSTG.yaml
- All trained model will be saved in
param/
. - Searching and training with the given shell script:
cd src/
.bash train_on_gpu0.sh
. The code will firstly load the best epoch fromparams/
, and then train the models for[epoch]
.
src/test_on_gpu0.sh gives an example to test the model on the two datasets using the trained models. Here are the instructions to execute the shell script:
cd src/
bash test_on_gpu0.sh
.
Note that: The given pre-trained models are trained under PyTorch 1.1.0 and can be loaded under other compatible versions (PyTorch 1.5.0 and 1.6.0 with Python 3.7 are tested) with expected behaviors (i.e., the same predicting error).
If you find this repository, e.g., the code and the datasets, useful in your research, please cite the following paper:
- Zheyi Pan, Songyu Ke, Xiaodu Yang, Yuxuan Liang, Yong Yu, Junbo Zhang and Yu Zheng. AutoSTG: Neural Architecture Search for Predictions of Spatio-Temporal Graphs. 2021. In Proceedings of the Web Conference.
AutoSTG is released under the MIT License (refer to the LICENSE file for details).