StructCast-Model

StructCast-Model is a configuration-driven toolkit that generates PyTorch, Flax (JAX), and Keras models — plus PyTorch training workflows — from YAML templates. Built on top of StructCast, it lets you describe model architecture, optimizer logic, dataset configuration, and training orchestration declaratively — then generates runnable Python code from those descriptions.

Model code generation is available for all three frameworks. Training workflow generation and the full training CLI (scm torch train) are currently PyTorch-only; Flax and Keras training support is planned (see Roadmap).

Table of Contents

• StructCast-Model
  • Table of Contents
  • What This Project Does
  • Installation
  • Project Structure
  • Core Workflow
  • StructCast Pattern Basics
  • Quick Start
  • Command Guide
    • 1. Format Templates
    • 2. Generate a Model Class
    • 3. Generate a Learner Class
    • 4. Inspect FLOPs and Parameters
    • 5. Measure Inference Time
    • 6. Train a Generated Model
      • Distributed Training with torchrun
        • How It Works
        • Single-Node Multi-GPU
        • Multi-Node Training
        • Dataset Configuration
        • Distributed Training Notes
  • Training Loop Anatomy
  • Configuration Examples
    • PyTorch
    • Flax
    • Keras
  • Development
  • Migration Notes
    • Upgrading from v1.x
  • Roadmap

What This Project Does

• Generate model code — Produce PyTorch nn.Module, Flax nnx.Module, and Keras Layer classes from YAML layer templates.
• Generate training code — Produce learner classes — the object owning the models, the optimizers, and the training and inference steps — from YAML templates (PyTorch only).
• Format reusable templates — Render parameterized YAML templates into concrete runtime configurations.
• Inspect model complexity — Compute FLOPs and parameter counts with ptflops and calflops (PyTorch only).
• Measure inference time — Benchmark average forward-pass latency of generated models across all three frameworks via scm [torch/flax/keras] time.
• Train end-to-end — Run PyTorch training with Automatic Mixed Precision (AMP), timm datasets, optional torch.compile, and MLflow or Weights & Biases experiment logging.
• Train programmatically — Use the same trainer directly from Python, without any YAML. See examples/ for a runnable tutorial.

Installation

StructCast-Model is installed with uv and exposes the scm CLI entry point.

uv sync --extra torch-cu130 --extra mlflow --extra flops

Each extra installs a group of optional dependencies. Pick the extras that match your target framework and accelerator. Keras is multi-backend and runs on top of JAX, PyTorch, or TensorFlow.

Omit any extra you do not need. For example, uv sync --extra torch-cu130 is sufficient if you only want to generate and train PyTorch models without FLOPs analysis or MLflow logging. To work with all three model frameworks on CPU:

uv sync --extra all-cpu

Project Structure

structcast-model/
├── cfg/
│   ├── torch/
│   │   ├── learners/      # learner, optimizer, loss, and metric templates
│   │   ├── models/        # model architecture templates
│   │   └── others/        # dataset, compile options, and other templates
│   ├── flax/
│   │   └── models/        # Flax model architecture templates
│   └── keras/
│       └── models/        # Keras model architecture templates
├── examples/
│   └── torch/         # runnable training tutorial and optimizer compositions
├── src/structcast_model/
│   ├── builders/      # generic and framework-specific code generators
│   ├── commands/      # Typer CLI entry points
│   ├── torch/         # trainer, layers, optimizer helpers
│   ├── flax/          # Flax layers and inference utilities
│   ├── keras/         # Keras layers and inference utilities
│   ├── utils/         # shared helpers
│   └── base_trainer.py
├── tests/             # CLI, builder, trainer, and layer tests
└── README.md

The main package areas are:

• builders/ — Converts validated YAML templates into intermediate representations, then renders Python source code for PyTorch, Flax, and Keras.
• commands/ — Exposes the scm CLI (built with Typer) with torch, flax, and keras sub-commands.
• torch/ — Runtime utilities used by the CLI and available for direct Python usage — training steps, trackers, timm wrappers, optimizer helpers.
• flax/ — Flax-specific layers (e.g. GlobalResponseNorm) and JAX inference helpers.
• keras/ — Keras-specific layers (e.g. GlobalResponseNormalization) and backend-agnostic inference helpers.
• cfg/torch/ — Declarative source of truth: YAML templates for PyTorch models, learners, datasets, and runtime presets.
• examples/torch/ — Runnable example code: a programmatic training tutorial, and optimizer + scheduler compositions that templates reference by file path.
• cfg/flax/ — YAML templates for Flax model architectures.
• cfg/keras/ — YAML templates for Keras model architectures.

Core Workflow

The repository follows a repeatable workflow:

1. Write or reuse YAML templates under cfg/[torch/flax/keras]/.
2. Render templates with scm format and -p/--parameter overrides to produce concrete configuration files.
3. Generate Python source files for the model (and, for PyTorch, the learner) using scm [torch/flax/keras] create.
4. Instantiate those generated modules at runtime through StructCast object patterns (see StructCast Pattern Basics).
5. Benchmark inference latency with scm [torch/flax/keras] time.
6. (PyTorch only) Train through scm torch train, which wires together datasets, models, the learner, AMP, and the experiment logger.

YAML templates  --->  scm format / scm [torch/flax/keras] create  --->  Generated .py files
                                                                              |
StructCast patterns  <--------------------------------------------------------+
       |
       v
scm [torch/flax/keras] time  --->  Inference benchmarks
scm torch train              --->  MLflow / wandb logs + model checkpoints

StructCast Pattern Basics

This repository relies heavily on StructCast object patterns to bridge generated source files and runtime commands. The minimum syntax you need to read the CLI examples is:

Example:

[_obj_, {_addr_: model.Model, _file_: model.py}, _call_]

This pattern does the following:

1. Import Model from the local file model.py.
2. Call Model() with no arguments and return the instance.

This pattern is the bridge between generated source files and runtime commands like ptflops, calflops, and train. For full documentation on StructCast patterns, see the StructCast README.

Quick Start

The following commands generate a ConvNeXtV2 model along with its learner and dataset configurations, then launch a training run on CIFAR-100.

# 1. Install
uv sync --extra torch-cu130 --extra mlflow --extra flops

# 2. Generate the model and the learner classes
scm torch create model cfg/torch/models/ConvNeXtV2.yaml -p 'DEFAULT: {backbone: femto}' -o model.py
scm torch create learner cfg/torch/learners/ConvNeXtV2.yaml -p 'DEFAULT: {epochs: 5}' -o learner.py

# 3. Render dataset configurations from templates
scm format cfg/torch/others/default_timm.yaml \
    -o dataset_train.yaml \
    -p 'DEFAULT: {training: true, epochs: 5, batch_size: 32, dataset: torch/cifar100, num_classes: 100, label_smoothing: 0.1, input_size: [3, 224, 224], image_dtype: bfloat16, download: true}'

scm format cfg/torch/others/default_timm.yaml \
    -o dataset_valid.yaml \
    -p 'DEFAULT: {training: false, epochs: 5, batch_size: 32, dataset: torch/cifar100, num_classes: 100, input_size: [3, 224, 224], image_dtype: bfloat16, download: true}'

# 4. Train
scm torch train \
    'model: [_obj_, {_addr_: model.Model, _file_: model.py}, _call_]' \
    -s 'image: [3, 224, 224]' \
    -d cuda \
    -L '[_obj_, {_addr_: learner.Learner, _file_: learner.py}]' \
    -c cfg/torch/others/compile_default.yaml \
    -e 5 \
    --training-dataset dataset_train.yaml \
    -V dataset_valid.yaml \
    -f 1 \
    -LC ce_loss \
    -LC val_ce_loss \
    -HC acc1 \
    -HC val_acc1 \
    -HC acc5 \
    -HC val_acc5 \
    -SC val_acc1 \
    --matmul-precision high \
    -E Test

Each step is explained in detail under Command Guide. To see the same training run built in plain Python instead of YAML, start from examples/ and run uv run python examples/torch/simple_training.py.

Command Guide

1. Format Templates

Use scm format to render a parameterized YAML template (such as cfg/torch/others/default_timm.yaml) into a concrete configuration file.

scm format cfg/torch/others/default_timm.yaml \
    -o dataset_train.yaml \
    -p 'DEFAULT: {training: true, epochs: 5, batch_size: 32, dataset: torch/cifar100, num_classes: 100, label_smoothing: 0.1, input_size: [3, 224, 224], image_dtype: bfloat16, download: true}'

scm format cfg/torch/others/default_timm.yaml \
    -o dataset_valid.yaml \
    -p 'DEFAULT: {training: false, epochs: 5, batch_size: 32, dataset: torch/cifar100, num_classes: 100, input_size: [3, 224, 224], image_dtype: bfloat16, download: true}'

What this does:

1. Loads the YAML template.
2. Merges any repeated -p/--parameter groups into a single parameter set.
3. Renders Jinja-based sections within the template.
4. Writes the resolved YAML to -o/--output (or prints to stdout if -o is omitted).

2. Generate a Model Class

Each framework has its own create model command that reads a YAML layer template and generates a framework-native module. The examples below use PyTorch; Flax and Keras share the same interface with minor differences noted afterward.

