SR-Forge

Trainers & Runners

The Writing Scripts page shows how to wire up a training script. This page explains what happens inside once you call trainer.train() or runner.run_epoch().

Trainer orchestrates the epoch loop — it owns the criteria (the training loss and the validation metrics) and delegates the actual forward/backward work to runners. Runners are stateless execution engines: they iterate over batches, fire hook points, and return aggregated metrics. The criterion is passed to run_epoch() at call time, not stored on the runner. Hooks attached to a runner or the trainer react at those points (progress bars, checkpoints, logging) — and can even influence the loop (gradient clipping, auxiliary losses).

---

Trainer lifecycle

PyTorchTrainer.train() runs the outer epoch loop. Here's the full flow:

flowchart TD
    start([trainer.train#40;epochs, train_loader, val_loader#41;]) --> began

    began["<b>TrainingBegan</b> event<br/><i>total_epochs, initial_epoch, batches per loader</i>"]
    began --> epoch_loop

    epoch_loop{"epoch < epochs?"}
    epoch_loop -- No --> done([Return])
    epoch_loop -- Yes --> loss_sched

    loss_sched["Update LossScheduler<br/><i>if criterion is LossScheduler</i>"]
    loss_sched --> train_run

    train_run["<b>training_runner.run_epoch</b>#40;model, train_loader, epoch, criterion#41;<br/>→ train_scores"]
    train_run --> val_run

    val_run["<b>validation_runner.run_epoch</b>#40;model, val_loader, epoch, criterion#41;<br/>→ val_scores"]
    val_run --> epoch_event

    epoch_event["<b>TrainerEpochFinished</b> event<br/><i>model, optimizer, lr_scheduler,<br/>train_loss, val_loss, epoch, scaler</i>"]
    epoch_event --> lr_step

    lr_step["Step LR scheduler<br/><i>ReduceLROnPlateau uses val_loss;<br/>all others use epoch count</i>"]
    lr_step --> stop_check

    stop_check{"StopCondition<br/>satisfied?"}
    stop_check -- Yes --> done
    stop_check -- No --> epoch_loop

What happens at each step

1. TrainingBegan — fired once before the loop starts. PyTorchModelSaver uses it to initialize best_loss from a resumed checkpoint.

2. LossScheduler update — if the criterion is a LossScheduler, it advances to the current epoch (e.g., switching from L1 to a combined L1+SSIM at epoch 50).

3. Training runner — runs one full pass through the training data with gradients enabled. Returns aggregated MetricScores.

4. Validation runner — runs one full pass through the validation data with torch.no_grad(). Returns aggregated MetricScores.

5. on_trainer_epoch_finished — the trainer's hook point fires after both runners complete. Trainer-level hooks react:

   • PyTorchModelSaver saves checkpoint_best.pth if validation loss improved, and always saves checkpoint_last.pth.
   • LossLogger logs all metrics to the tracker and calls tracker.commit().

6. LR scheduler step — the trainer steps the LR scheduler once per epoch. Metric-based schedulers receive the mean validation loss; epoch-based schedulers step by count.

7. Stop condition — checked after each epoch. Returns True to end training early (e.g., when validation loss plateaus). The default NoCondition never stops.

---

Runner lifecycle

Every runner follows the same pattern in run_epoch(). The difference is whether gradients and optimizer steps are involved.

flowchart TD
    start([run_epoch#40;model, data_loader, epoch, criterion#41;]) --> mode

    mode["Set model mode<br/><i>train#40;True#41; or train#40;False#41; + no_grad</i>"]
    mode --> epoch_start

    epoch_start["<b>RunnerEpochStarted</b> event<br/><i>epoch, dataset_size, batch_size, num_batches</i>"]
    epoch_start --> batch_loop

    batch_loop{"Next batch?"}
    batch_loop -- No --> epoch_end
    batch_loop -- Yes --> forward

    forward["<b>Forward pass</b><br/><i>entry.to#40;device#41; → model#40;entry#41;</i><br/>with autocast if mixed_precision"]
    forward --> merge

    merge["Merge output into Entry"]
    merge --> post

    post["Apply postprocessors<br/><i>e.g., clamp, denormalize</i>"]
    post --> loss

    loss["Compute loss<br/><i>batch_scores = criterion#40;entry#41;</i>"]
    loss --> backward

    backward{"Training<br/>runner?"}
    backward -- Yes --> grad
    backward -- No --> accum

    grad["<b>Backward pass</b><br/><i>loss / accumulation_steps → .backward#40;#41;</i>"]
    grad --> optim_check

    optim_check{"Accumulation<br/>step?"}
    optim_check -- Yes --> optim_step
    optim_check -- No --> accum

    optim_step["<b>Optimizer step</b><br/><i>scaler.step#40;optimizer#41;<br/>scaler.update#40;#41;<br/>optimizer.zero_grad#40;#41;</i>"]
    optim_step --> accum

    accum["Accumulate MetricScores"]
    accum --> batch_event

    batch_event["<b>RunnerBatchFinished</b> event<br/><i>epoch, batch, entry, batch_scores,<br/>criterion, epoch_scores</i>"]
    batch_event --> free

    free["Free memory<br/><i>del output, entry, batch_scores</i>"]
    free --> batch_loop

    epoch_end["<b>RunnerEpochFinished</b> event<br/><i>epoch, epoch_scores</i>"]
    epoch_end --> ret([Return epoch_scores])

Key details

• Model mode — TrainingEpochRunner sets model.train(True); validation and benchmark runners set model.train(False) inside torch.no_grad().

• Forward pass — the model receives an Entry and returns an Entry, GraphEntry, or dict. The runner merges the output back into the original entry so postprocessors and the loss see all fields.

• Postprocessors — a list of transforms applied after the model (e.g., clamping pixel values, undoing normalization). Configured per-runner.

• Accumulation step — the optimizer steps every gradient_accumulation_steps batches, or on the last batch of the epoch. The loss is divided by the accumulation factor before .backward().

• Hook points — the per-batch on_post_step context carries the full entry, so hooks like BatchImageSaver can extract and save predictions. on_epoch_end carries the aggregated scores for the entire epoch.

---

Hook timeline

This sequence diagram shows one complete epoch — where each component's hooks fire. TH are hooks attached to the training runner, VH to the validation runner, and H to the trainer:

sequenceDiagram
    participant T as PyTorchTrainer
    participant TR as TrainingRunner
    participant TH as TR's hooks
    participant VR as ValidationRunner
    participant VH as VR's hooks
    participant H as Trainer's hooks

    Note over T: Epoch begins

    T->>TR: run_epoch(model, train_loader, epoch, criterion=training_criterion)
    activate TR
    TR->>TH: on_epoch_start
    Note over TH: ProgressBar "T" initializes

    loop Each training batch
        TR->>TR: forward → postprocess → loss → backward → optimizer step
        TR->>TH: on_post_step
        Note over TH: ProgressBar updates
    end

    TR->>TH: on_epoch_end
    TR-->>T: train_scores
    deactivate TR

    T->>VR: run_epoch(model, val_loader, epoch, criterion=validation_criterion)
    activate VR
    VR->>VH: on_epoch_start

    loop Each validation batch
        VR->>VR: forward → postprocess → loss
        VR->>VH: on_post_step
        Note over VH: BatchImageLogger buffers predictions
    end

    VR->>VH: on_epoch_end
    VR-->>T: val_scores
    deactivate VR

    T->>H: on_trainer_epoch_finished
    Note over H: PyTorchModelSaver saves checkpoint<br/>LossLogger logs metrics

    Note over T: Step LR scheduler<br/>Check stop condition

No scoping — attachment is the scope

There is no scope filter to configure. A hook in the training runner's hooks: list fires only during training epochs; a hook in the validation runner's list fires only during validation; a hook in the trainer's list fires at trainer-level points (on_training_began, on_trainer_epoch_finished). To fire the same hook instance from multiple components, put it in each component's hooks: list.

Runner scope= was removed in 0.16.0

Up to 0.15.x a runner's scope="train" / scope="val" filtered which observers attached. With observers gone it had nothing to filter, so it was removed: a runner given scope refuses it, like any option it does not take.

---

Runner types

EpochRunner (ABC)
├── TrainingEpochRunner    — backward pass, gradient accumulation, optimizer, hooks
├── GANTrainingRunner      — alternating G/D updates, self-balancing step ratio, hooks
└── InferenceRunner        — torch.no_grad() loop, optional criterion
    ├── ValidationEpochRunner
    └── BenchmarkRunner

ValidationEpochRunner and BenchmarkRunner both inherit from InferenceRunner — the same inference loop with no gradients. They differ only in their scope tag and whether criterion is enforced.

	TrainingEpochRunner	GANTrainingRunner	ValidationEpochRunner	BenchmarkRunner
Scope	train	train	val	benchmark
Gradients	Yes	Yes (G + D)	No (torch.no_grad)	No
Optimizer	One	Two (G + D)	No	No
Mixed precision	Optional (one scaler)	Optional (two scalers)	Optional	Optional
Grad accumulation	Optional	No	N/A	N/A
Criterion	Required in run_epoch	d_criterion+g_criterion on constructor	Required in run_epoch	Optional in run_epoch
Hooks	5 base stages	5 base + 6 GAN stages	None	None
Used in	PyTorchTrainer	PyTorchTrainer	PyTorchTrainer	Benchmark scripts

Criterion is not stored on the runner

Runners are stateless with respect to criteria. The criterion is always passed as a parameter to run_epoch(criterion=...). The PyTorchTrainer owns both criteria and passes them at call time. This keeps runners as pure execution engines and avoids temporal coupling.

TrainingEpochRunner — the workhorse. Handles forward pass, backward pass, optimizer steps, gradient accumulation, and mixed-precision scaling. Used as the trainer's training_epoch_runner.

ValidationEpochRunner — same forward pass and loss computation, but with gradients disabled. Raises ValueError if criterion is not provided to run_epoch(). Used as the trainer's validation_epoch_runner.

BenchmarkRunner — designed for test/inference scripts. Criterion is optional — if you only need predictions (e.g., saving images), omit it. If provided, it computes and reports metrics. Used standalone, not inside a trainer.

GANTrainingRunner — alternating G/D updates with self-balancing step ratio. Requires a GANModel (srforge.models.GANModel). Field- agnostic: it never touches Entry fields directly — GANModel handles discriminator scoring and gradient control (detach, torch.no_grad), and every criterion reads from the Entry through the standard srforge.metrics.Metric interface. Constructor takes two optimizers (optimizer_G, optimizer_D), two criteria (d_criterion, g_criterion) and an optional list of hooks — the cadence of G and D updates is the StepRatio hook (ratio: 1.0 updates both every batch; 2.0 trains G twice as often). Uses two GradScalers under mixed precision. For multiple GPUs launch with torchrun (DDP), which wraps the generator and discriminator separately; a device list (DataParallel) runs a GAN on one GPU, because the G/D steps bypass that wrapper.

---

Training hooks

The side-channel mechanism for everything around the training step — logging, checkpointing, gradient clipping, extra loss terms. (It replaced the Observer / EventBus pattern, removed in 0.16.0.)

A Hook is a composable behaviour attached to a runner or the Trainer at construction. Each runner / trainer exposes its lifecycle stages as HookPoint attributes (on_epoch_start, on_post_step, …); the hook's on_<x> methods are auto-bound to matching HookPoints at attach time via Hook.bind_to(target).

Hooks receive a mutable Context and can:

• Skip a step entirely (ctx.run_d = False, ctx.run_g = False in GAN).
• Add an auxiliary regularisation loss: write a MetricEntry into ctx.scores from on_post_forward (the GAN runner uses ctx.d_extra_losses / ctx.g_extra_losses).
• Inspect or mutate gradients between backward and the optimizer step.
• Modify the Entry before forward (e.g. data augmentation).

HookPoints declared by each runner / trainer

Each class exposes its own set, and a hook can only bind to a point the target actually declares. The full table — plus how to check at runtime and what the warning looks like when you attach a hook to the wrong component — lives with the hook documentation: Which HookPoints can I attach to?.

User-defined runners declare new HookPoints with a one-line attribute annotation — see HookPoints in custom runners below.

Worked example

Putting attachment-is-the-scope into practice — each hook goes in the list of the component it should fire from:

PyTorchTrainer(
    model=model,
    training_epoch_runner=TrainingEpochRunner(
        optimizer=optimizer, hooks=[ProgressBar(name="T")]),   # training only
    validation_epoch_runner=ValidationEpochRunner(
        hooks=[ProgressBar(name="V")]),                        # validation only
    training_criterion=loss,
    validation_criterion=loss,
    hooks=[LossLogger(tracker=tracker)],                       # the trainer only
)

Hooks with persistent state implement state_dict() / load_state_dict(state) for checkpoint round-trip.

Built-in hooks

Ten ship with SR-Forge, all @register_class-decorated so they can be referenced by bare name in YAML or by dotted module path. Attach to is the column that matters: a hook only fires from a component that declares its HookPoints (see Which HookPoints can I attach to?).

Hook	Purpose	Binds to	Attach to
ProgressBar	Rich progress bar for one runner, with live metrics.	on_epoch_start, on_post_step, on_epoch_end	any runner
LossLogger	Train/val losses + learning rates to the tracker, once per epoch.	on_trainer_epoch_finished	trainer
PyTorchModelSaver	Best/last checkpoints and model weights.	on_training_began, on_trainer_epoch_finished	trainer
BatchImageSaver	Writes prediction images to the run's output directory.	on_post_step, on_epoch_end	any runner
BatchImageLogger	Logs prediction images to W&B.	on_post_step, on_epoch_end	any runner
GradientClip	Clips gradient norms; unscales under AMP via the runner's scaler first.	on_post_backward, on_post_d_backward, on_post_g_backward	training or GAN runner
StepRatio	Self-balancing G/D update ratio. Sets ctx.run_d / ctx.run_g.	on_pre_step	GAN runner
DWarmup	First N global batches are D-only, until the discriminator has learned something.	on_pre_step	GAN runner
R1GradientPenalty	R1 regularisation — gradient penalty on real images, mean-reduced so the weight is resolution-independent.	on_post_d_forward	GAN runner
EdgeEnhancedInput	Augments the input with a Sobel-style edge map as an extra channel.	on_pre_d_forward, on_pre_g_forward	GAN runner

GradientClip is the one hook that spans both: it declares a method per point, so the same instance clips a plain training runner or a GAN's D and G separately, and only the applicable methods bind.

# Example hooks block for a GAN training_runner:
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 }
    hooks:
      - _target: StepRatio
        params: { ratio: 1.0 }
      - _target: DWarmup
        params: { batches: 500 }
      - _target: R1GradientPenalty
        params: { weight: 0.5 }
      - _target: GradientClip
        params: { max_norm: 1.0 }
    device: ${system.device}
    mixed_precision: ${system.mixed_precision}

Writing a custom hook

from srforge.training.hooks import Hook, hooks_into
from srforge.training.context import Context
from srforge.training.runners import GANTrainingRunner
from srforge.registry import register_class

@register_class
class LogDStepCount(Hook):
    """Track how many discriminator steps fired this epoch.

    Attach it to the GAN training runner's ``hooks:`` list — where a hook
    is attached decides where it fires.
    """

    def __init__(self):
        super().__init__()
        self.d_steps_this_epoch = 0

    @hooks_into("on_epoch_start")
    def reset(self, ctx: Context):
        self.d_steps_this_epoch = 0

    # Type-annotating ctx as the runner's nested Ctx gives IDE
    # autocomplete on the GAN-specific fields (run_d, d_steps, …).
    @hooks_into("on_post_d_backward")
    def count_d_step(self, ctx: GANTrainingRunner.Ctx):
        self.d_steps_this_epoch += 1

    def state_dict(self):
        return {"d_steps_this_epoch": self.d_steps_this_epoch}

    def load_state_dict(self, state):
        self.d_steps_this_epoch = state["d_steps_this_epoch"]

HookPoints in custom runners

To add a new entry point to a custom runner, declare a class-level HookPoint annotation — the base class harvests them via __init_subclass__ and instantiates one per attribute on each instance.

from srforge.training.context import Context
from srforge.training.hookpoint import HookPoint
from srforge.training.runners import EpochRunner

class AdversarialAttackRunner(EpochRunner):
    """Custom runner with attack-specific stages."""

    class Ctx(Context):
        """IDE-typed Context flavour. Hooks targeting this runner can
        annotate ctx: AdversarialAttackRunner.Ctx to get autocomplete
        on the custom fields below."""
        attack_strength: float
        attack_iterations: int

    on_pre_attack:  HookPoint
    on_post_attack: HookPoint

    def run_epoch(self, model, loader, epoch, criterion=None):
        ctx = self.Ctx(
            epoch=epoch, model=model,
            attack_strength=0.1, attack_iterations=10,
        )
        self.on_epoch_start.fire(ctx)
        for i, batch in enumerate(loader):
            ctx.batch_idx, ctx.entry = i, batch
            self.on_pre_attack.fire(ctx)
            ...
            self.on_post_attack.fire(ctx)
        self.on_epoch_end.fire(ctx)

---

Key features

Mixed precision

Enable automatic mixed precision (AMP) per-runner:

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    mixed_precision: true
    # ...

When enabled, the forward pass and loss computation run inside torch.autocast('cuda'). A GradScaler handles loss scaling for the backward pass to prevent underflow in float16 gradients. The scaler state is saved and restored in checkpoints.

Gradient accumulation

Simulate larger batch sizes by accumulating gradients over multiple batches:

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    gradient_accumulation_steps: 4
    # ...

The loss is divided by gradient_accumulation_steps before .backward(). The optimizer steps every 4 batches (or on the last batch of the epoch, whichever comes first). Effective batch size = batch_size * gradient_accumulation_steps.

LR scheduling

The trainer steps the LR scheduler once per epoch, after both runners finish. Every scheduler in torch.optim.lr_scheduler works — epoch-based ones (e.g., StepLR, CosineAnnealingLR, ExponentialLR) and ReduceLROnPlateau, which watches the validation loss; the trainer works out which ones want that loss, so you never pass it yourself. The loss it passes is the epoch's mean weighted validation total — the same number the best model is chosen by.

Whatever you give it, the trainer wraps it in a SchedulerGroup — so trainer.lr_scheduler is always one. Passing nothing gives you an empty group: a schedule that steps nothing, i.e. a constant learning rate.

Several optimizers? A GAN runner owns optimizer_G and optimizer_D, and one scheduler can only drive one of them. Build a SchedulerGroup with one scheduler per optimizer and the trainer steps, checkpoints and logs all of them together:

lr_scheduler:
  _target: srforge.training.schedule.SchedulerGroup
  params:
    schedulers: [${ref:sched_G}, ${ref:sched_D}]

Each scheduler already names its own optimizer (optimizer: ${ref:optimizer_G}), so the group needs no keys of its own. A scheduler wired to an optimizer the runner doesn't own is rejected when the trainer is built, rather than quietly stepping something nothing trains.

Your own schedule. There is nothing to inherit. Any object with three methods works — step(), state_dict() and load_state_dict() — and the last two are what put it in checkpoints, so a resumed run continues the schedule instead of restarting it:

from srforge.registry import register_class

@register_class
class WarmupThenConstant:
    """Linear warm-up over the first `warmup` epochs, then the base LR."""

    def __init__(self, optimizer, warmup: int = 5):
        self.optimizer = optimizer          # lets the trainer check and name it
        self.warmup = warmup
        self.base = [g["lr"] for g in optimizer.param_groups]
        self.epoch = 0
        self._apply()

    def step(self):                         # called once per epoch
        self.epoch += 1
        self._apply()

    def _apply(self):
        scale = min(1.0, (self.epoch + 1) / self.warmup)
        for group, lr in zip(self.optimizer.param_groups, self.base):
            group["lr"] = lr * scale

    def state_dict(self):
        return {"epoch": self.epoch}

    def load_state_dict(self, state):
        self.epoch = state["epoch"]
        self._apply()

Keep the optimizer in self.optimizer: the trainer uses it to reject a schedule wired to an optimizer the runner doesn't train, and to file its checkpoint state under that optimizer's name.

Should your schedule react to the validation loss, like ReduceLROnPlateau does, add wants_metric = True to the class. step() then receives the loss — def step(self, metric): .... Without the flag, step() is called with no arguments.

Watching the learning rate. The W&B logger records each optimizer's current learning rate every epoch — read from the optimizer itself, so a run with no schedule logs it too. The chart names follow the optimizer names:

Runner's optimizer	Chart
optimizer	learning_rate
optimizer_G, optimizer_D (GAN)	learning_rate_G, learning_rate_D
an optimizer with several parameter groups	the name above plus _0, _1, … per group

The full reference for SchedulerGroup and the schedule protocol is in the API reference.

Early stopping

The trainer checks a StopCondition after each epoch. The StopCondition interface receives the current epoch, training loss, and validation loss, and returns True to stop training.

SR-Forge ships two implementations:

• NoCondition (default) — never stops; training runs for all epochs.
• ValidationLossDidNotImprove(patience, min_delta) — stops if validation loss hasn't improved by at least min_delta for patience consecutive epochs.

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    stop_condition:
      _target: srforge.training.stop.ValidationLossDidNotImprove
      params:
        patience: 10
        min_delta: 0.0001
    # ...

You can implement custom stop conditions by subclassing StopCondition and overriding is_condition_satisfied().

---

YAML example

A complete trainer + runners configuration:

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    optimizer: ${ref:optimizer}
    device: ${system.device}
    postprocessor: ${postprocessing}
    mixed_precision: true
    gradient_accumulation_steps: 1

validation_runner:
  _target: srforge.training.runners.ValidationEpochRunner
  params:
    device: ${system.device}
    postprocessor: ${postprocessing}
    mixed_precision: true

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    model: ${ref:model}
    training_epoch_runner: ${ref:training_runner}
    validation_epoch_runner: ${ref:validation_runner}
    training_criterion: ${ref:loss}
    validation_criterion: ${ref:loss}
    lr_scheduler: ${ref:lr_scheduler}
    stop_condition:
      _target: srforge.training.stop.NoCondition

For benchmark scripts, use BenchmarkRunner directly — no trainer needed:

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

Omit criterion for inference-only runs (no scoring):

runner = BenchmarkRunner(device=device, postprocessor=postprocessor)
runner.run_epoch(model=model, data_loader=test_loader, epoch=0)

---

Next: Configuration — Wire everything together in YAML