diff --git a/README.md b/README.md index 2c4c6ad..416e24d 100644 --- a/README.md +++ b/README.md @@ -1,27 +1,33 @@ # MTGNN + This is a PyTorch implementation of the paper: [Connecting the Dots: Multivariate Time Series Forecasting with Graph Neural Networks](https://arxiv.org/abs/2005.11650), published in KDD-2020. ## Requirements + The model is implemented using Python3 with dependencies specified in requirements.txt + ## Data Preparation + ### Multivariate time series datasets Download Solar-Energy, Traffic, Electricity, Exchange-rate datasets from [https://github.com/laiguokun/multivariate-time-series-data](https://github.com/laiguokun/multivariate-time-series-data). Uncompress them and move them to the data folder. ### Traffic datasets + Download the METR-LA and PEMS-BAY dataset from [Google Drive](https://drive.google.com/open?id=10FOTa6HXPqX8Pf5WRoRwcFnW9BrNZEIX) or [Baidu Yun](https://pan.baidu.com/s/14Yy9isAIZYdU__OYEQGa_g) provided by [Li et al.](https://github.com/liyaguang/DCRNN.git) . Move them into the data folder. ``` - # 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 +python generate_training_data.py --output_dir=data/PEMS-BAY --traffic_df_filename=data/pems-bay.h5 ``` ## Model Training @@ -32,14 +38,15 @@ python generate_training_data.py --output_dir=data/PEMS-BAY --traffic_df_filenam ``` python train_single_step.py --save ./model-solar-3.pt --data ./data/solar_AL.txt --num_nodes 137 --batch_size 4 --epochs 30 --horizon 3 -#sampling +# sampling python train_single_step.py --num_split 3 --save ./model-solar-sampling-3.pt --data ./data/solar_AL.txt --num_nodes 137 --batch_size 16 --epochs 30 --horizon 3 ``` -* Traffic + +* Traffic ``` python train_single_step.py --save ./model-traffic3.pt --data ./data/traffic.txt --num_nodes 862 --batch_size 16 --epochs 30 --horizon 3 -#sampling +# sampling python train_single_step.py --num_split 3 --save ./model-traffic-sampling-3.pt --data ./data/traffic --num_nodes 321 --batch_size 16 --epochs 30 --horizon 3 ``` @@ -47,7 +54,7 @@ python train_single_step.py --num_split 3 --save ./model-traffic-sampling-3.pt - ``` python train_single_step.py --save ./model-electricity-3.pt --data ./data/electricity.txt --num_nodes 321 --batch_size 4 --epochs 30 --horizon 3 -#sampling +# sampling python train_single_step.py --num_split 3 --save ./model-electricity-sampling-3.pt --data ./data/electricity.txt --num_nodes 321 --batch_size 16 --epochs 30 --horizon 3 ``` @@ -58,12 +65,15 @@ python train_single_step.py --save ./model/model-exchange-3.pt --data ./data/exc #sampling python train_single_step.py --num_split 3 --save ./model-exchange-3.pt --data ./data/exchange_rate.txt --num_nodes 8 --subgraph_size 2 --batch_size 16 --epochs 30 --horizon 3 ``` + ### Multi-step + * METR-LA ``` python train_multi_step.py --adj_data ./data/sensor_graph/adj_mx.pkl --data ./data/METR-LA --num_nodes 207 ``` + * PEMS-BAY ```