SR-Forge

Writing Scripts

SR-Forge ships with ready-to-use train and test scripts that work entirely from YAML config. Most users will never need to modify them. But if you need custom logic — a different training loop, a GAN setup, transfer learning — you can write your own script.

This page shows how to write a script from scratch and explains each building block.

---

The srforge CLI

Installing sr-forge registers a srforge console script (pyproject.toml → [project.scripts]) that handles project bootstrapping and config diagnostics. Two subcommands:

srforge init [--force]

Scaffolds a complete training project in the current directory by copying four files from srforge/_scaffold/:

File	Purpose
train.py	Hydra-wrapped training script — calls init(cfg), builds the trainer + datasets + tracker, runs resume_from_checkpoint against trainer.training_runner, hands control to trainer.train(...).
benchmark.py	Inference / metric-evaluation script using BenchmarkRunner.
configs/train-cfg.yaml	Full training config: system, model, optimizer, scheduler, loss, datasets, runners, trainer — with hooks attached inside the runner/trainer sections.
configs/benchmark-cfg.yaml	Inference-side config.

cd ~/my-project
srforge init                  # creates the files above, skipping any that already exist
srforge init --force          # overwrites existing copies

Each file is a starting point — edit the YAML to match your data and model. The scripts call init(cfg) (next section), so anything they do is already covered by this documentation.

srforge audit <path> [--format text|json|graphml]

Statically extracts every _target dispatch and every ${ref:...} reference from a YAML config without instantiating anything. It also runs full validation: cycle detection in the indirection graph, dangling-reference detection, missing-_target detection. Useful during refactors when you want to confirm a config still parses before booting CUDA or downloading checkpoints.

srforge audit configs/train-cfg.yaml                     # human-readable table
srforge audit configs/train-cfg.yaml --format json       # array of {kind,source_path,raw,resolved}
srforge audit configs/train-cfg.yaml --format graphml    # GraphML (load with networkx / graphify merge-graphs)

Under the hood the command instantiates a ConfigResolver (which builds an in-memory ConfigGraph) and emits each indirection edge as an IndirectionEdge — see Config Indirection for the full type surface. The same audit runs automatically inside init() (see step 4 below) — running it standalone is the way to catch problems without booting the rest of the framework.

---

Minimal Training Script

Here is the shortest working training script:

import hydra
import omegaconf

from srforge import GlobalSettings, init
from srforge.data.loader import DataLoaderFactory
from srforge.utils.checkpoint import resume_from_checkpoint


@hydra.main(config_path="configs", config_name="train-cfg", version_base=None)
def main(cfg):
    resolve = init(cfg)                              # (1)

    # -- Training objects --------------------------------------------------
    model     = resolve(cfg.model)                   # (2)
    optimizer = resolve(cfg.optimizer)
    scheduler = resolve(cfg.lr_scheduler)
    loss      = resolve(cfg.loss)
    model.to("cuda")

    # -- Datasets & loaders ------------------------------------------------
    train_loader = DataLoaderFactory(                # (3)
        resolve(cfg.dataset.training),
        batch_size=cfg.training.batch_size, shuffle=True,
    ).get_loader()
    val_loader = DataLoaderFactory(
        resolve(cfg.dataset.validation),
        batch_size=1, shuffle=False,
    ).get_loader()

    # -- Trainer (owns the training runner used by resume) -----------------
    trainer = resolve(cfg.trainer)                   # (4)

    # -- Tracker (after trainer so log_config can include any defaults) ----
    tracker = resolve(cfg.tracker)                   # (5)
    tracker.log_config(omegaconf.OmegaConf.to_container(cfg, resolve=False))

    # -- Resume from checkpoint (uses trainer.training_runner, NOT the
    # raw optimizer — optimizer + scaler state live inside the runner) ----
    ckpt = resume_from_checkpoint(                   # (6)
        model, trainer.training_runner, scheduler, tracker=tracker
    )
    trainer.restore(ckpt)

    # Hooks (progress bars, loggers, checkpointing) were attached when
    # the runners/trainer were resolved — they live in the config's
    # hooks: lists, so there's no separate subscription step here.

    trainer.train(cfg.training.epochs, train_loader, val_loader)

    tracker.finish(0)


if __name__ == "__main__":
    main()

Let's walk through each numbered step.

---

Step by Step

1. init(cfg) — Framework Setup

resolve = init(cfg)

srforge.init() is the boilerplate every script needs. It does six things and returns a ConfigResolver:

1. Configures the global logger with coloured output at the chosen log_level (default INFO; pass log_level=cfg.system.debug_level if your config has one).
2. Installs a hard-kill SIGINT/SIGBREAK handler that calls tracker.finish(1) on GlobalSettings().tracker and then os._exit(1). This is the only reliable way to stop a hung DataLoader-worker pool or an MKL/Fortran thread on Windows.
3. Prints the startup banner ("⚡ SR-Forge" panel via rich) and a PyPI update notice if a newer version exists.
4. Strips internal Hydra/OmegaConf bookkeeping keys from the config (clear_defaults) and runs a pre-flight _target audit — every _target string is best-effort imported up-front. Typos and missing classes surface as a single warning at startup instead of mid-run.
5. Wires GlobalSettings — stores the config and the Hydra-resolved output directory.
6. Returns a ConfigResolver bound to the config — your main tool: pass any subtree to it and the recursive instantiation produces the Python objects.

