Upload STGFORMER_PRETRAIN model trained on METR-LA

README.md

@@ -0,0 +1,57 @@
+---
+tags:
+- traffic-forecasting
+- time-series
+- graph-neural-network
+- stgformer_pretrain
+datasets:
+- metr-la
+---
+
+# Spatial-Temporal Graph Transformer (Pretrain) - METR-LA
+
+Spatial-Temporal Graph Transformer (Pretrain) (STGFORMER_PRETRAIN) trained on METR-LA dataset for traffic speed forecasting.
+
+## Model Description
+
+STGFormer with masked node pretraining (curriculum: per-timestep -> per-node)
+
+
+
+## Dataset
+
+**METR-LA**: Traffic speed data from highway sensors.
+
+## Usage
+
+```python
+from utils.stgformer import load_from_hub
+
+# Load model from Hub
+model, scaler = load_from_hub("METR-LA", hf_repo_prefix="STGFORMER_PRETRAIN")
+
+# Get predictions
+from utils.stgformer import get_predictions
+predictions = get_predictions(model, scaler, test_dataset)
+```
+
+## Training
+
+Model was trained using the STGFORMER_PRETRAIN implementation with default hyperparameters.
+
+## Citation
+
+If you use this model, please cite the original STGFORMER_PRETRAIN paper:
+
+```bibtex
+@inproceedings{lan2022stgformer,
+  title={STGformer: Spatial-Temporal Graph Transformer for Traffic Forecasting},
+  author={Lan, Shengnan and Ma, Yong and Huang, Weijia and Wang, Wanwei and Yang, Hui and Li, Peng},
+  booktitle={IEEE Transactions on Neural Networks and Learning Systems},
+  year={2022}
+}
+```
+
+## License
+
+This model checkpoint is released under the same license as the training code.

config.json

@@ -0,0 +1,32 @@
+{
+  "num_nodes": 207,
+  "in_steps": 12,
+  "out_steps": 12,
+  "input_dim": 1,
+  "output_dim": 1,
+  "steps_per_day": 288,
+  "input_embedding_dim": 24,
+  "tod_embedding_dim": 24,
+  "dow_embedding_dim": 0,
+  "spatial_embedding_dim": 0,
+  "adaptive_embedding_dim": 80,
+  "num_heads": 4,
+  "num_layers": 3,
+  "dropout_a": 0.3,
+  "use_mixed_proj": true,
+  "model_dim": 128,
+  "dataset": "METR-LA",
+  "graph_mode": "learned",
+  "lambda_hybrid": 0.5,
+  "sparsity_k": null,
+  "propagation_mode": "power",
+  "temporal_mode": "transformer",
+  "mamba_d_state": 16,
+  "mamba_d_conv": 4,
+  "mamba_expand": 2,
+  "tcn_num_layers": 3,
+  "tcn_kernel_size": 3,
+  "tcn_dilation_base": 2,
+  "tcn_dropout": 0.1,
+  "use_zero_init": true
+}

hub_metadata.json

@@ -0,0 +1,7 @@
+{
+  "dataset": "METR-LA",
+  "metrics": {},
+  "framework": "PyTorch",
+  "hf_repo_prefix": "STGFORMER_PRETRAIN",
+  "implementation": "internal"
+}

metadata.json

@@ -0,0 +1,7 @@
+{
+  "dataset": "METR-LA",
+  "upload_date": "2025-12-07T20:12:44.070342",
+  "metrics": {},
+  "framework": "PyTorch",
+  "model_type": "STGFORMER_PRETRAIN"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bc3b7dbba0b56513a78ded78ffbedaa87b9810f06136daaf40974c12c7d2d0f
+size 7231112

pretrained_checkpoint/hub_metadata.json

@@ -0,0 +1,6 @@
+{
+  "dataset": "METR-LA",
+  "checkpoint_type": "pretrained",
+  "framework": "PyTorch",
+  "hf_repo_prefix": "STGFORMER_PRETRAINED"
+}

pretrained_checkpoint/imputation_head.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dad72f93fc148ecdd1897ad4b92e8cc347adc38bbdc2ebcae7fa1b6d1ddd3890
+size 652

pretrained_checkpoint/metadata.json

@@ -0,0 +1,7 @@
+{
+  "dataset": "METR-LA",
+  "upload_date": "2025-12-07T19:53:28.410992",
+  "metrics": {},
+  "framework": "PyTorch",
+  "model_type": "PRETRAINED"
+}

pretrained_checkpoint/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:545fe5933952da9b4be422274667c7e590360565596b6ed5cf1bf328231411b8
+size 7231112

pretrained_checkpoint/pretrain_config.json

@@ -0,0 +1,14 @@
+{
+  "dataset_name": "METR-LA",
+  "pretrain_config": {
+    "stage1_epochs": 5,
+    "stage1_mask_ratio": 0.15,
+    "stage2_epochs": 5,
+    "stage2_mask_ratio": 0.1,
+    "pretrain_window": 24,
+    "learning_rate": 0.001,
+    "save_to": "STGFORMER_PRETRAINED"
+  },
+  "model_dim": 128,
+  "num_nodes": 207
+}