# davidham3/stsgcn architecture

> Architecture diagram of the GitHub repository davidham3/stsgcn, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/davidham3/stsgcn
- Repository: https://github.com/davidham3/stsgcn

## Overview

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`](https://github.com/davidham3/stsgcn/blob/master/main.py) (Training runner), [`utils.py`](https://github.com/davidham3/stsgcn/blob/master/utils.py) (Model assembly), [`models/stsgcn.py`](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py) (Position embeddings), [`load_params.py`](https://github.com/davidham3/stsgcn/blob/master/load_params.py) (Checkpoint loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Training workflow"]
  node_runner["Training runner<br/>[main.py]"]
  node_config["Dataset configs"]
  node_mxnet["MXNet runtime"]
  node_training["Model training<br/>[main.py]"]
end

subgraph group_data_graph["Data and graph"]
  node_data_loader["Data preparation<br/>[utils.py]"]
  node_signal_data["Graph signal data"]
  node_graph_builder["Graph construction<br/>[utils.py]"]
  node_graph_files["Adjacency inputs"]
end

subgraph group_model["Forecast model"]
  node_model_builder["Model assembly<br/>[utils.py]"]
  node_embedding["Position embeddings<br/>[stsgcn.py]"]
  node_stsgcl["STSGC layers<br/>[stsgcn.py]"]
  node_stsgcm["Graph convolution<br/>[stsgcn.py]"]
  node_prediction["Forecast output<br/>[stsgcn.py]"]
  node_loss["Weighted loss<br/>[stsgcn.py]"]
end

subgraph group_results["Evaluation and artifacts"]
  node_validation["Validation selection<br/>[main.py]"]
  node_test_metrics["Test metrics<br/>[main.py]"]
  node_metric_helpers["Masked metrics<br/>[utils.py]"]
  node_checkpoint["Model checkpoint"]
  node_parameter_loader["Checkpoint loader<br/>[load_params.py]"]
end

node_user(("Operator"))

node_user -->|"starts"| node_runner
node_runner -->|"reads"| node_config
node_runner -->|"constructs model"| node_model_builder
node_runner -->|"loads splits"| node_data_loader
node_data_loader -->|"reads"| node_signal_data
node_model_builder -->|"builds adjacency"| node_graph_builder
node_graph_builder -->|"reads"| node_graph_files
node_model_builder -->|"assembles"| node_embedding
node_embedding -->|"embeds inputs"| node_stsgcl
node_stsgcl -->|"applies layers"| node_stsgcm
node_stsgcm -->|"produces features"| node_prediction
node_prediction -->|"compares labels"| node_loss
node_runner -->|"binds model"| node_mxnet
node_mxnet -->|"runs batches"| node_training
node_training -->|"evaluates"| node_validation
node_validation -->|"predicts"| node_mxnet
node_validation -->|"selects best"| node_test_metrics
node_test_metrics -->|"predicts"| node_mxnet
node_test_metrics -->|"computes errors"| node_metric_helpers
node_runner -->|"saves optionally"| node_checkpoint
node_parameter_loader -->|"loads"| node_checkpoint

click node_runner "https://github.com/davidham3/stsgcn/blob/master/main.py"
click node_config "https://github.com/davidham3/stsgcn/tree/master/config"
click node_data_loader "https://github.com/davidham3/stsgcn/blob/master/utils.py"
click node_graph_builder "https://github.com/davidham3/stsgcn/blob/master/utils.py"
click node_model_builder "https://github.com/davidham3/stsgcn/blob/master/utils.py"
click node_embedding "https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py"
click node_stsgcl "https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py"
click node_stsgcm "https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py"
click node_prediction "https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py"
click node_loss "https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py"
click node_training "https://github.com/davidham3/stsgcn/blob/master/main.py"
click node_validation "https://github.com/davidham3/stsgcn/blob/master/main.py"
click node_test_metrics "https://github.com/davidham3/stsgcn/blob/master/main.py"
click node_metric_helpers "https://github.com/davidham3/stsgcn/blob/master/utils.py"
click node_parameter_loader "https://github.com/davidham3/stsgcn/blob/master/load_params.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_runner,node_config,node_mxnet,node_training toneBlue
class node_data_loader,node_signal_data,node_graph_builder,node_graph_files toneAmber
class node_model_builder,node_embedding,node_stsgcl,node_stsgcm,node_prediction,node_loss toneMint
class node_validation,node_test_metrics,node_metric_helpers,node_checkpoint,node_parameter_loader toneRose
class node_user toneIndigo
```

## Components

### Training workflow

- [Training runner](https://github.com/davidham3/stsgcn/blob/master/main.py): `main.py`
- [Dataset configs](https://github.com/davidham3/stsgcn/tree/master/config): `config`
- **MXNet runtime**
- [Model training](https://github.com/davidham3/stsgcn/blob/master/main.py): `main.py`

### Data and graph

- [Data preparation](https://github.com/davidham3/stsgcn/blob/master/utils.py): `utils.py`
- **Graph signal data**
- [Graph construction](https://github.com/davidham3/stsgcn/blob/master/utils.py): `utils.py`
- **Adjacency inputs**

### Forecast model

- [Model assembly](https://github.com/davidham3/stsgcn/blob/master/utils.py): `utils.py`
- [Position embeddings](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py): `models/stsgcn.py`
- [STSGC layers](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py): `models/stsgcn.py`
- [Graph convolution](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py): `models/stsgcn.py`
- [Forecast output](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py): `models/stsgcn.py`
- [Weighted loss](https://github.com/davidham3/stsgcn/blob/master/models/stsgcn.py): `models/stsgcn.py`

### Evaluation and artifacts

- [Validation selection](https://github.com/davidham3/stsgcn/blob/master/main.py): `main.py`
- [Test metrics](https://github.com/davidham3/stsgcn/blob/master/main.py): `main.py`
- [Masked metrics](https://github.com/davidham3/stsgcn/blob/master/utils.py): `utils.py`
- **Model checkpoint**
- [Checkpoint loader](https://github.com/davidham3/stsgcn/blob/master/load_params.py): `load_params.py`

### Other components

- **Operator**

## Connections

- 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