scm torch create model cfg/torch/models/ConvNeXtV2.yaml
scm torch create model cfg/torch/models/ConvNeXtV2.yaml -p 'DEFAULT: {backbone: femto}'
scm torch create model cfg/torch/models/ConvNeXtV2.yaml -p 'DEFAULT: {backbone: atto}' -o torch_model.py

Common options — All three framework commands share the same options:

• -p/--parameter: override template parameters
• -c/--classname: set the generated class name, default Model
• --structured-output/--no-structured-output: force the root model's return type. scm torch defaults to the template's STRUCTURED_OUTPUT (a plain tuple-like return unless the template sets it); scm flax and scm keras default to a structured output mapping
• -s/--sublayer: generate a named sublayer from the template instead of the root model
• -o/--output: output file path; if omitted, defaults to the snake-cased class name in the current directory (e.g., model.py for the default class name Model)

The ConvNeXtV2 template uses Jinja parameter groups to switch between backbone variants such as atto, femto, tiny, and base.

Flax and Keras — Replace scm torch with scm flax or scm keras. Templates live under cfg/flax/models/ and cfg/keras/models/ respectively. Flax generates nnx.Module classes; Keras generates Layer classes. Both use channel-last tensor layout (H × W × C) instead of PyTorch's channel-first (C × H × W).

3. Generate a Learner Class

The learner is the object that owns the models and defines how they learn: when an update happens, how a training step runs, and how an inference step runs. Losses and metrics are part of it — they are declared inline in the learner's flow, so there is no separate loss or metric command.

scm torch create learner cfg/torch/learners/ConvNeXtV2.yaml -p 'DEFAULT: {epochs: 5}' -o learner.py

Options: -p/--parameter overrides template parameters, -c/--classname sets the generated class name (default Learner), and -o/--output sets the output path.

The generated class manages:

• a training-time execution graph (FLOW) and an inference-time execution graph (INFERENCE_FLOW) per learner entry
• inline layer instantiation (loss layers, metric layers, and arbitrary modules can be defined directly in the flow)
• one or more LEARNERS entries, each with its own optimizer and trainable layers — enabling multi-optimizer training (e.g., GAN generator + discriminator)
• optimizer construction via StructCast patterns, including file-addressed optimizer + scheduler compositions such as examples/torch/optimizers.py
• optional gradient scaler creation (MIXED_PRECISION)
• optional gradient clipping (CLIP)
• optional gradient accumulation (ACCUMULATE_GRADIENTS)
• optimizer stepping, zeroing, and automatic train/eval mode switching
• learning-rate and parameter-group inspection helpers

The result implements the Learner protocol — the models, optimizers, and learning_rates properties plus update, training_step, and inference_step — and the optional grad_scalers, weight_decays, and param_group_names properties the toolkit reads when present (the loggers merge learning_rates and weight_decays into the epoch metrics). Any object with those members can be trained, generated or hand-written; see examples/torch/simple_training.py.

For example, a CycleGAN learner template defines three LEARNERS entries — one for the generator pair and one for each discriminator — each with its own flow, optimizer, and trainable layers:

scm torch create learner cfg/torch/learners/CycleGAN.yaml -o learner.py

4. Inspect FLOPs and Parameters

Once a model has been generated, you can instantiate it from a StructCast pattern and measure its computational complexity.

scm torch ptflops '[_obj_, {_addr_: model.Model, _file_: model.py}, _call_]' \
    -s 'image: [3, 224, 224]' \
    --backend pytorch

scm torch calflops '[_obj_, {_addr_: model.Model, _file_: model.py}, _call_]' \
    -s 'image: [3, 224, 224]'

What these commands do internally:

1. Instantiate the model from the _obj_ pattern.
2. Create dummy tensors from the -s/--shape specification.
3. Run one initialization forward pass via initial_model(...).
4. Pass the initialized model to ptflops or calflops for complexity analysis.

5. Measure Inference Time

Use scm [torch/flax/keras] time to benchmark the average forward-pass latency of a generated model. All three frameworks share the same basic options:

PyTorch example:

scm torch create model cfg/torch/models/ConvNeXtV2.yaml \
    -p 'DEFAULT: {backbone: atto}' -o torch_model.py

scm torch time \
    '[_obj_, {_addr_: model.Model, _file_: torch_model.py}, _call_]' \
    -s 'image: [3, 224, 224]' \
    -c cfg/torch/others/compile_default.yaml \
    -d cuda

PyTorch-specific option: --matmul-precision (highest, high, medium) controls torch.set_float32_matmul_precision.

Flax and Keras — Replace scm torch with scm flax or scm keras. Both use channel-last shapes (e.g., 'image: [224, 224, 3]'). Flax additionally accepts --training-mode-kwargs to override keyword arguments for nnx.view. Keras compilation uses keras.Model.compile. When using the Keras JAX backend on GPU, you may need to set LD_LIBRARY_PATH to include NVIDIA shared libraries from your virtual environment.

