Skip to content
SR-Forge

Hooks

Hooks are the way to inject behaviour into the training loop — logging, checkpointing, gradient clipping, progress bars, custom loss terms. They replaced the Observer / EventBus system, which was removed in 0.16.0. This page covers what hooks are, how to write one, and — if you are upgrading from 0.15.x — how to turn an old Observer into a Hook.

For where hooks plug into the runner / trainer machinery, see Trainers & Runners.


The model in one paragraph

A Hook is a class with methods. Each method declares (via the @hooks_into("on_<x>") decorator) which HookPoint on the runner / trainer it should fire at. At construction the trainer walks the hook, asks each method "where do you bind?", and registers it with the matching HookPoint. When the runner reaches that stage and calls HookPoint.fire(ctx), every registered method runs with the mutable Context.

That's it. No event bus, no scope tag at construction time, no name-derived event lookup.


A first hook

from srforge.training.hooks import Hook, hooks_into
from srforge.training.context import Context


class EpochBanner(Hook):
    """Print a banner at the start of every training epoch."""

    @hooks_into("on_epoch_start")
    def print_banner(self, ctx: Context) -> None:
        print(f"▶ epoch {ctx.epoch}")

Wire it up:

trainer = PyTorchTrainer(
    ...,
    hooks=[EpochBanner()],
)

Or in YAML:

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    hooks:
      - _target: my_module.EpochBanner

Which HookPoints can I attach to?

A HookPoint only exists on the component that declares it, so the name you put in @hooks_into(...) has to match one the target exposes. These are all of them:

Component HookPoints
PyTorchTrainer on_training_began, on_trainer_epoch_finished
TrainingEpochRunner on_epoch_start, on_pre_step, on_post_forward, on_post_backward, on_post_step, on_epoch_end
InferenceRunner — and ValidationEpochRunner / BenchmarkRunner on_epoch_start, on_post_step, on_epoch_end
GANTrainingRunner on_epoch_start, on_pre_step, on_pre_d_forward, on_post_d_forward, on_post_d_backward, on_pre_g_forward, on_post_g_forward, on_post_g_backward, on_post_step, on_epoch_end

Two things that trip people up:

  • The trainer's two points are trainer-only. on_trainer_epoch_finished fires once per epoch after both runners finish and after cross-rank metric reduction — it is where per-epoch logging and checkpointing belong. No runner has it.
  • GANTrainingRunner is not TrainingEpochRunner plus extras. It declares its own set: there is no on_post_forward or on_post_backward on a GAN runner, because forward/backward happen separately for D and G.

Checking at runtime

You do not have to trust this table. Ask the class what it declares:

>>> from srforge.training.runners import TrainingEpochRunner
>>> TrainingEpochRunner._hookpoint_names
['on_epoch_start', 'on_pre_step', 'on_post_forward', 'on_post_backward',
 'on_post_step', 'on_epoch_end']

and ask a hook where it will bind:

>>> from srforge.training.hooks import LossLogger
>>> LossLogger(tracker=None).list_bindings()
[('flush_epoch_metrics', 'on_trainer_epoch_finished')]     # → attach to the trainer

And you are told when you get it wrong. LossLogger binds to on_trainer_epoch_finished, which no runner declares — so it belongs to the trainer and will never fire from a runner:

# WRONG — a runner has no on_trainer_epoch_finished, so this never fires
ValidationEpochRunner(device="cpu", hooks=[LossLogger(tracker=tracker)])

That does not fail silently. It warns at construction, before training starts:

Hook LossLogger declares handler method(s) (['flush_epoch_metrics']) but
ValidationEpochRunner exposes no matching HookPoint. Was this hook attached
to the wrong target?
# RIGHT — trainer-level hook, attached to the trainer
PyTorchTrainer(..., hooks=[LossLogger(tracker=tracker)])

A misspelled name in @hooks_into("on_post_stpe") warns the same way, so a hook that never fires announces itself rather than leaving you to wonder.

Adding your own

A custom runner declares new points with a one-line annotation, and they work exactly like the built-in ones:

class MyAdversarialRunner(EpochRunner):
    on_pre_attack:  HookPoint
    on_post_attack: HookPoint

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

The base class harvests the annotations at subclass creation and instantiates one HookPoint per name on each instance — see Trainers & Runners.

The three binding sources

Hook.bind_to resolves each method's target HookPoint from three sources, in this order:

  1. self.bindings — instance- or class-level dict {method_name: target_point_name | [target_point_names] | None}. Set via constructor kwarg, class attribute, or YAML. None / [] explicitly disables a default binding.
  2. @hooks_into("on_<x>", ...) — decorator metadata on the method itself. Stack the decorator or pass multiple names to bind one method to several points.
  3. Nothing — a method without a decorator and without a bindings entry is not a handler. Naming a method on_<x> is not enough; SR-Forge does not introspect names.

Multiple targets

Bind one method to several HookPoints by passing the names together. The usual reason is a hook that should work on more than one kind of runner: on_post_backward exists on TrainingEpochRunner, while a GAN runner splits the same moment into on_post_d_backward and on_post_g_backward.

class GradientNormLogger(Hook):
    @hooks_into("on_post_backward", "on_post_d_backward", "on_post_g_backward")
    def log_grad_norm(self, ctx):
        ...

Attached to a TrainingEpochRunner only the first binds; on a GANTrainingRunner only the other two do. Names that the target does not declare are simply not bound — the warning fires only when none of a method's targets match.

Or stack decorators — equivalent:

    @hooks_into("on_post_backward")
    @hooks_into("on_post_d_backward")
    def log_grad_norm(self, ctx):
        ...

Remapping at the call site

bindings overrides where a hook's methods attach, without subclassing it. GradientClip.clip_all_grads ships on on_post_backward; to clip only the discriminator on a GAN runner, point it at on_post_d_backward instead:

from srforge.training.hooks import GradientClip

clip = GradientClip(
    max_norm=1.0,
    bindings={"clip_all_grads": "on_post_d_backward"},
)

In YAML, the same thing:

hooks:
  - _target: srforge.training.hooks.GradientClip
    params:
      max_norm: 1.0
      bindings:
        clip_all_grads: on_post_d_backward

This is also how you attach an existing hook to a HookPoint your own runner declares — the name on the right just has to exist on the target.

Disabling a default binding

Pass None to suppress a method's default binding without subclassing:

hooks:
  - _target: my_module.CheckpointSaver
    params:
      path: ${ref:out}/best.pt
      bindings:
        restore_best_loss: null     # skip the restore step

Attaching hooks: directly, not via cascade

Pass each hook to the component it should fire from. There is no distribution layer — a hook attached to a runner fires from that runner's HookPoints; a hook attached to the trainer fires from the trainer's HookPoints; nothing more.

train_runner = TrainingEpochRunner(
    optimizer=opt,
    hooks=[ProgressBar(name="T"), GradientClip(max_norm=1.0)],
)
val_runner = ValidationEpochRunner(
    hooks=[ProgressBar(name="V")],
)
trainer = PyTorchTrainer(
    training_epoch_runner=train_runner,
    validation_epoch_runner=val_runner,
    hooks=[LossLogger(tracker=tracker)],   # fires only from the Trainer
    ...,
)

If you need the same hook instance to fire from several components, put the same object in each hooks=[...] list — hook instances are plain Python objects, nothing prevents reuse:

progress = ProgressBar(name="P")
train_runner = TrainingEpochRunner(..., hooks=[progress])
val_runner   = ValidationEpochRunner(hooks=[progress])

In YAML, declare the hook once and point each component at the built instance with ${ref:...} — SR-Forge's object reference, the same mechanism as ${ref:model} or ${ref:optimizer}:

hooks:
  progress:
    _target: srforge.training.hooks.ProgressBar
    params: {name: "P"}

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    hooks:
      - ${ref:hooks.progress}

validation_runner:
  _target: srforge.training.runners.ValidationEpochRunner
  params:
    hooks:
      - ${ref:hooks.progress}

