SR-Forge

LR Schedules

LR-schedule seam: one interface for zero, one, or many schedulers.

A runner may own more than one optimizer — :class:GANTrainingRunner owns optimizer_G and optimizer_D — but the trainer takes a single lr_scheduler. :class:SchedulerGroup is the composite that closes that gap: the trainer always holds one of these, whether the run configured no schedule, one, or one per optimizer.

See docs/design/lr-schedule-group-rfc.md for the full rationale.

LRSchedule

Bases: Protocol

What the trainer, checkpointer and LR logger need from a schedule.

Structural, not nominal: every torch.optim.lr_scheduler satisfies it without inheriting anything from SR-Forge — which is the point, since inheriting :class:~torch.optim.lr_scheduler.LRScheduler forces a subclass to supply an optimizer it may not have.

Static-only, deliberately: nothing isinstance-checks against this, so it is not @runtime_checkable. A runtime check would compare method names and not signatures, which is exactly the false assurance a schedule protocol must not give.

The protocol carries no behaviour; the method bodies are ... and it is never instantiated. All runtime behaviour lives in :class:SchedulerGroup.

LR reporting is deliberately absent: a learning rate belongs to the optimizer, not to the thing that mutates it. The LR loggers read ctx.optimizers.

SchedulerGroup

Zero or more LR schedules over a runner's optimizers, driven as one.

An empty group is the constant-LR case — step() iterating nothing is the whole of a null schedule's behaviour, so no separate null class is needed.

Reporting learning rates is deliberately not here. A learning rate belongs to the optimizer; a schedule is only what mutates it. The LR loggers read ctx.optimizers directly (see :func:srforge.training.hooks.learning_rate_metrics). The optimizers this class does hold are for validating its members against the runner and naming their checkpoint state.

Parameters:

Name	Type	Description	Default
schedulers	Iterable[LRSchedule] | None	Schedules to drive. Each is expected to satisfy :class:LRSchedule and to carry the optimizer it was built on. None means a constant learning rate.	None
optimizers	Mapping[str, Optimizer] | None	The runner's optimizers, keyed as :meth:~srforge.training.runners.EpochRunner.optimizers keys them ("optimizer", or "optimizer_G" / "optimizer_D"). Usually supplied by :meth:of rather than by hand.	None

of(value: LRSchedule | SchedulerGroup | None, runner: EpochRunner | None = None) -> SchedulerGroup classmethod

Return value as a group, whatever shape it arrived in.

A run can express its LR schedule three different ways, but the trainer wants one type to hold and step. This is the single place that converts:

• None — no schedule configured — becomes an empty group: nothing to step, so the learning rate stays constant.
• a single scheduler becomes a group of one.
• a :class:SchedulerGroup, which is what a run with several optimizers builds, is returned unchanged.

Passing runner additionally binds its optimizers to the group. Binding is what lets checkpoint state be keyed by optimizer name rather than list position, and what catches a scheduler wired to an optimizer this runner does not train. A group that already carries optimizers keeps the ones it has.

Example — the trainer's entire scheduler setup is this one line::

self.lr_scheduler = SchedulerGroup.of(lr_scheduler, training_runner)

Parameters:

Name	Type	Description	Default
value	LRSchedule | SchedulerGroup | None	Whatever the caller configured: None, one torch.optim.lr_scheduler, or a group they assembled for a multi-optimizer run.	required
runner	EpochRunner | None	Training runner to bind against; only its optimizers() is read. None leaves the group unbound, meaning no validation and positional checkpoint keys.	None

Returns:

Type	Description
SchedulerGroup	A group — never None, so callers need no further checks.

Raises:

Type	Description
ValueError	A scheduler in value drives an optimizer that runner does not own.

step(metric: float | Tensor | None = None) -> None

Step every member, passing metric only to those that want it.

The routing is not cosmetic: LRScheduler.step takes epoch and ReduceLROnPlateau.step takes metrics, so handing a validation loss to the wrong one does not raise — it lands in the epoch slot and silently corrupts last_epoch.

state_dict() -> dict

Member states keyed by optimizer name.

Keyed by name rather than by position so that reordering the schedulers: list in a config between runs cannot silently swap G's and D's state on resume.

load_state_dict(state_dict: dict) -> None

Restore member states, accepting pre-group checkpoints.

A checkpoint written before this class existed holds one torch scheduler's state directly, with no "schedulers" key. It is routed to the single member; a group that has since gained or lost schedulers warns and starts them fresh rather than failing the resume.