6. Train a Generated Model

Below is the complete training command from the included ConvNeXtV2 example.

scm torch train \
    'model: [_obj_, {_addr_: model.Model, _file_: model.py}, _call_]' \
    -s 'image: [3, 224, 224]' \
    -d cuda \
    -L '[_obj_, {_addr_: learner.Learner, _file_: learner.py}]' \
    -c cfg/torch/others/compile_default.yaml \
    -e 5 \
    --training-dataset dataset_train.yaml \
    -V dataset_valid.yaml \
    -f 1 \
    -LC ce_loss \
    -LC val_ce_loss \
    -HC acc1 \
    -HC val_acc1 \
    -HC acc5 \
    -HC val_acc5 \
    -SC val_acc1 \
    --matmul-precision high \
    --logger mlflow \
    -E Test \
    -A model.py \
    -A learner.py \
    -A cfg/torch/others/compile_default.yaml \
    -A dataset_train.yaml \
    -A dataset_valid.yaml

Key arguments:

• positional model patterns: one or more named model definitions
• -s/--shape: dummy input shapes used for model initialization
• -d/--device: cpu or cuda
• -L/--learner: StructCast pattern for the learner class; it is called with the instantiated models as keyword arguments
• -LO/--learner-outputs: criterion names to track, when the learner exposes no outputs attribute
• -c/--compile: boolean, YAML file, or inline dict for torch.compile
• --training-dataset: training dataset pattern or rendered dataset YAML
• -V/--validation-dataset: validation dataset pattern or rendered dataset YAML; omit it to skip validation
• -f/--validation-frequency: run validation every N epochs
• -LC/--lower-criterion: criteria where lower is better
• -HC/--higher-criterion: criteria where higher is better
• -SC/--save-criterion: criteria that should trigger best-model saving
• --logger: experiment tracking service, mlflow (default) or wandb
• -E/--experiment: experiment name passed to the logger
• -A/--log-artifacts: files to store as run artifacts
• --trainer: StructCast pattern for a TorchTrainer replacement, when the default loop is not enough
• --strategy: StructCast pattern for the DistributedStrategy; it is called with the resolved device and local_rank. Defaults to DistributedDataParallelStrategy when a distributed environment is detected, and SingleDeviceStrategy otherwise
• --resume: training state to restore before the loop starts; the reference is resolved by the active --logger, so a local path always works, a runs:/<run_id>/<artifact> MLflow URI requires --logger mlflow, and a wandb://<entity>/<project>/<run_id>/<file> reference requires --logger wandb — resuming across services is not supported. Models, optimizers, and gradient scalers are restored and training continues from the saved epoch (--start-epoch is overridden, with a warning)

What the train command does internally:

1. Instantiates the datasets and composes them into a SimpleDataProvider, which reports steps_per_epoch and validation_steps. The trainer scans the provider datasets for event protocols, so a dataset implementing one receives the lifecycle events it defines.
2. Builds the models from their patterns on the training device, initializes them with optional dummy-input forward passes, applies the initializers on rank 0 and broadcasts the result (sync_initial_weights), then compiles each model where the strategy places the units and hands it to the strategy, which wraps it. The learner is built from the already-wrapped models.
3. Builds a TorchTracker from the learner's output names, still inside the device scope so its buffers live on the training device.
4. Compiles the learner's generated _flow_* functions — the pure-compute part of each step — on a single device only. train()/eval(), backward, optimizer steps, and zero_grad() stay eager.
5. Creates the TorchTrainer with the learner, the tracker, and the data provider.
6. Collects the callbacks from the trainer's prefixes: a ProgressBar (or a Printer under --ci) and the logger on rank 0 only, plus a training-state saver and one TorchBestCriterion per monitored criterion on every rank — producing their states is a collective, and off rank 0 they hold a NullLogger and write nothing. They join the trainer's events on first use, and the resulting routing is printed.
7. Runs fit() inside the logger's run context, recording metrics, arguments, model states, optimizer states, gradient scaler states, and best checkpoints.

Distributed Training with torchrun

scm torch train supports multi-GPU and multi-node distributed data parallel (DDP) training out of the box via torchrun. No changes to your generated code, YAML templates, or dataset configurations are required — the same scm torch train command works for both single-GPU and distributed training.

⚠️ SyncBatchNorm Warning

When using multi-GPU training, scm torch train does not automatically convert BatchNorm layers to SyncBatchNorm. Standard BatchNorm computes statistics per-GPU, which can cause inconsistent behavior across ranks — especially with small per-GPU batch sizes. If your model contains BatchNorm layers and you are training distributed, apply torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) at model construction time, since the CLI wraps the models with the distributed strategy right after the initializers run. This conversion must happen in user code or in the model definition; the CLI will not perform it for you.