Persistent state

If a hook carries state that should survive a checkpoint reload, override state_dict() and load_state_dict():

class StepCounter(Hook):
    def __init__(self):
        self.steps = 0

    @hooks_into("on_post_step")
    def tick(self, ctx):
        self.steps += 1

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

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

The trainer collects each hook's state_dict into the overall training state.


Reaching run-wide state

Most of what a hook needs arrives on the Context, and anything else is best taken as a constructor parameter — tracker: ${ref:tracker} in YAML, rather than reached for at runtime. It is explicit, testable, and lets one script drive two trackers.

The output directory — the one fact about the run most hooks that write files need — arrives on the trainer's Context:

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

It comes from srforge.init()

ctx.output_directory is the directory init() set and created, or the trainer's output_directory= argument. With neither it is None, so a hook that writes files should say so clearly — PyTorchModelSaver raises, naming both fixes. A runner-level hook can read GlobalSettings().output_directory, which raises if init() was not called. Full table: GlobalSettings.


Upgrading from observers (removed in 0.16.0)

srforge.observers and srforge.events were removed in 0.16.0. If you wrote your own Observer on 0.15.x, it has to become a Hook; the built-in ones already exist as hooks under srforge.training.hooks. The migration is mechanical. Most hooks land in 5–10 minutes.

Side-by-side example

Before — Observer (0.15.x; these imports no longer exist):

import srforge.events.trainer
from srforge.observers.base import Observer


class LossLogger(Observer):
    EVENTS = [srforge.events.trainer.TrainerEpochFinished]

    def __init__(self, tracker=None, **kwargs):
        super().__init__(**kwargs)              # forwards `scope` kwarg
        self.tracker = tracker

    def on_epoch_finished(
        self,
        event: srforge.events.trainer.TrainerEpochFinished,
    ) -> None:
        self.tracker.log_metrics(
            event.train_loss.mean_weighted(), step=event.epoch
        )

After — Hook:

from srforge.training.hooks import Hook, hooks_into
from srforge.training.context import Context


class LossLogger(Hook):
    def __init__(self, tracker=None):
        self.tracker = tracker

    @hooks_into("on_trainer_epoch_finished")
    def flush_epoch_metrics(self, ctx: Context) -> None:
        self.tracker.log_metrics(
            ctx.train_loss.mean_weighted(), step=ctx.epoch
        )

Step-by-step

  1. Replace the base class. class Foo(Observer) → class Foo(Hook). Drop the EVENTS = [...] line.
  2. Drop the super().__init__(**kwargs) boilerplate unless you have your own kwargs to forward. Hook.__init__ only accepts bindings=.
  3. Rename methods to describe what they do, not which event they handle (on_epoch_finished → flush_epoch_metrics, save_checkpoint, update_ema, …). The method name is documentation now; the binding is metadata.
  4. Add @hooks_into("on_<x>") above each handler. Map old on_<event_name> methods to the corresponding HookPoint name:
  5. on_runner_epoch_started → @hooks_into("on_epoch_start")
  6. on_runner_batch_finished → @hooks_into("on_post_step")
  7. on_runner_epoch_finished → @hooks_into("on_epoch_end")
  8. on_training_began → @hooks_into("on_training_began") (same)
  9. on_trainer_epoch_finished → @hooks_into("on_trainer_epoch_finished") (same)
  10. Swap event: SomeEvent for ctx: Context in the signature. The fields you read off the event almost all live on ctx — event.epoch becomes ctx.epoch, event.train_loss becomes ctx.train_loss, and so on. If a field is missing, check srforge/training/context.py.
  11. Drop the EventBus plumbing. No more bus.subscribe(observer); pass the hook to Trainer(hooks=[...]) or EpochRunner(hooks=[...]) directly — to the specific component it should fire from.
  12. If the old observer used scope="train" etc., simply pass the hook to that runner's hooks=[...] and don't pass it to the others. There is no scope filter on Hook — direct attachment is the contract.

Porting the config

