This is the code of ST-MoE for paper 'ST-MoE: Spatio-Temporal Mixture-of-Experts for Debiasing in Traffic Prediction'. We present here the code of the ST-MoE framework and use Graph WaveNet as the base model example.
- python 3
- see
requirements.txt
-
Download METR-LA and PEMS-BAY data from Google Drive or Baidu Yun links provided by DCRNN.
# 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
Base model train and predict:
python train.py
ST-MoE train and predict:
python main_moe.py