How It Works

When launched through torchrun, the environment variables RANK, LOCAL_RANK, WORLD_SIZE, MASTER_ADDR, and MASTER_PORT are set automatically. scm torch train detects these and enables distributed mode:

1. Process group initialization — The NCCL backend is initialized via torch.distributed.init_process_group.
2. Per-rank device assignment — Each process is assigned to cuda:<LOCAL_RANK>.
3. Strategy model wrapping — Every model is wrapped by the selected DistributedStrategy before the learner is built. The default in a distributed environment is DistributedDataParallel; SingleDeviceStrategy and FullyShardedDataParallelStrategy (FSDP2, requires torch>=2.6) are selectable through --strategy.
4. Distributed data loading — The example TimmDataLoaderWrapper automatically creates a DistributedSampler when a distributed environment is detected. Per-epoch reshuffling additionally needs the sampler's set_epoch(), which the wrapper issues from its own on_epoch_begin; the trainer scans the provider datasets for event protocols on every rank, so the hook runs everywhere it must.
5. Metric synchronization — TorchTracker uses all_reduce to average loss and metric values across all ranks.
6. Rank-0 logging — Experiment logging and progress bars run only on rank 0. Checkpoint states are produced on every rank, because the strategy's state dict is a collective, and written only by rank 0.
7. Gradient sync gating — Generated learners precede every model call with a sync_gate(model, armed) statement. Gradients synchronize only on the last call of a model owned by the running optimizer segment, on steps that update; every other call runs without synchronization, which covers gradient accumulation.
8. Cleanup — torch.distributed.destroy_process_group() is called when training finishes.

Single-Node Multi-GPU

To train on all GPUs of a single machine, prefix your scm torch train command with torchrun:

# Use all available GPUs on the current machine
torchrun --nproc_per_node=gpu \
    -m structcast_model.commands.main \
    torch train \
    'model: [_obj_, {_addr_: model.Model, _file_: model.py}, _call_]' \
    -s 'image: [3, 224, 224]' \
    -d cuda \
    -L '[_obj_, {_addr_: learner.Learner, _file_: learner.py}]' \
    -c cfg/torch/others/compile_default.yaml \
    -e 5 \
    --training-dataset dataset_train.yaml \
    -V dataset_valid.yaml \
    -f 1 \
    -LC ce_loss -LC val_ce_loss \
    -HC acc1 -HC val_acc1 -HC acc5 -HC val_acc5 \
    -SC val_acc1 \
    --matmul-precision high \
    -E Test

Or specify an exact GPU count:

# Use exactly 4 GPUs
torchrun --nproc_per_node=4 \
    -m structcast_model.commands.main \
    torch train ...

Note: torchrun launches the training script as a Python module (-m structcast_model.commands.main) rather than through the scm entry point. This is because torchrun requires a module or script path, not a console script wrapper.

Multi-Node Training

For training across multiple machines, provide the node topology to torchrun on each node:

# On node 0 (master)
torchrun \
    --nproc_per_node=4 \
    --nnodes=2 \
    --node_rank=0 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    -m structcast_model.commands.main \
    torch train ...

# On node 1
torchrun \
    --nproc_per_node=4 \
    --nnodes=2 \
    --node_rank=1 \
    --master_addr=192.168.1.100 \
    --master_port=29500 \
    -m structcast_model.commands.main \
    torch train ...

This creates 8 total processes (4 GPUs × 2 nodes) training with DDP.

torchrun parameters:

scm torch train distributed-related options:

Dataset Configuration

Dataset YAML files do not need per-rank customization. A single device: cuda value in the dataset configuration works for all ranks — the example TimmDataLoaderWrapper internally resolves it to the correct cuda:<LOCAL_RANK> device for each process.

# The same dataset YAML works for single-GPU and distributed training
scm format cfg/torch/others/default_timm.yaml \
    -o dataset_train.yaml \
    -p 'DEFAULT: {training: true, epochs: 5, batch_size: 32, dataset: torch/cifar100, num_classes: 100, label_smoothing: 0.1, input_size: [3, 224, 224], image_dtype: bfloat16, download: true}'

Tip: The batch_size in the dataset template is the per-GPU batch size. With 4 GPUs and batch_size: 32, the effective global batch size is 128.

Distributed Training Notes

