The code repository is now available. This is the official repository of our work STG-Mamba, the paper is currently posted on ArXiV.
- PyTorch==1.11.0
- Python==3.8.10
- numpy==1.22.4
- pandas==2.0.3
- einops==0.7.0
- argparse
- dataclasses
- typing
- time
- math
/root/STG_Mamba/
|
|Know_Air_Dataset/
| |knowair_adj_mat.npy
| |knowair_temperature.csv
|PEMS04_Dataset/
| |pems04_adj.npy
| |pems04_flow.csv
|HZ_Metro_Dataset/
| |hzmetro_adj.npy
| |hzmetro_flow.csv
|
|main.py
|modules.py
|prepare.py
|STGMamba.py
|train_STGmamba.py
|train_rnn.py
(1) PEMS04: PEMS04 dataset is among the most popular benchmark in ST Data Mining & Traffic Forecasting. You can find the source data at here.
(2) KnowAir: KnowAir is an open-sourced weather dataset introduced by Shuo Wang et al. at SIGSPATIAL' 20. We've already put the knowair data in this repository, with the extracted Graph Adjacency Matrix based on weather station's geographical location. You can also access the raw data at here.
(3) HZ-Metro: HZ-Metro is an open-sourced Metro Crowd-Flow dataset collected from 80 Metro Stations in HangZhou, China. You can find the raw data at here. Also, there is a reference link of the 80 Metro Stations' ID and their Name in real-world.
Using the following commands to Train/Test STG_Mamba model on KnowAir, PEMS04, HZ_Metro respectively. You can also optionally change the model parameters in the file main.py.
# KnowAir
python main.py -dataset=know_air -model=STGmamba -mamba_features=184
# PEMS04
python main.py -dataset=pems04 -model=STGmamba -mamba_features=307
#HZ_Metro
python main.py -dataset=hz_metro -model=STGmamba -mamba_features=80If you find this repository useful in your own research, please cite our work.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent