Echo-Ji / HeST

Pytorch implementation of the HeST framework.

HeST (a.k.a ST-SSL)

This is a Pytorch implementation of the Heterogeneity-aware Spatial Temporal (HeST) framework in the following paper:

• J. Ji, J. Wang, C. Huang, et al. "Spatio-Temporal Self-Supervised Learning for Traffic Flow Prediction". in Thirty-Seventh AAAI Conference on Artificial Intelligence, 2023.

Please visit the newer version at ST-SSL repo.

Requirement

We build this project by Python 3.8 with the following packages:

numpy==1.21.2
pandas==1.3.5
PyYAML==6.0
torch==1.10.1

Model training and Evaluation

If the environment is ready, please run the following commands to train model on the specific dataset from {NYCBike1, NYCBike2, NYCTaxi, BJTaxi}.

>> cd HeST
>> ./runme 0 NYCBike1   # 0 gives the gpu id

This repo contains the NYCBike1 data. If you are interested in other datasets, please download from ST-SSL_Dataset.