STMAE

Revealing the Power of Masked Autoencoders in Traffic Forecasting (CIKM'2024)

by Jiarui Sun et. al.

[Paper]

This repository contains the official PyTorch implementation of STMAE.

We propose Spatial-Temporal Masked AutoEncoders (STMAE), a plug-and-play framework designed to enhance existing spatial-temporal models on traffic forecasting.

Installation

1. Environment

Python/conda/mamba environment

Coming Soon!

2. Datasets

PEMS-03/04/07/08

We follow https://github.com/liuxu77/STGCL for dataset preparation. The generated data files should be placed inside ./data/pems_0<3/4/7/8> directories.

Evaluation

Run the following scripts to evaluate STMAE:

./test_all.sh

Training

Run the following scripts to train STMAE:

python train_pf.py --cfg <configuration_file_name>

Citation

If you find our work useful in your research, please consider citing our paper:

@inproceedings{sun2024revealing,
  title={Revealing the power of masked autoencoders in traffic forecasting},
  author={Sun, Jiarui and Fan, Yujie and Yeh, Chin-Chia Michael and Zhang, Wei and Chowdhary, Girish},
  booktitle={Proceedings of the 33rd ACM International Conference on Information and Knowledge Management},
  pages={4071--4075},
  year={2024}
}

Note: We borrow parts from AGCRN, DCRNN, MTGNN and STGCL.