• Seed reproducibility — Each rank's random seed is offset by global_rank to ensure different data augmentation across processes while remaining reproducible.
• Learning rate scaling — When scaling to multiple GPUs, consider adjusting the learning rate. A common practice is linear scaling: multiply the base learning rate by the number of GPUs. This must be configured in the learner template or optimizer settings — scm torch train does not scale the learning rate automatically.
• SyncBatchNorm — scm torch train does not automatically convert BatchNorm layers to SyncBatchNorm. If your model uses BatchNorm and you are training with DDP, consider applying torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) in the model definition. See the SyncBatchNorm warning for details.
• torch.compile and the strategy — with --compile, the strategy decides where its compile units sit: the model root in place by default, the matched shard_modules blocks under per-block FSDP2 — always before wrapping, so the strategy wrapper stays outermost. The learner's generated _flow_* functions compile on a single device only (distributed wrappers graph-break inside them); the eager step methods are never compiled.
• Checkpoint saving — State dicts are produced through torch.distributed.checkpoint.state_dict, so the keys are wrapper-free for raw, compiled, DDP, and FSDP2 models alike. Producing them is a collective that runs on every rank; only rank 0 writes them to the experiment tracking service. --resume loads the same training state on all ranks.

Training Loop Anatomy

Whether it is built by the CLI or by hand, a training run is the same five objects handed to a trainer at construction:

trainer = TorchTrainer(
    device="cpu",
    learner=learner,
    tracker=tracker,
    data=SimpleDataProvider(training_dataset=training_dataset, validation_dataset=validation_dataset),
    callbacks=[Printer(), BestCriterion(target="val_loss", mode="min")],
)
trainer.fit(epochs=3)

There is no registration call and no global registry. Every participant — the learner, the learner's optimizers, the tracker, the data provider and its datasets, then the callbacks in the order given — is scanned once on first use — the first dispatched event; describe() only previews the routing — and is routed into each lifecycle event whose protocol it implements:

on_update, on_training_begin, on_training_end, on_training_step_begin, on_training_step_end, on_validation_begin, on_validation_end, on_validation_step_begin, on_validation_step_end, on_epoch_begin, on_epoch_end.

An object joins an event simply by defining the matching method; trainer.describe() shows the resulting routing. This is how an optimizer + scheduler composition steps its schedule, how TorchTracker resets its averages between training and validation, and how a logger records epoch metrics — all through the same mechanism.

Datasets arrive at construction through the data provider, so fit(epochs, start_epoch, validation_frequency) takes loop parameters only. train(dataset) and evaluate(dataset) remain available for a single pass over a dataset.

For a complete, commented program built from these pieces, see examples/.

Configuration Examples

The cfg/ directory contains working YAML templates that demonstrate each part of the workflow. Templates are organized by framework under cfg/torch/, cfg/flax/, and cfg/keras/. For schema details on every key used below, see REFERENCE.md.

PyTorch

cfg/torch/models/ConvNeXtV2.yaml — Demonstrates the model-building style used throughout the project. The root model defines the top-level execution flow, and sublayer keys (Backbone, Block, etc.) define reusable nested modules:

# Root model: routes tensors through backbone → pooling → classifier
INPUTS: [image]
OUTPUTS: [cls]
FLOW:
  - [image, {feature: feat4}, backbone, {TYPE: Backbone}]
  - [feature, _, [_obj_, {_addr_: torch.nn.AdaptiveAvgPool2d}, {_call_: {output_size: 1}}]]
  - [_, _, [_obj_, {_addr_: torch.nn.Flatten}, _call_]]
  - # ... LayerNorm (Jinja-expanded from backbone dims) ...
  - [_, cls, head, [_obj_, {_addr_: torch.nn.LazyLinear}, {_call_: {out_features: 1000}}]]

Parameter groups define multiple backbone sizes, and Jinja rendering expands blocks based on depths and dims:

PARAMETERS:
  DEFAULT:
    backbone: atto
  SHARED:
    stem_kernel_size: 4
    kernel_size: 7
    drop_path_rate: 0.0
    num_classes: 1000
  atto:
    dims: [40, 80, 160, 320]
    depths: [2, 2, 6, 2]
  femto:
    dims: [48, 96, 192, 384]
    depths: [2, 2, 6, 2]
  # ... tiny, small, base, large, huge ...

The Block sublayer shows how a single convolutional block is defined with depthwise convolution, normalization, MLP expansion, GRN, and residual addition:

