{"id":29248918,"url":"https://github.com/hkuds/flashst","last_synced_at":"2025-07-04T00:09:27.751Z","repository":{"id":241214783,"uuid":"804657920","full_name":"HKUDS/FlashST","owner":"HKUDS","description":"[ICML'2024] \"FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction\"","archived":false,"fork":false,"pushed_at":"2024-06-10T06:36:38.000Z","size":21113,"stargazers_count":22,"open_issues_count":1,"forks_count":1,"subscribers_count":2,"default_branch":"main","last_synced_at":"2024-06-10T08:24:28.386Z","etag":null,"topics":["pre-training","prompt-tuning","smart-cities","spatio-temporal-prediction","traffic-flow-prediction","urban-computing"],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2405.17898","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/HKUDS.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE.txt","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-05-23T02:34:57.000Z","updated_at":"2024-06-10T06:36:41.000Z","dependencies_parsed_at":"2024-06-05T08:43:16.343Z","dependency_job_id":null,"html_url":"https://github.com/HKUDS/FlashST","commit_stats":null,"previous_names":["hkuds/flashst"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/HKUDS/FlashST","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/HKUDS%2FFlashST","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/HKUDS%2FFlashST/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/HKUDS%2FFlashST/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/HKUDS%2FFlashST/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/HKUDS","download_url":"https://codeload.github.com/HKUDS/FlashST/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/HKUDS%2FFlashST/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":263421915,"owners_count":23464052,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["pre-training","prompt-tuning","smart-cities","spatio-temporal-prediction","traffic-flow-prediction","urban-computing"],"created_at":"2025-07-04T00:09:18.511Z","updated_at":"2025-07-04T00:09:27.656Z","avatar_url":"https://github.com/HKUDS.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction\n\n\u003cimg src='https://github.com/LZH-YS1998/GPT-ST_img/blob/main/FlashST_repo.png' /\u003e\n\nA pytorch implementation for the paper: [FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction](https://arxiv.org/abs/2405.17898)\u003cbr /\u003e  \n\n[Zhonghang Li](https://scholar.google.com/citations?user=__9uvQkAAAAJ), [Lianghao Xia](https://akaxlh.github.io/), [Yong Xu](https://scholar.google.com/citations?user=1hx5iwEAAAAJ), [Chao Huang](https://sites.google.com/view/chaoh)* (*Correspondence)\u003cbr /\u003e  \n\n**[Data Intelligence Lab](https://sites.google.com/view/chaoh/home)@[University of Hong Kong](https://www.hku.hk/)**, [South China University of Technology](https://www.scut.edu.cn/en/), PAZHOU LAB\n\n-----------\n\n## Introduction\n\n\u003cp style=\"text-align: justify\"\u003e\nIn this work, we introduce a simple and universal spatio-temporal prompt-tuning framework, which addresses the significant challenge posed by distribution shift in this field. \nTo achieve this objective, we present FlashST, a framework that adapts pretrained models to the specific characteristics of diverse downstream datasets, thereby improving generalization across various prediction scenarios.\nWe begin by utilizing a lightweight spatio-temporal prompt network for in-context learning, capturing spatio-temporal invariant knowledge and facilitating effective adaptation to diverse scenarios. Additionally, we incorporate a distribution mapping mechanism to align the data distributions of pre-training and downstream data, facilitating effective knowledge transfer in spatio-temporal forecasting. Empirical evaluations demonstrate the effectiveness of our FlashST across different spatio-temporal prediction tasks.\n\n\u003c/p\u003e\n\n![The detailed framework of the proposed FlashST.](https://github.com/LZH-YS1998/GPT-ST_img/blob/main/FlashST.png)\n\n-----------\n\u003cspan id='Usage'/\u003e\n\n\n\n## Getting Started\n\n\u003cspan id='all_catelogue'/\u003e\n\n### Table of Contents:\n* \u003ca href='#Code Structure'\u003e1. Code Structure\u003c/a\u003e\n* \u003ca href='#Environment'\u003e2. Environment \u003c/a\u003e\n* \u003ca href='#Run the codes'\u003e3. Run the codes \u003c/a\u003e\n\n****\n\n\n\u003cspan id='Code Structure'/\u003e\n\n### 1. Code Structure \u003ca href='#all_catelogue'\u003e[Back to Top]\u003c/a\u003e\n\n\n* **conf**: This folder includes parameter settings for FlashST (`config.conf`) as well as all other baseline models.\n* **data**: The documentation encompasses all the datasets utilized in our work, alongside prefabricated files and the corresponding file generation codes necessary for certain baselines.\n* **lib**: Including a series of initialization methods for data processing, as follows:\n\t* `data_process.py`: Load, split, generate data, normalization method, slicing, etc.\n    * `logger.py`: For output printing.\n\t* `metrics.py`: Method for calculating evaluation indicators.\n\t* `predifineGraph.py`: Predefined graph generation method.\n\t* `TrainInits.py`: Training initialization, including settings of optimizer, device, random seed, etc.\n* **model**: Includes the implementation of FlashST and all baseline models, along with the necessary code to support the framework's execution. The `args.py` script is utilized to generate the required prefabricated data and parameter configurations for different baselines. Additionally, the `SAVE` folder serves as the storage location for saving the pre-trained models.\n* **SAVE**: This folder serves as the storage location for saving the trained models, including pretrain, eval and ori.\n\n\n```\n│  README.md\n│  requirements.txt\n│\n├─conf\n│  ├─AGCRN\n│  ├─ASTGCN\n│  ├─FlashST\n│  │  │  config.conf\n│  │  │  Params_pretrain.py\n│  ├─GWN\n│  ├─MSDR\n│  ├─MTGNN\n│  ├─PDFormer\n│  ├─ST-WA\n│  ├─STFGNN\n│  ├─STGCN\n│  ├─STSGCN\n│  └─TGCN\n│\n├─data\n│  ├─CA_District5\n│  ├─chengdu_didi\n│  ├─NYC_BIKE\n│  ├─PEMS03\n│  ├─PEMS04\n│  ├─PEMS07\n│  ├─PEMS07M\n│  ├─PEMS08\n│  ├─PDFormer\n│  ├─STFGNN\n│  └─STGODE\n│\n├─lib\n│  │  data_process.py\n│  │  logger.py\n│  │  metrics.py\n│  │  predifineGraph.py\n│  │  TrainInits.py\n│\n├─model\n│  │  FlashST.py\n│  │  PromptNet.py\n│  │  Run.py\n│  │  Trainer.py\n│  │\n│  ├─AGCRN\n│  ├─ASTGCN\n│  ├─DMSTGCN\n│  ├─GWN\n│  ├─MSDR\n│  ├─MTGNN\n│  ├─PDFormer\n│  ├─STFGNN\n│  ├─STGCN\n│  ├─STGODE\n│  ├─STSGCN\n│  ├─ST_WA\n│  └─TGCN\n│\n└─SAVE\n    └─pretrain\n        ├─GWN\n        │      GWN_P8437.pth\n        │\n        ├─MTGNN\n        │      MTGNN_P8437.pth\n        │\n        ├─PDFormer\n        │      PDFormer_P8437.pth\n        │\n        └─STGCN\n                STGCN_P8437.pth\n            \n```\n\n---------\n\n\u003cspan id='Environment'/\u003e\n\n### 2.Environment \u003ca href='#all_catelogue'\u003e[Back to Top]\u003c/a\u003e\nThe code can be run in the following environments, other version of required packages may also work.\n* python==3.9.12\n* numpy==1.23.1\n* pytorch==1.9.0\n* cudatoolkit==11.1.1  \n\nOr you can install the required environment, which can be done by running the following commands:\n```\n# cteate new environmrnt\nconda create -n FlashST python=3.9.12\n\n# activate environmrnt\nconda activate FlashST\n\n# Torch with CUDA 11.1\npip install torch==1.9.1+cu111 torchvision==0.10.1+cu111 torchaudio==0.9.1 -f https://download.pytorch.org/whl/torch_stable.html\n\n# Install required libraries\npip install -r requirements.txt\n```\n\n---------\n\n\u003cspan id='Run the codes'/\u003e\n\n### 3. Run the codes \u003ca href='#all_catelogue'\u003e[Back to Top]\u003c/a\u003e\n\n* First, download \"data\" folder from [hugging face (data.zip)](https://huggingface.co/datasets/bjdwh/FlashST-DATA/tree/main) or [wisemodel (data.zip)](https://wisemodel.cn/datasets/BJDWH/FlashST-data/file), put it in the FlashST-main directory, unzip and then enter \"model\" folder:\n```\ncd model\n```\n* To test different models in various modes, you can execute the Run.py code. There are some examples:\n```\n# Evaluate the performance of MTGNN enhanced by FlashST on the PEMS07M dataset\npython Run.py -dataset_test PEMS07M -mode eval -model MTGNN\n\n# Evaluate the performance of STGCN enhanced by FlashST on the CA_District5 dataset\npython Run.py -dataset_test CA_District5 -mode eval -model STGCN\n\n# Evaluate the original performance of STGCN on the chengdu_didi dataset\npython Run.py -dataset_test chengdu_didi -mode ori -model STGCN\n\n# Pretrain from scratch with MTGNN model, checkpoint will be saved in FlashST-main/SAVE/pretrain/MTGNN(model name)/xxx.pth\npython Run.py -mode pretrain -model MTGNN\n```\n\n* Parameter setting instructions. The parameter settings consist of two parts: the pre-training model and the baseline model. To avoid any confusion arising from potential overlapping parameter names, we employ a hyphen (-) to specify the parameters of FlashST and use a double hyphen (--) to specify the parameters of the baseline model. Here is an example:\n```\n# Set first_layer_embedding_size and out_layer_dim to 32 in STFGNN\npython Run.py -model STFGNN -mode ori -dataset_test PEMS08 --first_layer_embedding_size 32 --out_layer_dim 32\n```\n\n---------\n\n\n## Citation\n\nIf you find FlashST useful in your research or applications, please kindly cite:\n\n```\n@misc{li2024flashst,\n      title={FlashST: A Simple and Universal Prompt-Tuning Framework for Traffic Prediction}, \n      author={Zhonghang Li and Lianghao Xia and Yong Xu and Chao Huang},\n      year={2024},\n      eprint={2405.17898},\n      archivePrefix={arXiv},\n      primaryClass={cs.LG}\n}\n```\n---------\n\n\n## Acknowledgements\nWe developed our code framework drawing inspiration from [AGCRN](https://github.com/LeiBAI/AGCRN) and [GPT-ST](https://github.com/HKUDS/GPT-ST). Furthermore, the implementation of the baselines primarily relies on a combination of the code released by the original author and the code from [LibCity](https://github.com/LibCity/Bigscity-LibCity). We extend our heartfelt gratitude for their remarkable contribution.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhkuds%2Fflashst","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fhkuds%2Fflashst","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fhkuds%2Fflashst/lists"}