This is the MXNet implementation of STFGNN in the paper: [Spatial-Temporal Fusion Graph Neural Networks for Traffic Flow Forecasting, AAAI 2021] (https://arxiv.org/abs/2012.09641). This framework is built based on framework of STSGCN(AAAI-20). Being familiar with its pipeline is strongly recommended.
- python 3
- see
requirements.txt
STFGNN is implemented on eight public traffic datasets.
- METR-LA and PEMS-BAY from DCRNN (ICLR-18). Download the data from Google Drive or Baidu Yun.
Then using preprocessing pipeline form Graph WaveNet (IJCAI-19):
# Create data directories
mkdir -p data/{METR-LA,PEMS-BAY}
# METR-LA
python generate_training_data.py --output_dir=data/METR-LA --traffic_df_filename=data/metr-la.h5
# PEMS-BAY
python generate_training_data.py --output_dir=data/PEMS-BAY --traffic_df_filename=data/pems-bay.h5
-
PEMS03, PEMS04, PEMS07 and PEMS08 from STSGCN (AAAI-20). Download the data STSGCN_data.tar.gz with password:
p72zand uncompress data file usingtar -zxvf data.tar.gz -
PeMSD7(Middle) and PeMSD7(Large) from STGCN (IJCAI-18). Download the data PeMSD7_data.zip with password:
9527and uncompress data file usingunzip PeMSD7_data.zip
The temporal graph of eight traffic dataset are available at ./data/adj_xxx_00x.csv. If traffic data is available, its temporal graph could also be generated by code:
cd ./data/
python fastDTW_adj_gen.py
METR-LA and PEMS-BAY:
python main_4n0_3layer_12T_712_res_npz.py --config config/XXXX/individual_3layer_12T.json
PEMS03, PEMS04, PEMS07, PEMS08 and PeMSD7M, PeMSD7L:
python main_4n0_3layer_12T_res.py --config config/XXXX/individual_3layer_12T.json
The authors would like to thank Prof Huaiyu Wan for his nice email interaction during submission of this paper, also like to thank Chao Song for his great mxnet implementation of STSGCN.