Block:
  OUTPUTS: [out]
  _jinja_yaml_: |-
    FLOW:
      - INPUTS: inp
        OUTPUTS: _
        LAYER:
          - _obj_
          - _addr_: torch.nn.LazyConv2d
          - _call_: {out_channels: {{fout}}, kernel_size: {{kernel_size}}, groups: {{fout}}, padding: "eval: {{kernel_size}} // 2"}
      - [_, _, [_obj_, {_addr_: structcast_model.torch.layers.ToChannelLast}, _call_]]
      - [_, _, [_obj_, {_addr_: timm.layers.LayerNorm}, {_call_: {num_channels: {{fout}}, eps: {{norm_eps}}}}]]
      - [_, _, [_obj_, {_addr_: torch.nn.LazyLinear}, {_call_: {out_features: "eval: {{fout}} * {{mlp_ratio}}"}}]]
      - [_, _, [_obj_, {_addr_: "timm.layers.{{activation}}"}, {_call_: {inplace: true}}]]
      - [_, _, [_obj_, {_addr_: timm.layers.grn.GlobalResponseNorm}, {_call_: {dim: "eval: {{fout}} * {{mlp_ratio}}"}}]]
      - [_, _, [_obj_, {_addr_: torch.nn.LazyLinear}, {_call_: {out_features: {{fout}}}}]]
      - [_, _, [_obj_, {_addr_: structcast_model.torch.layers.ToChannelFirst}, _call_]]
      - [_, feat, {TYPE: DropPath, PARAM: {DEFAULT: {drop_prob: {{drop_path}}}}}]
      - ["eval: inp + feat", out, null]

cfg/torch/learners/ConvNeXtV2.yaml — Demonstrates a single-optimizer learner with mixed precision, gradient accumulation, cosine LR scheduling, and inline loss/metric definitions in the flow. The optimizer is a file-addressed composition from examples/torch/optimizers.py: the package builds optimizers (create_opt), while optimizer + scheduler combinations are example code you can copy and adapt:

MIXED_PRECISION:
  init_scale: "eval: 2.0**16"
  growth_factor: 2.0
  backoff_factor: 0.5
  growth_interval: 2000
  enabled: True
MIXED_PRECISION_TYPE: bfloat16
OUTPUTS: [ce_loss, acc1, acc5]
LEARNERS:
  - LOSS: ce_loss
    TRAINABLE_LAYERS: [model]
    NAME: optimizer
    OPTIMIZER:
      - _obj_
      - _addr_: AdamWWithCosine
        _file_: examples/torch/optimizers.py
      - _bind_:
          optimizer_kwargs: {opt: adamw, lr: 4.0e-3, weight_decay: 0.001}
          scheduler_kwargs: {sched: cosine, num_epochs: 300, criterion: ce_loss}
    FLOW:
      - [image, cls, model]
      - [{target: label, input: cls}, ce_loss, cross_entropy_loss, [_obj_, _addr_: torch.nn.CrossEntropyLoss, _call_]]
      - [{y_true: label, y_pred: cls}, acc1, accuracy, [_obj_, _addr_: torch.no_grad, _call_, _call_: [[_obj_, {_addr_: structcast_model.torch.layers.sparse_categorical_accuracy}]]]]
      - [{y_true: label, y_pred: cls, k: 5}, acc5, top_5_accuracy, [_obj_, _addr_: torch.no_grad, _call_, _call_: [[_obj_, {_addr_: structcast_model.torch.layers.sparse_top_k_categorical_accuracy}]]]]

cfg/torch/learners/CycleGAN.yaml — Demonstrates a multi-optimizer learner for GAN-style training with three LEARNERS entries (generator pair + two discriminators), each with its own flow, optimizer, and trainable layers.

cfg/torch/models/CycleGAN_generator.yaml and CycleGAN_discriminator.yaml — Pair of model templates for the CycleGAN architecture:

• Generator — uses ResidualBlock, DownBlock, and UpBlock sublayers with reflection padding, instance normalization, and Jinja-driven residual block expansion (n_residual_blocks parameter)
• Discriminator — uses a DiscriminatorBlock sublayer with conditional instance normalization controlled by a normalize parameter
• both templates use LazyConv2d for automatic input channel inference

cfg/torch/others/default_timm.yaml — Formats directly into a TimmDataLoaderWrapper.model_validate(...) pattern, loading the wrapper from the example file examples/torch/data.py by path. The template covers timm dataset and dataloader construction, device and prefetch settings, mixup/cutmix configuration, and train/validation split generation — all from a single parameterized template:

_obj_:
  - _addr_: TimmDataLoaderWrapper
    _file_: examples/torch/data.py
  - _attr_: model_validate
  - - _call_
    - spec: {image: "0", label: "1"}
      dataset:
        input_img_mode: RGB
        _jinja_yaml_: |-
          batch_size: {{batch_size}}
          name: {{dataset}}
          root: {{dataset_dir}}
          is_training: {{training}}
          split: {{"train" if training else "validation"}}
          # ...
      use_prefetcher: true
      mixup_alpha: 0.0
      cutmix_alpha: 0.0
      # ...

Flax

cfg/flax/models/ConvNeXtV2.yaml — Generates a Flax nnx.Module equivalent of the PyTorch ConvNeXtV2 model. The template mirrors the same parameter groups (atto through huge) and uses GlobalResponseNorm as a custom Flax layer. Key differences from the PyTorch variant:

