STLGRU: Spatio-Temporal Lightweight Graph GRU for Traffic Flow Prediction

This is the official implementation of STLGRU: Spatio-Temporal Lightweight Graph GRU for Traffic Flow Prediction:
Kishor Kumar Bhaumik, Fahim Faisal Niloy, Saif mahmud and Simon S. Woo STLGRU: Spatio-Temporal Lightweight Graph GRU for Traffic Flow Prediction.

Dependency can be installed using the following command:

pip install -r requirement.txt

Data Preparation

Download the dataset(PEMS03, PEMS04, PEMS07, PEMS08) from here, Baidu Drive, and the password is 1s5t. Download METR-LA and PEMS-BAY data from Google Drive or Baidu Yun links provided by DCRNN.

Process raw data for METR-LA and PEMS-BAY

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

Train Commands for PEMS08

python train.py --device cuda:0

Test Commands for PEMS08

python test.py --checkpoint garage8/PEMS08_epoch_158_16.75.pth --batch_size 1 --device cuda:0