SR-Forge

Runners

Individual runners — classes that execute one epoch of a training process.

Each runner declares its lifecycle stages as :class:HookPoint attributes (on_epoch_start, on_post_step etc.). The :class:EpochRunner base class harvests these declarations from subclasses (via :meth:__init_subclass__) and instantiates them in :meth:__init__, then auto-binds every attached :class:Hook to matching points.

To add a new entry point to a custom runner, declare an attribute::

class MyAdversarialRunner(EpochRunner):
    class Ctx(Context):
        attack_strength: float
    on_pre_attack:  HookPoint
    on_post_attack: HookPoint

    def run_epoch(self, ...):
        ctx = self.Ctx(epoch=..., attack_strength=0.1, ...)
        self.on_pre_attack.fire(ctx)

EpochRunner

Bases: ABC

Base class defining the interface for running one pass through a dataset.

Subclasses declare lifecycle stages as :class:HookPoint-typed class attributes; the base class introspects them at subclass creation and instantiates one :class:HookPoint per attribute on each instance. Attached :class:Hook s are auto-bound at construction.

training_state() -> dict

Return all optimizer/scaler state needed for checkpoint resume.

Subclasses override to include their specific training state. Non-training runners (validation, benchmark) return empty dict.

load_training_state(state: dict) -> None

Restore optimizer/scaler state from a checkpoint dict.

Subclasses override to restore their specific training state.

optimizers() -> dict

Return this runner's optimizers keyed by config-aligned name.

The keys match the config nodes that build them ("optimizer" for a standard runner; "optimizer_G" / "optimizer_D" for a GAN runner), so callers can act on every optimizer a runner owns without an if gan branch, and name what they find after the config the user wrote. The trainer uses it twice: to bind a :class:~srforge.training.schedule.SchedulerGroup to the optimizers it may schedule, and to publish ctx.optimizers for LR logging. Non-training runners return an empty dict.

zero_grad() -> None

Zero gradients on all optimizers managed by this runner.

Subclasses override for their specific optimizers.

run_epoch(model: nn.Module, data_loader: data.DataLoader, epoch: int) -> MetricScores abstractmethod

:param model: Model to run epoch on. :param data_loader: Loader return batches of tuples (input, label) :param epoch: Current epoch index. :return: Epoch loss

TrainingEpochRunner

Bases: EpochRunner

Runner that executes one training epoch with backward pass and optimizer steps.

The criterion is not stored on the runner — it is passed to :meth:run_epoch at call time by the :class:PyTorchTrainer.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Optimizer instance (e.g. torch.optim.AdamW).	required
device	Union[device, str]	Device to move entries to before the forward pass (e.g. "cpu", "cuda:0", or a torch.device).	'cpu'
postprocessor	List	List of callables applied to the entry after the forward pass and before loss computation (e.g. clamping, denormalization). None means no postprocessing.	None
mixed_precision	Optional[bool]	Legacy AMP toggle. True -> fp16 autocast + enabled GradScaler; False -> fp32. Left None (the default) it defers to precision / the global policy. Kept for backward compatibility with existing configs.	None
precision	Optional[Union[PrecisionPolicy, str, Mapping]]	Preferred precision selector — a :class:~srforge.training.precision.PrecisionPolicy, a mode string ("fp16"/"bf16"/"fp32"), or a mapping ({mode, tf32}). Overrides mixed_precision for this runner; when unset the runner uses the global system.precision policy.	None
gradient_accumulation_steps	int	Number of batches over which to accumulate gradients before an optimizer step. The loss is divided by this value before .backward(). Effective batch size = batch_size * gradient_accumulation_steps.	1

InferenceRunner

Bases: EpochRunner

Base for non-training runners (no backward pass).

Runs the model under torch.no_grad() and optionally scores each batch with a criterion. The criterion is passed to :meth:run_epoch at call time, not stored on the runner.

Parameters:

Name	Type	Description	Default
device	Union[str, device]	Device to move entries to before the forward pass (e.g. "cpu", "cuda:0", or a torch.device).	'cpu'
postprocessor	List	List of callables applied to the entry after the forward pass (e.g. clamping, denormalization). None means no postprocessing.	None
mixed_precision	Optional[bool]	Legacy AMP toggle (True fp16 / False fp32); None defers to precision / the global policy.	None
precision	Optional[Union[PrecisionPolicy, str, Mapping]]	Preferred precision selector (PrecisionPolicy, mode string, or {mode, tf32} mapping). Overrides mixed_precision.	None

ValidationEpochRunner

Bases: InferenceRunner

Inference runner for validation.

Identical to :class:InferenceRunner but enforces that a criterion is passed to :meth:run_epoch (raises ValueError otherwise).

Constructor accepts the same parameters as :class:InferenceRunner.

BenchmarkRunner

Bases: InferenceRunner

Inference runner for benchmarking.

Criterion is optional — omit it for inference-only runs where only model outputs are needed (e.g. saving images).

Constructor accepts the same parameters as :class:InferenceRunner.

GANTrainingRunner

Bases: EpochRunner

Training runner for GAN with alternating generator/discriminator updates.

Requires a :class:~srforge.models.GANModel (checked at runtime). The runner is field-agnostic — it never accesses Entry fields directly. :class:GANModel handles discriminator scoring and gradient control via :meth:~srforge.models.GANModel.discriminator_step / :meth:~srforge.models.GANModel.generator_step, and all losses read from the Entry through the standard :class:~srforge.metrics.Metric interface.

Hooks declared with GAN-specific on_<x> method names attach to the corresponding HookPoints below. Hooks needing typed IDE access to the runner-specific Context fields (run_d, run_g, d_extra_losses, …) type-annotate the parameter as GANTrainingRunner.Ctx.

Parameters:

Name	Type	Description	Default
optimizer_G	Optimizer	Generator optimizer.	required
optimizer_D	Optimizer	Discriminator optimizer.	required
d_criterion	Metric	Discriminator adversarial loss (e.g. :class:~srforge.metrics.adversarial.RaGANDiscriminatorLoss).	required
g_criterion	Metric	Generator adversarial loss (e.g. :class:~srforge.metrics.adversarial.RaGANGeneratorLoss). Its weight controls the adversarial contribution.	required
device	Union[device, str]	Device for computation.	'cpu'
postprocessor	List	Transforms applied after G forward, before loss.	None
mixed_precision	Optional[bool]	Legacy AMP toggle (True fp16 / False fp32); None defers to precision / the global policy. Drives both the G and D scalers.	None
precision	Optional[Union[PrecisionPolicy, str, Mapping]]	Preferred precision selector (PrecisionPolicy, mode string, or {mode, tf32} mapping). Overrides mixed_precision.	None
empty_cache_every	int	Clear CUDA cache every N batches (0 = disabled).	0

Ctx

Bases: Context

GAN-specific Context flavour. Provides IDE autocomplete for the GAN fields the runner sets on the Context.