• uses channel-last tensor layout (H × W × C)
• constructor accepts a rngs: flax.nnx.Rngs argument for parameter initialization
• __call__ propagates a training flag to sub-modules
• layer APIs differ (e.g., flax.nnx.Conv instead of torch.nn.LazyConv2d)

Keras

cfg/keras/models/ConvNeXtV2.yaml — Generates a Keras Layer equivalent of the ConvNeXtV2 model. Shares the same backbone parameter groups and uses GlobalResponseNormalization as a custom Keras layer. Key differences:

• uses channel-last tensor layout (H × W × C)
• follows the Keras call(self, ..., *, training=None, **kwargs) convention
• runs on any Keras backend (JAX, PyTorch, or TensorFlow)
• uses keras.layers.Add for residual connections instead of "eval: inp + feat" expressions

Development

Set up the development environment with:

uv sync --extra torch-cpu --dev --group tox

Run the test suite:

pytest

Run static type checks:

mypy src
mypy tests

Run linting and formatting:

ruff check src tests
ruff format src tests

Run all checks in parallel with:

tox run-parallel --parallel all

The repository includes tests for:

• CLI behavior
• Builder code generation
• Schema validation
• Trainer utilities
• timm dataset and dataloader wrappers
• Custom torch layers

Migration Notes

Upgrading to v2.x

The training loop was redesigned around protocol-routed callbacks. The rationale is recorded in docs/adr/0002-protocol-routed-training-loop.md; the vocabulary in CONTEXT.md. There are no compatibility aliases:

• Backward is now Learner — The rename cascades through the runtime, the CLI (scm torch create learner, --learner/-L), the builder and schema names (LEARNERS, LearnerBehavior, UserDefinedLearner), and the template directory (cfg/torch/learners/).
• Callbacks are routed by protocol — The GLOBAL_CALLBACKS registry, the callbacks_session context manager, and NamedCallbackList.register() are gone. Pass participants to the trainer as callbacks=[...]; each one joins the events whose on_* method it defines. Ad-hoc lambdas become small callback classes — ProgressBar and Printer ship with the package.
• Datasets are given at construction — fit() no longer takes datasets. Build a DataProvider (SimpleDataProvider, or your own object with training_dataset, validation_dataset, steps_per_epoch, and validation_steps — the dataset properties must return the same object on every read, since the trainer reads them for the event scan and again in fit()) and pass it as data=. The trainer also scans the provider datasets for event protocols, so a dataset with an on_* hook (e.g. a distributed sampler wrapper) takes part in the loop without being passed as a callback. fit() keeps epochs, start_epoch, and validation_frequency; train(dataset) and evaluate(dataset) are unchanged.
• create_with_scheduler is removed — The package keeps create_opt (regex weight-decay and layer-decay grouping over torch.optim and timm engines). Optimizer + scheduler combinations move to example code referenced by file path; AdamWWithCosine (timm schedules) and OptimizerWithNativeScheduler (per-epoch native schedules) in examples/torch/optimizers.py cover the cosine and per-epoch native cases and also keep the schedule in their state_dict; metric-driven (ReduceLROnPlateau), per-update, and composite schedules need a wrapper of their own modeled on these.
• Loggers own the run — MLflowLogger (structcast_model.torch.mlflow_logger) and WandbLogger (structcast_model.torch.wandb_logger) are context managers that start and end the run and log epoch metrics; both follow the Logger protocol in structcast_model.torch.logger. Select the backend with --logger mlflow|wandb.
• Trackers reset themselves — TorchTracker clears its averages from on_training_begin and on_validation_begin; the explicit reset() call in the loop is gone.

Upgrading from v1.x

The following breaking changes were introduced by the learner-template restructure for multi-optimizer GAN training support:

• EMA support removed — TimmEmaWrapper, the cfg/torch/others/ema.yaml configuration, and all InferenceWrapper-based EMA integration in cmd_torch.py and torch/trainer.py have been removed. If your training workflow relied on built-in EMA, you will need to manage EMA externally.
• Learner template schema restructured — The LEARNERS key expects a list of LearnerBehavior entries (each with its own NAME, LOSS, TRAINABLE_LAYERS, OPTIMIZER, FLOW, and optional INFERENCE_FLOW). Previous single-optimizer configurations must be wrapped in a single-entry list.
• Separate loss and metric templates removed — Losses and metrics are declared inline in the learner's flow, so scm torch train no longer takes --loss or --metric.

Roadmap

• PyTorch model construction from YAML configuration files
• PyTorch training workflow generation from YAML configuration files
• JAX (Flax) model construction from YAML configuration files
• JAX (Flax) training workflow generation from YAML configuration files
• Keras model construction from YAML configuration files
• Keras training workflow generation from YAML configuration files