This is a PyTorch implementation of Propagation Delay-aware Dynamic Long-range Transformer for Traffic Flow Prediction (PDFormer) for traffic flow prediction, as described in our paper: Jiawei Jiang*, Chengkai Han*, Wayne Xin Zhao, Xiaohan Jiang, JingyuanWang, Propagation Delay-aware Dynamic Long-range Transformer for Traffic Flow Prediction, AAAI2023.

* Equal Contributions.

Our code is based on Python version 3.9.7 and PyTorch version 1.10.1. Please make sure you have installed Python and PyTorch correctly. Then you can install all the dependencies with the following command by pip:

pip install -r requirements.txt

The dataset link is Google Drive. You can download the datasets and place them in the raw_data directory.

All 6 datasets come from the LibCity repository, which are processed into the atomic files format. The only difference with the datasets provided by origin LibCity repository here is that the filename of the datasets are differently.

Note that our model would calculate a DTW matrix and a traffic pattern set for each dataset, which is time-consuming. Therefore, we have provided DTW matrices and traffic pattern sets of all datasets in ./libcity/cache/dataset_cache/.

You can train and test PDFormer through the following commands for 6 datasets.

python run_model.py --task traffic_state_pred --model PDFormer --dataset PeMS04 --config_file PeMS04
python run_model.py --task traffic_state_pred --model PDFormer --dataset PeMS08 --config_file PeMS08
python run_model.py --task traffic_state_pred --model PDFormer --dataset PeMS07 --config_file PeMS07
python run_model.py --task traffic_state_pred --model PDFormer --dataset NYCTaxi --config_file NYCTaxi --evaluator TrafficStateGridEvaluator
python run_model.py --task traffic_state_pred --model PDFormer --dataset CHIBike --config_file CHIBike --evaluator TrafficStateGridEvaluator
python run_model.py --task traffic_state_pred --model PDFormer --dataset T-Drive --config_file T-Drive --evaluator TrafficStateGridEvaluator

If you have trained a model as above and only want to test it, you can set it as follows (taking PeMS08 as an example, assuming the experiment ID during training is $ID):

python run_model.py --task traffic_state_pred --model PDFormer --dataset PeMS08 --config_file PeMS08 --train false --exp_id $ID

Code based on LibCity framework development, an open source library for traffic prediction.

If you find the paper useful, please cite as following:

@inproceedings{pdformer,
  title={PDFormer: Propagation Delay-aware Dynamic Long-range Transformer for Traffic Flow Prediction},
  author={Jiawei Jiang and 
  		  Chengkai Han and 
  		  Wayne Xin Zhao and 
  		  Jingyuan Wang},
  booktitle = {{AAAI}},
  publisher = {{AAAI} Press},
  year      = {2023}
}

If you find LibCity useful, please cite as following:

@inproceedings{libcity,
  author    = {Jingyuan Wang and
               Jiawei Jiang and
               Wenjun Jiang and
               Chao Li and
               Wayne Xin Zhao},
  title     = {LibCity: An Open Library for Traffic Prediction},
  booktitle = {{SIGSPATIAL/GIS}},
  pages     = {145--148},
  publisher = {{ACM}},
  year      = {2021}
}