init() has no config key assumptions beyond a Hydra-managed output dir — it doesn't care whether your config has model, tracker, or any other specific key. That's your script's job.

Pre-flight audit

Step 4's audit is the same one you can run manually via srforge audit configs/train-cfg.yaml (see top of this page). Inside init() it logs warnings but doesn't abort — the srforge audit CLI prints the same edge list explicitly.

2. Tracker Setup

tracker = resolve(cfg.tracker)
tracker.log_config(omegaconf.OmegaConf.to_container(cfg, resolve=False))

The tracker provides experiment logging (metrics, images, checkpoints). Resolve it from config and log the full config as metadata.

Hooks that need the tracker (e.g., LossLogger, PyTorchModelSaver, BatchImageLogger) receive it as a constructor parameter via tracker: ${ref:tracker} in the YAML config. No need to store it on GlobalSettings.

See Experiment Tracking for details on the tracker abstraction.

3. Resolve Training Objects

model     = resolve(cfg.model)
optimizer = resolve(cfg.optimizer)
scheduler = resolve(cfg.lr_scheduler)
loss      = resolve(cfg.loss)

Each resolve() call instantiates the class specified by _target with the params from YAML. Results are cached — calling resolve(cfg.model) twice returns the same Python object. This is how ${ref:model} references work: when the optimizer config says params: ${ref:model}.trainable_params(), the resolver returns the already-created model instance.

Order doesn't matter. The resolver handles forward references lazily.

4. Resume from Checkpoint

ckpt = resume_from_checkpoint(
    model, trainer.training_runner, scheduler, tracker=tracker
)

Signature: resume_from_checkpoint(model, runner, lr_scheduler, tracker=None). The 2^nd argument is the runner, not the optimizer — sr-forge keeps the optimizer + AMP scaler state inside the TrainingEpochRunner, so the runner's load_training_state(state) method is what reads them. Use trainer.training_runner to get the runner the trainer created.

One call that handles the entire resume flow:

• If the tracker says the run was resumed (tracker.is_resumed), it loads the last checkpoint, restores runner state (optimizer + scaler), loads model weights, restores LR-scheduler state, and restores RNG states for exact reproducibility.
• If not resumed (or tracker is None), it returns None.

You pass the result to trainer.restore() later — it accepts None as a no-op.

5. Hooks — no script step needed

Progress bars, loss logging, checkpoint saving, and image previews are all hooks, and they live in the config — inside the hooks: lists of training_runner, validation_runner, and trainer. When resolve() builds those components, the hooks are constructed and attached automatically. The script does nothing.

To add or remove monitoring, edit the YAML — see Hooks.

Migrating from the legacy observers: pattern

Scripts written for older versions had an explicit step here: observers = resolve(cfg.observers) followed by GlobalSettings().event_bus.subscribe(observers). That mechanism still works via a compatibility adapter but is deprecated — removal in v0.16.0. See the porting guide and Observers & Events for the legacy reference.

6. Trainer

trainer = resolve(cfg.trainer)
trainer.restore(ckpt)
trainer.train(cfg.training.epochs, train_loader, val_loader)

The trainer orchestrates the epoch loop. restore(ckpt) sets the initial epoch, best losses, and restores the GradScaler state from the checkpoint. It accepts None (fresh run) as a no-op. See Trainers & Runners for a detailed walkthrough of the trainer and runner internals.

7. Cleanup

tracker.finish(0)

Finalizes the tracker (flushes logs, uploads remaining files). Always call this at the end.

---

GlobalSettings

GlobalSettings is a singleton that stores application-wide state. The framework uses it to find the event bus and output directory.

