How thuml/corrformer works
1. What this repo is
- This is a PyTorch implementation of Corrformer, a unified deep model for large-scale weather forecasting.
- The code is organized around a single training/inference pipeline that can run on different station-scale datasets, especially the Global Temp and Global Wind benchmarks.
2. Top-level execution path
- run.py is the main entry point for launching experiments.
- The actual experiment logic is delegated into the exp/ package, which handles the lifecycle of train, validation, test, and demo/evaluation runs.
3. Experiment orchestration boundary
- exp/exp_basic.py defines shared experiment behavior and common setup.
- exp/exp_main.py contains the concrete experiment implementation used for Corrformer.
- This layer is the control plane: it wires together data loading, model creation, optimization, evaluation, and checkpoint usage.
4. Data subsystem
- data_provider/ encapsulates dataset loading and batching.
- data_factory.py selects the appropriate dataset/dataloader setup for each experiment.
- data_loader.py likely defines dataset classes and split handling for train/val/test.
- The repo expects external data under ./dataset, so raw data is not stored here.
5. Model subsystem
- models/Corrformer.py is the top-level model definition.
- It assembles the architecture from reusable components in layers/.
- The model is designed around spatiotemporal forecasting across many stations, not a generic sequence classifier.
6. Core architectural building blocks
- layers/Embed.py handles input embedding and feature projection.
- layers/Corrformer_EncDec.py defines the encoder-decoder structure used in the paper.
- layers/Causal_Conv.py supports causal convolution for cross-correlation modeling.
- layers/Multi_Correlation.py contains the Multi-Correlation mechanism that unifies temporal and spatial dependencies.
- layers/ is therefore the main research/novelty layer of the repo.
7. Utility layer
- utils/ holds shared helpers that are not model-specific.
- metrics.py provides evaluation metrics.
- timefeatures.py builds time-based covariates.
- masking.py likely supports attention or causal masking.
- tools.py contains general training utilities such as logging, seeding, or checkpoint helpers.
8. Training and evaluation outputs
- The repo writes learned weights and pretrained artifacts to ./checkpoints.
- The README indicates both full training runs and demo runs can load pretrained checkpoints.
- This creates a clear boundary between offline model training and later evaluation/demo inference.
9. Scripted experiment entry points
- scripts/Global_Temp/ and scripts/Global_Wind/ contain the production training scripts for the two main benchmarks.
- scripts/Demo/ contains shorter runs for testing pretrained models.
- These scripts are the operational interface for reproducing paper results.
10. End-to-end data flow
- External dataset files are placed in ./dataset.
- data_provider/ loads and prepares them into model-ready batches.
- exp/ drives the run loop and feeds batches into models/Corrformer.py.
- The model uses the specialized layers/ modules to compute forecasts.
- utils/metrics.py evaluates predictions, and results/checkpoints are persisted for later reuse.
11. Technology and runtime assumptions
- The project is built for Python + PyTorch and was reproduced on a single GPU environment.
- The README notes the intended setup is Linux with a CUDA-capable device.
- This matters architecturally because the code and scripts are geared toward batch GPU training rather than distributed serving.
12. Architectural shape to diagram
- Draw this as a pipeline with four main zones: data ingestion, experiment orchestration, Corrformer model internals, and evaluation/checkpointing.
- Treat layers/ as the internal model backbone, exp/ as the control layer, and data_provider/ + utils/ as supporting infrastructure around them.
Main entry points: exp_main.py Main exp · run.py run.py · Corrformer.py Corrformer · Corrformer_EncDec.py Enc-Dec
Components 18
Entry & scripts
Command-line entry points and benchmark launch scripts.
- run.py · Main launcher for training, validation, testing, and demo runs.
run.py - Temp script · Training script for the Global Temp benchmark.
scripts/Global_Temp/Corrformer.sh - Wind script · Training script for the Global Wind benchmark.
scripts/Global_Wind/Corrformer.sh - Temp demo · Short demo run for pretrained evaluation on Global Temp.
scripts/Demo/Global_Temp_demo.sh - Wind demo · Short demo run for pretrained evaluation on Global Wind.
scripts/Demo/Global_Wind_demo.sh
Data ingestion
Dataset selection, loading, and batch preparation.
- Data factory · Chooses the dataset and dataloader configuration for each experiment.
data_provider/data_factory.py - Data loader · Defines dataset splits and batch construction for train, validation, and test.
data_provider/data_loader.py
Experiment control
Shared experiment lifecycle and Corrformer-specific run logic.
Corrformer model
Model assembly and core architectural blocks.
- Corrformer · Top-level model that assembles the full forecasting architecture.
models/Corrformer.py - Embedding · Projects raw station and time features into model space.
layers/Embed.py - Enc-Dec · Encoder-decoder structure used by Corrformer.
layers/Corrformer_EncDec.py - Causal conv · Causal convolution block for cross-correlation modeling.
layers/Causal_Conv.py - Multi-corr · Core mechanism for unifying temporal and spatial dependencies.
layers/Multi_Correlation.py
Evaluation & utilities
Metrics, time features, masking, and training helpers.
- Metrics · Forecast error metrics and result scoring.
utils/metrics.py - Time feats · Builds time-based covariates for the input pipeline.
utils/timefeatures.py - Masking · Mask construction for causal or attention-based computation.
utils/masking.py - Tools · General helpers such as logging, seeding, and checkpoint support.
utils/tools.py
Connections 20
- run.py → Base exp: launches · CLI entry point instantiates the experiment layer.
- run.py → Temp script: used by · Benchmark script drives the main launcher for Global Temp.
- run.py → Wind script: used by · Benchmark script drives the main launcher for Global Wind.
- run.py → Temp demo: used by · Demo script drives pretrained evaluation for Global Temp.
- run.py → Wind demo: used by · Demo script drives pretrained evaluation for Global Wind.
- Base exp → Main exp: extends · Concrete Corrformer experiment inherits shared experiment behavior.
- Main exp → Data factory: loads data · Experiment selects the appropriate dataset pipeline.
- Data factory → Data loader: builds · Factory instantiates the dataset and batching logic.
- Main exp → Corrformer: builds model · Experiment creates the Corrformer network for training or inference.
- Main exp → Metrics: evaluates · Experiment computes forecast quality during validation and test.
- Main exp → Tools: uses · Experiment relies on shared helpers for runtime support.
- Main exp → Time feats: uses · Experiment consumes time covariates from the utility layer.
- Corrformer → Embedding: assembles · Top-level model starts with input embedding.
- Corrformer → Enc-Dec: assembles · Model backbone uses the encoder-decoder structure.
- Enc-Dec → Causal conv: uses · Backbone incorporates causal convolution for temporal correlation.
- Enc-Dec → Multi-corr: uses · Backbone centers on the multi-correlation mechanism.
- Multi-corr → Masking: uses · Correlation module depends on masking for constrained computation.
- Data loader → Time feats: augments · Dataset processing enriches inputs with temporal features.
- Main exp → Corrformer: trains · Control plane feeds batches into the model for optimization and inference.
- Metrics → Tools: reports via · Metric results are handled alongside shared training utilities.