This is the MXNet implementation of ST-MetaNet in the following paper:
Dependency can be installed using the following command:
pip install -r requirements.txt
After that, change the backend deep learning framework for dgl:
echo 'export DGLBACKEND=mxnet' >> ~/.bashrc
. ~/.bashrc
Unzip the data files in:
flow-prediction/data.zip
traffic-prediction/data.zip
(the data is pre-processed in repository DCRNN)The flow data is collected from 32x32 grids of Beijing city.
(row, colume, feature_id)
. This data represents the node feature for each grid. It consists of POI features and road features saved in BJ_POI.h5 and BJ_ROAD.h5, respectively. (date, hour, row, colume, flow_type)
. The flow_type
represents inflow or outflow of the region. (grid_1, grid_2, feature_id)
. This data represents the edge feature from grid_1 to grid_2.The description please refers to the repository of DCRNN.
Given the flow prediction task as example (the traffic prediction task is exactly the same as the flow prediction task):
cd flow-prediction/
.src/model_setting
, saved as yaml format. Three models are provided: seq2seq
, gat-seq2seq
, and st-metanet
. Other baselines refers to DCRNN and ST-ResNet, respectively.param/
. There are two types of files in this folder:
model.yaml
: the model training log (the result on evaluation dataset of each epoch). This file records the number of the best epoch for the model.model-xxxx.params
: the saved model parameters of the best evaluation epoch, where xxxx
is the epoch number.cd src/
.python train.py --file model_settting/[model_name].yaml --gpus [gpu_ids] --epochs [num_epoch]
. The code will firstly load the best epoch from params/
, and then train the models for [num_epoch]
. Our code can be trained with multiple gpus. An example of [gpu_ids]
is 0,1,2,3
if you have four gpus. But we recommend to use a single gpu to train & evaluate the model if possible (currently, our implementation using DGL library for meta graph attention is not efficient when using multiple gpus). param/
, otherwise the code will train from the best pre-trained (saved) model.python train.py --file model_setting/st-metanet.yaml --gpus 0 --epochs 200
. python train.py --file model_setting/st-metanet.yaml --gpus 0,1 --epochs 200
.python train.py --file model_setting/st-metanet.yaml --gpus 0 --epochs 0
. The code will directly give evaluation results on evaluation dataset and test dataset, without training.If you find this repository, e.g., the code and the datasets, useful in your research, please cite the following paper:
ST-MetaNet is released under the MIT License (refer to the LICENSE file for details).