Upload STGFORMER_PRETRAIN model trained on METR-LA

metadata.json

@@ -1,6 +1,6 @@
 {
   "dataset": "METR-LA",
-  "upload_date": "2025-12-07T20:12:44.070342",
+  "upload_date": "2025-12-08T00:58:01.654463",
   "metrics": {},
   "framework": "PyTorch",
   "model_type": "STGFORMER_PRETRAIN"

model.safetensors

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

pretrained_checkpoint/README.md

@@ -0,0 +1,57 @@
+---
+tags:
+- traffic-forecasting
+- time-series
+- graph-neural-network
+- stgformer_pretrained
+datasets:
+- metr-la
+---
+
+# Spatial-Temporal Graph Transformer (Pretrained) - METR-LA
+
+Spatial-Temporal Graph Transformer (Pretrained) (STGFORMER_PRETRAINED) trained on METR-LA dataset for traffic speed forecasting.
+
+## Model Description
+
+STGFormer pretrained checkpoint for METR-LA. This checkpoint contains pretrained model weights and imputation head from masked node pretraining. Use with load_from config option.
+
+
+
+## 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_PRETRAINED")
+
+# Get predictions
+from utils.stgformer import get_predictions
+predictions = get_predictions(model, scaler, test_dataset)
+```
+
+## Training
+
+Model was trained using the STGFORMER_PRETRAINED implementation with default hyperparameters.
+
+## Citation
+
+If you use this model, please cite the original STGFORMER_PRETRAINED 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.

pretrained_checkpoint/imputation_head.safetensors

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

pretrained_checkpoint/metadata.json

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

pretrained_checkpoint/model.safetensors

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

pretrained_checkpoint/pretrain_config.json

@@ -5,8 +5,10 @@
     "stage1_mask_ratio": 0.15,
     "stage2_epochs": 5,
     "stage2_mask_ratio": 0.1,
-    "pretrain_window": 24,
+    "pretrain_window": 12,
     "learning_rate": 0.001,
+    "pretrain_batch_size": 100,
+    "imputation_iterations": 3,
     "save_to": "STGFORMER_PRETRAINED"
   },
   "model_dim": 128,