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 loss criteria 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.

Legacy event scoping

Older versions routed events through a global EventBus with scope="train" / scope="val" filters. That mechanism is deprecated (removal in v0.16.0) — see Observers & Events for the legacy reference.

---

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  (scope="val")
    └── BenchmarkRunner        (scope="benchmark")

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 all losses read from the Entry through the standard srforge.loss.Loss interface. Constructor takes two optimizers (optimizer_G, optimizer_D), two criteria (d_criterion, g_criterion), an optional list of hooks, and step_ratio (G updates per D update — 1.0 updates both every batch; 2.0 trains G twice as often). Uses two GradScalers under mixed precision. Single-GPU only (the G/D step methods bypass DataParallel) — set system.device to a single device when configuring.

---

Training hooks

The canonical side-channel mechanism — replaces the legacy Observer / EventBus pattern (see Observers & Events for the legacy reference).

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 (ctx.extra_losses["r1"] = ...).
• 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

Class	HookPoints
Trainer (PyTorchTrainer)	on_training_began, on_trainer_epoch_finished
TrainingEpochRunner	on_epoch_start, on_pre_step, on_post_backward, on_post_step, on_epoch_end
InferenceRunner (val / benchmark)	on_epoch_start, on_post_step, on_epoch_end
GANTrainingRunner	the five base + on_pre_d_forward, on_post_d_forward, on_post_d_backward, on_pre_g_forward, on_post_g_forward, on_post_g_backward

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

Two attachment surfaces

Trainer(hooks=[…]) distributes each hook to itself and to every runner whose scope matches hook.scope:

hook.scope	Attaches to
None (default)	Trainer + every runner (cross-runner)
"train"	Trainer + the training runner
"val"	Trainer + the validation runner
["train", "val"]	Trainer + listed runners

Runner(hooks=[…]) is the local-to-this-runner form — useful for hooks that only ever fire on one runner (e.g. a runner-specific GradientClip).

PyTorchTrainer(
    model=model,
    training_epoch_runner=train_runner,
    validation_epoch_runner=val_runner,
    training_criterion=loss,
    validation_criterion=loss,
    hooks=[
        LossLogger(tracker=tracker),         # cross-runner + trainer
        ProgressBar(scope="train", name="T"),  # training runner only
        ProgressBar(scope="val", name="V"),    # validation runner only
    ],
)

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

Built-in hooks

All five built-in hooks are @register_class-decorated, so they can be referenced bare-name in YAML or via the dotted module path.

Hook	Purpose	Stages it fires on
StepRatio	Self-balancing G/D update ratio for GAN training. Reads ratio (target G:D updates). Sets ctx.run_d / ctx.run_g in pre_step.	pre_step (GAN)
DWarmup	First N global batches: D-only training. Suppresses G updates until the discriminator has learned something.	pre_step (GAN)
R1GradientPenalty	R1 regularisation — gradient penalty on real images for the discriminator. Mean-reduced so weight is resolution-independent.	post_d_forward (GAN)
GradientClip	Clips gradient norms via torch.nn.utils.clip_grad_norm_. Unscales under AMP via the runner's scaler first.	post_backward (base) or post_d_backward / post_g_backward (GAN)
EdgeEnhancedInput	Augments the input with a Sobel-style edge map as an extra channel.	pre_step (base)

# 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.loss.adversarial.RaGANDiscriminatorLoss
    g_criterion:
      _target: srforge.loss.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 }
    step_ratio: 1.0
    device: ${system.device}
    mixed_precision: ${system.mixed_precision}

Writing a custom hook

from srforge.training.hooks import Hook
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."""

    scope = "train"   # attach to training runners only

    def __init__(self):
        self.d_steps_this_epoch = 0

    def on_epoch_start(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, …).
    def on_post_d_backward(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. Any torch.optim.lr_scheduler.LRScheduler subclass works — epoch-based schedulers (e.g., StepLR, CosineAnnealingLR, ExponentialLR) and metric-based schedulers (e.g., ReduceLROnPlateau) are both supported. If no scheduler is provided, the trainer uses a no-op BlankLRScheduler internally.

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