OpenCity: Open Spatio-Temporal Foundation Models for Traffic Prediction

A pytorch implementation for the paper: OpenCity: Open Spatio-Temporal Foundation Models for Traffic Prediction

Zhonghang Li, Long Xia, Lei Shi, Yong Xu, Dawei Yin, Chao Huang* (*Correspondence)

undefinedData Intelligence Lab@University of Hong Kong, South China University of Technology, Baidu Inc

This repository hosts the code, data, and model weights of OpenCity.

🎉 News

• [2024.08.21] Release the full paper.
• [2024.08.20] Add video.
• [2024.08.15] 🚀🚀 Release the code, model weights and datasets of OpenCity.
• [2024.08.15] Release baselines codes.

🎯🎯📢📢 We upload the models and data used in our OpenCity on 🤗 Huggingface. We highly recommend referring to the table below for further details:

👉 TODO

…

Introduction

In this work, we aim to unlock new possibilities for building versatile, resilient and adaptive spatio-temporal foundation models for traffic prediction. To achieve this goal, we introduce a novel foundation model, named OpenCity, that can effectively capture and normalize the underlying spatio-temporal patterns from diverse data characteristics, facilitating zero-shot generalization across diverse urban environments. OpenCity integrates the Transformer architecture with graph neural networks to model the complex spatio-temporal dependencies in traffic data. By pre-training OpenCity on large-scale, heterogeneous traffic datasets, we enable the model to learn rich, generalizable representations that can be seamlessly applied to a wide range of traffic forecasting scenarios. Experimental results demonstrate that OpenCity exhibits exceptional zero-shot predictive performance in various traffic prediction tasks.

Main Results

undefinedOutstanding Zero-shot Prediction Performance. OpenCity achieves significant zero-shot learning breakthroughs, outperforming most baselines even without fine-tuning. This highlights the approach’s robustness and effectiveness at learning complex spatio-temporal patterns in large-scale traffic data, extracting universal insights applicable across downstream tasks.

Demo Video

https://github.com/user-attachments/assets/39265dc5-0126-483b-951e-518c6cb210e0

├── conf/
│   ├── AGCRN/
│   │   └── AGCRN.conf
│   ├── ASTGCN/
│   │   └── ASTGCN.conf
│   ├── general_conf/
│   │   ├── global_baselines.conf
│   │   └── pretrain.conf
│   ├── GWN/
│   │   └── GWN.conf
│   ├── MSDR/
│   │   └── MSDR.conf
│   ├── MTGNN/
│   │   └── MTGNN.conf
│   ├── OpenCity/
│   │   └── OpenCity.conf
│   ├── PDFormer/
│   │   └── PDFormer.conf
│   ├── STGCN/
│   │   └── STGCN.conf
│   ├── STSGCN/
│   │   └── STSGCN.conf
│   ├── STWA/
│   │   └── STWA.conf
│   └── TGCN/
│       └── TGCN.conf
├── data/
│   ├── generate_ca_data.py
│   └── README.md
├── lib/
│   ├── data_process.py
│   ├── logger.py
│   ├── metrics.py
│   ├── Params_predictor.py
│   ├── Params_pretrain.py
│   ├── predifineGraph.py
│   └── TrainInits.py
├── model/
│   ├── AGCRN/
│   │   ├── AGCN.py
│   │   ├── AGCRN.py
│   │   ├── AGCRNCell.py
│   │   └── args.py
│   ├── ASTGCN/
│   │   ├── args.py
│   │   └── ASTGCN.py
│   ├── GWN/
│   │   ├── args.py
│   │   └── GWN.py
│   ├── MSDR/
│   │   ├── args.py
│   │   ├── gmsdr_cell.py
│   │   └── gmsdr_model.py
│   ├── MTGNN/
│   │   ├── args.py
│   │   └── MTGNN.py
│   ├── OpenCity/
│   │   ├── args.py
│   │   └── OpenCity.py
│   ├── PDFormer/
│   │   ├── args.py
│   │   └── PDFormer.py
│   ├── ST_WA/
│   │   ├── args.py
│   │   ├── attention.py
│   │   └── ST_WA.py
│   ├── STGCN/
│   │   ├── args.py
│   │   └── stgcn.py
│   ├── STSGCN/
│   │   ├── args.py
│   │   └── STSGCN.py
│   └── TGCN/
│       ├── args.py
│       └── TGCN.py
│   ├── Model.py
│   ├── BasicTrainer.py
│   ├── Run.py
└── model_weights/
    ├── OpenCity/
    └── README.md

conda create -n opencity python=3.9.13

conda activate opencity

# Torch (other versions are also ok)
pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

# Clone our OpenCity or download it
git clone https://github.com/HKUDS/OpenCity.git
cd OpenCity-main

# Install required libraries
pip install -r requirements.txt

# OpenCity-plus
python Run.py -mode pretrain -model OpenCity -save_pretrain_path OpenCity-plus2.0.pth -batch_size 4 --embed_dim 512 --skip_dim 512 --enc_depth 6

# OpenCity-base
python Run.py -mode pretrain -model OpenCity -save_pretrain_path OpenCity-base2.0.pth -batch_size 8 --embed_dim 256 --skip_dim 256 --enc_depth 3

# OpenCity-mini
python Run.py -mode pretrain -model OpenCity -save_pretrain_path OpenCity-mini2.0.pth -batch_size 16 --embed_dim 128 --skip_dim 128 --enc_depth 3

# Use OpenCity-plus to evaluate, please use only one dataset to test (e.g. dataset_use = ['PEMS07M'] in pretrain.config).
python Run.py -mode test -model OpenCity -load_pretrain_path OpenCity-plus.pth -batch_size 2 --embed_dim 512 --skip_dim 512 --enc_depth 6

# Run STGCN in ori mode
python Run.py -mode ori -model STGCN -batch_size 64 --real_value False

@misc{li2024opencity,
      title={OpenCity: Open Spatio-Temporal Foundation Models for Traffic Prediction}, 
      author={Zhonghang Li and Long Xia and Lei Shi and Yong Xu and Dawei Yin and Chao Huang},
      year={2024},
      eprint={2408.10269},
      archivePrefix={arXiv}
}