AAAI-2021 Spatial-Temporal Fusion Graph Neural Networks for Traffic Flow Forecasting

Requirements

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

Data Preparation

STFGNN is implemented on those several public traffic datasets.

PEMS03, PEMS04, PEMS07 and PEMS08 from STSGCN (AAAI-20). Download the data STSGCN_data.tar.gz with password: p72z and uncompress data file usingtar -zxvf data.tar.gz

Model Training

PEMS03, PEMS04, PEMS07, PEMS08:

python main_4n0_3layer_12T_res.py --config config/XXXX/individual_3layer_12T.json

Train your own dataset

(1) Temporal Graph Construction

If traffic data is available, its temporal graph could also be generated by code:

cd ./data/
python Temporal_Graph_gen.py

The shape of input traffic data should be "(Total_Time_Steps, Node_Number). For example, in PEMS08 dataset, it has 170 roads and 62 days data. Thus its shape is (62*288, 170).

The calculation uses CPU, which should be prepared for enough computation resources.

(2) Configuration of dataset

vi ../config/newdataset/individual_3layer_12T.json

where could set path of spatial graph, temporal graph, input data and other parameters:

    "adj_filename": spatial adjacency matrix,
    "adj_dtw_filename": temporal adjacency matrix,
    "graph_signal_matrix_filename": input data (multivariate time series),
    "num_of_vertices": node number

It could be a spatial information free framework when you set "adj_filename" equals to adj_dtw_filename.

(3) Model training

python main_4n0_3layer_12T_res.py --config config/XXXX/individual_3layer_12T.json

Acknowledgments

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.

Update (January, 2021):

The numerical results on PEMS-BAY, METR-LA, PeMSD7(M) and PeMSD7(L) are not correct because of different loss metric. After discussion with AAAI-2021 comitte by e-mail, results of PEMS0X are kept. The final camera-ready arxiv would be modifed as soon as possible.

Very sorry for misunderstanding results on partial datasets.

mcdragon/STFGNN