STGFORMER_PEMS-BAY

STGformer Model - PEMS-BAY

Spatio-Temporal Graph Transformer (STGformer) trained on PEMS-BAY dataset for traffic speed forecasting.

Model Description

This model uses a transformer-based graph neural network architecture that combines:

• Self-attention mechanisms for capturing temporal dependencies
• Spatial graph convolution for modeling spatial relationships
• Adaptive embeddings for learning node-specific patterns
• Time-of-day embeddings for capturing daily patterns

Evaluation Metrics

• Test MAE (15 min): 1.0933
• Test MAPE (15 min): 0.0221
• Test RMSE (15 min): 2.2381

Dataset

PEMS-BAY: Traffic speed data from highway sensors.

Usage

from utils.stgformer import load_from_hub

# Load model from Hub
model, scaler = load_from_hub("PEMS-BAY")

# Get predictions
import numpy as np
x = np.random.randn(10, 12, 207, 2)  # (batch, seq_len, nodes, [value, tod])
predictions = model.predict(x)

Training

Model was trained using the STGformer implementation with configuration:

• Input features: 2 [speed, time-of-day]
• Time-of-day embedding dimension: 24
• Day-of-week embedding dimension: 0 (disabled)
• Adaptive embedding dimension: 80
• Number of attention heads: 4
• Number of layers: 3

Citation

If you use this model, please cite the STGformer paper:

@article{stgformer,
  title={STGformer: Spatio-Temporal Graph Transformer for Traffic Forecasting},
  author={Author names},
  journal={Conference/Journal},
  year={Year}
}

License

This model checkpoint is released under the same license as the training code.