Skip to content
SR-Forge

Writing Scripts

srforge init writes a working train.py for you, and most users never change it. This page is about the times you do: it walks that script block by block, so you know what each line is for before you move it, and then shows the customisations that come up most — a different training loop, a GAN, a test script, multiple GPUs.

It assumes the generated script's shape (Hydra + YAML). For whether to author your experiments that way at all — versus building objects in Python, or splitting the difference — see Code or Config?.


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.


The Generated Script

This is what srforge init writes, and what the rest of this page dissects. Every block is explained in Step by Step below.

Before anything else: every SR-Forge script calls srforge.init()

The generated script below, a copy of it you have edited beyond recognition, or one you started from an empty file — all three call it, and call it first. It is not part of the config machinery: it sets up the process, and none of that depends on whether you wrote any YAML.

resolve = init(cfg)   # config-driven — also pre-flights every _target,
                      # installs the precision policy, returns a resolver
srforge.init()        # hand-wired — no config, no Hydra, no argument

It configures logging, picks and creates the run's output directory, disables OpenCV's thread pool (the classic DataLoader deadlock), installs the Ctrl+C handler, and populates GlobalSettings — the singleton your hooks read for the output directory, the precision policy and the tracker.

Skip it and none of that exists — the first thing to touch GlobalSettings stops with "GlobalSettings not initialized. Call srforge.init() before accessing 'output_directory'", which for a script that saves checkpoints means at the end of the first epoch. Step 1 lists everything it does.

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 all of the following 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 the active tracker (every tracker registers itself when built) 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. Disables OpenCV's internal thread pool (cv2.setNumThreads(0)) — the classic DataLoader deadlock. Forked workers inherit that pool's locks mid-flight and freeze in bootstrap before their first Python instruction: a hang with no traceback. Parallelism comes from the workers, not intra-image threading, so nothing is lost. Opt out with SRFORGE_CV2_THREADS=1 (e.g. pure-inference scripts that never fork). With DEBUG_HANG=1 it also dumps every thread's stack every five minutes, so a silent hang reports its own blocking frame.
  4. Prints the startup banner ("⚡ SR-Forge" panel via rich) and a PyPI update notice if a newer version exists.
  5. 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.
  6. Installs the global precision policy from system.precision, so every runner shares one fp16/bf16/TF32 decision instead of each parsing its own flag. TF32 backend flags are only touched when the config asks for them.
  7. Wires GlobalSettings — stores the run's output directory, and creates that directory. Hydra's run dir when Hydra launched the script; Hydra's own outputs/<date>/<time>/ convention when it didn't; or whatever you pass as output_dir=.
  8. Seeds Python, NumPy and torch when a seed is given — seed= or system.seed. Under torchrun each rank adds its rank, so replicas draw different augmentation while DDP still starts them from rank 0's weights. Without a seed nothing is seeded.
  9. 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 — it doesn't care whether your config has model, tracker, or any other specific key. That's your script's job.

Writing the script without Hydra

A script that builds its objects in Python needs no config and no @hydra.main — call srforge.init() bare. Three things change, and nothing else:

  • init() takes no argument, so it skips the config-dependent steps (5 and 6) and binds the returned resolver to an empty config. Everything else above — logging, output directory, OpenCV fork hygiene, Ctrl+C handling — you still get, and still need.
  • You choose the output directory with output_dir=, or let init() pick Hydra's outputs/<date>/<time>/ convention. Nothing else sets it, so a script that skips init() dies at its first checkpoint.
  • No CLI overrides. python train.py optimizer.params.lr=1e-4 is Hydra's doing; without it you edit the file or hand-roll argparse.

A complete listing is the first tab of Code or Config?, which also covers resolving part of your objects from YAML while keeping the plumbing in Python.

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 2nd 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.

Scripts written for 0.15.x and earlier

Older scripts had an explicit step here: observers = resolve(cfg.observers) followed by GlobalSettings().event_bus.subscribe(observers). Both were removed in 0.16.0 — delete those lines and see upgrading from observers.

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.


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.metrics.adversarial.RaGANDiscriminatorLoss
        g_criterion:
          _target: srforge.metrics.adversarial.RaGANGeneratorLoss
          params: { weight: 0.01 }
        device: ${system.device}
        mixed_precision: ${system.mixed_precision}
        hooks:
          - _target: StepRatio          # G updates per D update
            params: { ratio: 1.0 }
    validation_epoch_runner:
      _target: srforge.training.runners.ValidationEpochRunner
      params: { device: ${system.device} }
    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; the StepRatio hook 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).


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.


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.


GlobalSettings

GlobalSettings is a process-wide singleton holding the run's shared state. srforge.init() is what populates it — that is the main reason every script has to call it, and the reason a custom hook can rely on it being there.

Attribute Type Populated by
output_directory str init() — and created on disk
precision PrecisionPolicy init() from system.precision

config, debug_mode and tracker were removed in 0.16.0: the first two had no readers, and the tracker a Ctrl+C drains is now private — every tracker registers itself when it is built.

Reading any of them before init() tells you so

>>> GlobalSettings().output_directory
RuntimeError: GlobalSettings not initialized. Call srforge.init() before
accessing 'output_directory'.

The singleton deliberately assigns no defaults in __init__ — a default would shadow that check and hand back a value instead. So a hook of yours that reads GlobalSettings fails immediately and says which call is missing, rather than returning None and dying somewhere unrelated.

Writing first is still fine (gs.output_directory = ...), which is how tests and embedders set up a partial environment.

The output directory in a custom hook

The trainer hands it to trainer-level hooks as ctx.output_directory — the directory init() set and created, or the trainer's own output_directory= for a script that doesn't call init(). PyTorchModelSaver is the in-tree example: it reads it there rather than taking it as a parameter, because the run decides where files go, not the hook:

class MyArtefactSaver(Hook):
    @hooks_into("on_trainer_epoch_finished")
    def dump(self, ctx: Context) -> None:
        out_dir = ctx.output_directory
        ...

A runner-level hook can read GlobalSettings().output_directory.

Prefer a constructor parameter where you can. Hooks that need the tracker take tracker: ${ref:tracker} in YAML instead of reaching for anything global — it is explicit, testable, and lets one script drive two trackers.


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.


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