The class is only half of it. A config with an observers: block moves to hooks: lists — and this is where the one genuine judgement call lives, because scope: does not survive the trip.

Before (0.15.x) — one flat list, filtered at delivery by scope::

observers:
  - _target: srforge.observers.ProgressBar
    params: {name: "T", scope: train}     # only training-runner events
  - _target: srforge.observers.ProgressBar
    params: {name: "V", scope: val}       # only validation-runner events
  - _target: srforge.observers.LossLogger  # no scope = every event
    params: {tracker: ${ref:tracker}}
  - _target: srforge.observers.PyTorchModelSaver
    params: {tracker: ${ref:tracker}}

After — no top-level block at all. Each hook goes into the hooks: list of the thing it should fire from:

training_runner:
  _target: srforge.training.runners.TrainingEpochRunner
  params:
    optimizer: ${ref:optimizer}
    hooks:
      - {_target: srforge.training.hooks.ProgressBar, params: {name: "T"}}

validation_runner:
  _target: srforge.training.runners.ValidationEpochRunner
  params:
    hooks:
      - {_target: srforge.training.hooks.ProgressBar, params: {name: "V"}}

trainer:
  _target: srforge.training.trainers.PyTorchTrainer
  params:
    # …model, runners, criteria, lr_scheduler…
    hooks:
      - {_target: srforge.training.hooks.LossLogger,        params: {tracker: ${ref:tracker}}}
      - {_target: srforge.training.hooks.PyTorchModelSaver, params: {tracker: ${ref:tracker}}}

Three rules cover every case:

Old New
scope: train put it in training_runner.params.hooks
scope: val put it in validation_runner.params.hooks
no scope: (fires on everything) decide where it belongs — see below

The no-scope case is the one to think about. An unscoped observer received events from the training runner, the validation runner and the trainer, and most of them only ever acted on one of those. LossLogger and PyTorchModelSaver handle on_trainer_epoch_finished, so they belong to the trainer, not to a runner.

You do not have to guess, and you will not fail silently: attaching a hook somewhere with no matching HookPoint warns at construction —

Hook LossLogger declares handler method(s) (['flush_epoch_metrics']) but
ValidationEpochRunner exposes no matching HookPoint. Was this hook attached
to the wrong target?

— so a misplaced hook announces itself before training starts rather than quietly never running. To check up front, ask the hook where it binds:

>>> from srforge.training.hooks import LossLogger
>>> LossLogger(tracker=None).list_bindings()
[('flush_epoch_metrics', 'on_trainer_epoch_finished')]     # → trainer

Attachment is the scope

There is no scope: parameter on a Hook, and adding one back would be a step backwards: with observers you had to read the filter to know what ran where, and a typo in scope: silently delivered to everything. Now the list a hook sits in is the answer, visible in the config.

One hook instance may appear in more than one list — ${ref:my_hook} in both runners attaches the same object to both, which is fine because the training and validation runners run sequentially. Delete the observers: key when you are done: nothing reads it, and the generated train.py no longer resolves it.

Field mapping cheat-sheet

Event field Context attribute
event.epoch ctx.epoch
event.batch ctx.batch_idx
event.entry ctx.entry
event.batch_scores ctx.scores
event.epoch_scores ctx.epoch_scores
event.train_loss ctx.train_loss
event.val_loss ctx.val_loss
event.lr_scheduler ctx.lr_scheduler
event.runner_state ctx.runner_state
event.total_epochs ctx.total_epochs
event.initial_epoch ctx.initial_epoch
event.best_losses ctx.best_losses

Inspecting a hook's bindings

Call Hook.list_bindings() to see the resolved (method, target_point) pairs — useful when porting and when something doesn't fire:

>>> LossLogger(tracker=...).list_bindings()
[('flush_epoch_metrics', 'on_trainer_epoch_finished')]

Where to go next

  • Trainers & Runners — the full lifecycle, including which HookPoints each runner exposes.
  • API reference: srforge.training.hooks (built-in hooks), srforge.training.hookpoint (the HookPoint / Handle primitives).