How davidham3/stsgcn works
STSGCN is a traffic-network forecasting model and training program. The principal workflow is to load a dataset-specific configuration and graph signal data, construct a spatial-temporal graph model, train it, select the best validation checkpoint, and report test forecasts using masked error metrics. The model supports individual or shared graph-convolution modules, temporal and spatial embeddings, and GLU or ReLU activations. The README documents PEMS configurations and OpenPAI job descriptions; sampled source shows training and data/model construction, but not the internals of the unsampled model sections or cluster dispatch.
Main entry points: main.py Training runner · utils.py Model assembly · stsgcn.py Position embeddings · load_params.py Checkpoint loader
Components 20
Training workflow
- Training runner
main.py - Dataset configs
config - MXNet runtime
- Model training
main.py
Data and graph
- Data preparation
utils.py - Graph signal data
- Graph construction
utils.py - Adjacency inputs
Forecast model
- Model assembly
utils.py - Position embeddings
models/stsgcn.py - STSGC layers
models/stsgcn.py - Graph convolution
models/stsgcn.py - Forecast output
models/stsgcn.py - Weighted loss
models/stsgcn.py
Evaluation and artifacts
- Validation selection
main.py - Test metrics
main.py - Masked metrics
utils.py - Model checkpoint
- Checkpoint loader
load_params.py
Other components
- Operator
Connections 21
- Operator → Training runner: starts
- Training runner → Dataset configs: reads
- Training runner → Model assembly: constructs model
- Training runner → Data preparation: loads splits
- Data preparation → Graph signal data: reads
- Model assembly → Graph construction: builds adjacency
- Graph construction → Adjacency inputs: reads
- Model assembly → Position embeddings: assembles
- Position embeddings → STSGC layers: embeds inputs
- STSGC layers → Graph convolution: applies layers
- Graph convolution → Forecast output: produces features
- Forecast output → Weighted loss: compares labels
- Training runner → MXNet runtime: binds model
- MXNet runtime → Model training: runs batches
- Model training → Validation selection: evaluates
- Validation selection → MXNet runtime: predicts
- Validation selection → Test metrics: selects best
- Test metrics → MXNet runtime: predicts
- Test metrics → Masked metrics: computes errors
- Training runner → Model checkpoint: saves optionally
- Checkpoint loader → Model checkpoint: loads