GATraj: A Graph- and Attention-based Multi-Agent Trajectory Prediction Model

Code for "GATraj: A Graph- and Attention-based Multi-Agent Trajectory Prediction Model"

Environment

pip install -r requirements.txt

Train

The Default settings are to train on ETH-eth. Data cache and models will be in the subdirectory "./savedata/0/".

git clone https://github.com/mengmengliu1998/GATraj.git
cd GATraj
python train.py --test_set <dataset to train> --num_epochs 1000 --x_encoder_layers 3 --eta_min 1e-5  --batch_size 32\
  --learning_rate 5e-4  --randomRotate True --final_mode 20 --neighbor_thred 10\
  --using_cuda True --clip 1 --pass_time 2 --ifGaussian False --SR True --input_offset True 

Configuration files are also created after the first run, arguments could be modified through configuration files or command line. Priority: command line > configuration files > default values in script.

The datasets are selected on arguments '--test_set'. Five datasets in ETH/UCY are corresponding to the value of [0,1,2,3,4] ([eth, hotel, zara1, zara2, univ]).

Example

This command is to train model for ETH-eth

python train.py --test_set 0 --num_epochs 1000 --x_encoder_layers 3 --eta_min 1e-5  --batch_size 32\
  --learning_rate 5e-4  --randomRotate True --final_mode 20 --neighbor_thred 10\
  --using_cuda True --clip 1 --pass_time 2 --ifGaussian False --SR True --input_offset True

Test

We provide the trained model weights in the subdirectory "./savedata/". This command is to test model for ETH-eth, just add --phase test --load_model 1000 to the end of this training command.

python train.py --test_set 0 --num_epochs 1000 --x_encoder_layers 3 --eta_min 1e-5  --batch_size 32\
  --learning_rate 5e-4  --randomRotate True --final_mode 20 --neighbor_thred 10\
  --using_cuda True --clip 1 --pass_time 2 --ifGaussian False --SR True --input_offset True --phase test  --load_model 1000

Cite GATraj

If you find this repo useful, please consider citing our paper

@article{cheng2023gatraj,
  title={Gatraj: A graph-and attention-based multi-agent trajectory prediction model},
  author={Cheng, Hao and Liu, Mengmeng and Chen, Lin and Broszio, Hellward and Sester, Monika and Yang, Michael Ying},
  journal={ISPRS Journal of Photogrammetry and Remote Sensing},
  volume={205},
  pages={163--175},
  year={2023},
  publisher={Elsevier}
}

Reference

The code base of dataloader is heavily adopted from SR-LSTM. The visulization code base for nuScenes is adopted from PGP.