Attribute	Type	Description
config	DictConfig	The full experiment config (set by init)
output_directory	str	Hydra's output directory for the current run (set by init)
event_bus	EventBus	(legacy) global event bus for deprecated Observer dispatch — removal in v0.16.0
tracker	ExperimentTracker \| None	The active tracker (set by the tracker's own __init__)
debug_mode	bool	True when the logger is at DEBUG level

init(cfg) sets config and output_directory automatically. Each tracker assigns itself to GlobalSettings().tracker when constructed (see WandbTracker.__init__) — this is how the SIGINT handler can call tracker.finish(1) to flush logs before the hard kill. Hooks that need the tracker still receive it as a constructor parameter via tracker: ${ref:tracker} in YAML (preferred over reading the global), but the global slot exists.

---

ConfigResolver

ConfigResolver is the engine that turns YAML into Python objects. Key behaviors:

• _target + params — instantiate any class
• Recursive — nested _target blocks are resolved bottom-up
• Cached — each config path is resolved once and reused
• ${ref:path} — references to already-instantiated objects
• io: key — automatically calls set_io() on IOModule subclasses
• Runtime kwargs — resolve(cfg.section, key=value) injects extra constructor arguments

See Configuration for the full reference.

---

Writing a Test Script

A test/benchmark script is simpler — no optimizer, scheduler, or resume:

import hydra
import omegaconf
import torch

from srforge import GlobalSettings, init
from srforge.data.loader import DataLoaderFactory
from srforge.training.runners import BenchmarkRunner


@hydra.main(config_path="configs", config_name="test-cfg", version_base=None)
def main(cfg):
    resolve = init(cfg)

    tracker = resolve(cfg.tracker)
    tracker.log_config(omegaconf.OmegaConf.to_container(cfg, resolve=False))

    model         = resolve(cfg.model)
    metrics       = resolve(cfg.test_metrics)
    postprocessor = resolve(cfg.postprocessing)
    test_dataset  = resolve(cfg.dataset)

    device = torch.device(cfg.system.device)
    model.to(device)

    test_loader = DataLoaderFactory(
        test_dataset, batch_size=1, shuffle=False,
    ).get_loader()

    runner = BenchmarkRunner(device=device, postprocessor=postprocessor,
                             hooks=resolve(cfg.hooks))
    runner.run_epoch(model=model, data_loader=test_loader, epoch=0, criterion=metrics)

    tracker.finish(0)


if __name__ == "__main__":
    main()

The pattern is the same: init -> resolve objects -> set up data -> run.

---

Custom Training Loops

If PyTorchTrainer's default TrainingEpochRunner doesn't fit your needs, you have two options. Important: for GAN-style alternating updates, sr-forge ships a ready-made GANTrainingRunner — don't reinvent it.

Option A: Swap in an alternative runner

PyTorchTrainer accepts any EpochRunner subclass as its training_epoch_runner. For GANs, plug in GANTrainingRunner:

# A GAN trainer is still a PyTorchTrainer — only the runner changes.
trainer:
  _target: srforge.training.PyTorchTrainer
  params:
    model: ${ref:gan_model}             # an srforge.models.GANModel
    training_epoch_runner:
      _target: srforge.training.runners.GANTrainingRunner
      params:
        optimizer_G: ${ref:optimizer_G}
        optimizer_D: ${ref:optimizer_D}
        d_criterion:
          _target: srforge.loss.adversarial.RaGANDiscriminatorLoss
        g_criterion:
          _target: srforge.loss.adversarial.RaGANGeneratorLoss
          params: { weight: 0.01 }
        step_ratio: 1.0                 # G updates per D update
        device: ${system.device}
        mixed_precision: ${system.mixed_precision}
        hooks: []                       # plug in StepRatio, DWarmup, etc.
    validation_epoch_runner:
      _target: srforge.training.runners.ValidationEpochRunner
      params: { device: ${system.device}, criterion: ${ref:val_metrics} }
    training_criterion: ${ref:g_criterion}
    validation_criterion: ${ref:val_metrics}
    lr_scheduler: ${ref:lr_scheduler}

The script stays the same — trainer = resolve(cfg.trainer) builds both the trainer and the runner. The runner handles alternating G/D updates automatically; step_ratio controls the cadence.

If you do need a genuinely different trainer abstraction (one that doesn't conform to the train/validate-per-epoch shape), subclass srforge.training.trainers.Trainer (abc.ABC, Observable) and implement train(epochs, train_loader, val_loader). Most users should not need this — pick a different runner first.

Option B: Inline the Loop in the Script

For one-off experiments, you can skip the trainer entirely and write the loop directly:

for epoch in range(resume.initial_epoch, cfg.training.epochs):
    model.train()
    for batch in train_loader:
        batch = batch.to(device)
        output = model(batch)
        loss_scores = loss(output)
        loss_scores.total_weighted().mean().backward()
        optimizer.step()
        optimizer.zero_grad()

    scheduler.step()

This gives you full control but loses the runner/hook infrastructure (progress bars, checkpointing, and logging hooks fire from runners and the trainer).

---

Multi-GPU

For multi-GPU DataParallel training, the device config is a list:

system:
  device: [0, 1]

The script uses setup_device() which handles everything — device detection, model.to(), and DataParallel wrapping:

from srforge.utils.multigpu import setup_device

model, device = setup_device(model, cfg.system.device, train_dataset)

For single-GPU or CPU, it simply moves the model. For multi-GPU, it auto-detects whether to use torch.nn.DataParallel or torch_geometric.nn.DataParallel based on the dataset's element type.

---

Summary

Building Block	What It Does
init(cfg)	Framework setup, returns ConfigResolver
resolve(cfg.section)	Instantiate any config section
resolve(cfg.section, key=val)	Inject runtime values
hooks: lists on runner/trainer configs	Attach hooks (progress bars, loggers, checkpointing)
resume_from_checkpoint(model, runner, lr_scheduler, tracker=)	One-call checkpoint resume (runner exposes optimizer+scaler state)
setup_device(model, device_config, dataset)	Device placement + DataParallel
DataLoaderFactory(...).get_loader()	Create data loaders
trainer.train(epochs, train_loader, val_loader)	Run the training loop
tracker.finish(0)	Finalize the run

---

Next: Configuration — Define entire experiments in YAML