# shawnzhang0828/polynormer-reimplementation architecture

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

- Interactive diagram: https://gitdiagram.com/shawnzhang0828/polynormer-reimplementation
- Repository: https://github.com/shawnzhang0828/polynormer-reimplementation

## Overview

Polynormer-Reimplementation trains and evaluates a graph node-classification model. The principal workflow loads a supported dataset and its splits, trains local layers during warm-up, then trains with global attention enabled, evaluates the selected metric, and saves the best checkpoints. The model combines graph-local attention with linear global attention. The README also identifies a notebook experiment runner; its implementation was not sampled, so its internal wiring is not shown.

Main entry points: [`train.py`](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py) (CLI trainer), [`utils/data_loaders.py`](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/data_loaders.py) (Dataset loader), [`models/polynormer.py`](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/polynormer.py) (Polynormer), [`utils/io.py`](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/io.py) (Checkpoint I/O)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_entry["Experiment entry"]
  node_cli["CLI trainer<br/>[train.py]"]
  node_notebook["Notebook runner"]
  node_config["Default config<br/>[config.py]"]
  node_seed["Seed control<br/>[seed.py]"]
end

subgraph group_data["Graph data"]
  node_loader["Dataset loader<br/>[data_loaders.py]"]
  node_dataset_libs["Dataset providers"]
  node_graph[("Graph data")]
end

subgraph group_model["Polynormer model"]
  node_model["Polynormer<br/>[polynormer.py]"]
  node_local["Local attention<br/>[local_attention.py]"]
  node_global["Global attention"]
end

subgraph group_results["Training results"]
  node_warmup["Local warm-up<br/>[train.py]"]
  node_joint["Local-global training<br/>[train.py]"]
  node_metrics["Metric calculation<br/>[metrics.py]"]
  node_checkpoint["Checkpoint I/O<br/>[io.py]"]
  node_checkpoint_store[("Checkpoint files")]
end

node_user(("Experimenter"))

node_user -->|"starts training"| node_cli
node_cli -->|"loads defaults"| node_config
node_cli -->|"sets seeds"| node_seed
node_cli -->|"loads dataset"| node_loader
node_loader -->|"requests dataset"| node_dataset_libs
node_dataset_libs -->|"returns graph"| node_loader
node_loader -->|"prepares splits"| node_graph
node_cli -->|"reads features"| node_graph
node_cli -->|"runs first"| node_warmup
node_warmup -->|"trains frozen-global"| node_model
node_cli -->|"saves best warm-up"| node_checkpoint
node_cli -->|"loads best warm-up"| node_checkpoint
node_cli -->|"runs second"| node_joint
node_joint -->|"trains full model"| node_model
node_model -->|"calls local layers"| node_local
node_model -->|"calls global layers"| node_global
node_cli -->|"computes metrics"| node_metrics
node_cli -->|"saves best model"| node_checkpoint
node_checkpoint -->|"writes and reads"| node_checkpoint_store
node_user -->|"runs experiments"| node_notebook

click node_cli "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py"
click node_notebook "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/run_experiments.ipynb"
click node_config "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/config.py"
click node_seed "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/seed.py"
click node_loader "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/data_loaders.py"
click node_model "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/polynormer.py"
click node_local "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/local_attention.py"
click node_global "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/global_attention.py"
click node_warmup "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py"
click node_joint "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py"
click node_metrics "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/metrics.py"
click node_checkpoint "https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/io.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_cli,node_notebook,node_config,node_seed toneBlue
class node_loader,node_dataset_libs,node_graph toneAmber
class node_model,node_local,node_global toneMint
class node_warmup,node_joint,node_metrics,node_checkpoint,node_checkpoint_store toneRose
class node_user toneIndigo
```

## Components

### Experiment entry

- [CLI trainer](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py): `train.py`
- [Notebook runner](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/run_experiments.ipynb): `run_experiments.ipynb`
- [Default config](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/config.py): `config.py`
- [Seed control](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/seed.py): `utils/seed.py`

### Graph data

- [Dataset loader](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/data_loaders.py): `utils/data_loaders.py`
- **Dataset providers**
- **Graph data**

### Polynormer model

- [Polynormer](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/polynormer.py): `models/polynormer.py`
- [Local attention](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/local_attention.py): `models/local_attention.py`
- [Global attention](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/models/global_attention.py): `models/global_attention.py`

### Training results

- [Local warm-up](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py): `train.py`
- [Local-global training](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/train.py): `train.py`
- [Metric calculation](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/metrics.py): `utils/metrics.py`
- [Checkpoint I/O](https://github.com/shawnzhang0828/polynormer-reimplementation/blob/main/utils/io.py): `utils/io.py`
- **Checkpoint files**

### Other components

- **Experimenter**

## Connections

- Experimenter → CLI trainer: starts training
- CLI trainer → Default config: loads defaults
- CLI trainer → Seed control: sets seeds
- CLI trainer → Dataset loader: loads dataset
- Dataset loader → Dataset providers: requests dataset
- Dataset providers → Dataset loader: returns graph
- Dataset loader → Graph data: prepares splits
- CLI trainer → Graph data: reads features
- CLI trainer → Local warm-up: runs first
- Local warm-up → Polynormer: trains frozen-global
- CLI trainer → Checkpoint I/O: saves best warm-up
- CLI trainer → Checkpoint I/O: loads best warm-up
- CLI trainer → Local-global training: runs second
- Local-global training → Polynormer: trains full model
- Polynormer → Local attention: calls local layers
- Polynormer → Global attention: calls global layers
- CLI trainer → Metric calculation: computes metrics
- CLI trainer → Checkpoint I/O: saves best model
- Checkpoint I/O → Checkpoint files: writes and reads
- Experimenter → Notebook runner: runs experiments