HookBase

class HookBase

Bases: object

Base class for all training hooks registered with the trainer.

A hook plugs custom behaviour into fixed points of the training loop. The trainer calls each lifecycle method on every registered hook, in registration order, at the corresponding point in the run. Subclasses override only the methods they need; the defaults are no-ops. Every hook in pimm derives from this class.

The lifecycle methods, in the order the trainer invokes them, are:

• modify_config(cfg): called once after all hooks are registered but before the writer (W&B/TensorBoard) is built, so a hook may still mutate the config (e.g. set cfg.wandb_run_name).

• before_train(): called once before the training loop starts.

• before_epoch(): called at the start of every epoch.

• before_step(): called before each training step (before the batch is consumed by the model).

• after_step(): called after each training step (after forward/backward).

• after_epoch(): called at the end of every epoch.

• after_train(): called once after the training loop finishes.

Note

self.trainer is a weak reference (proxy) to the owning trainer, assigned by the trainer at registration time. Through it a hook reaches shared state such as trainer.model, trainer.optimizer, trainer.storage, trainer.writer, trainer.logger, trainer.comm_info, and trainer.cfg. In distributed runs every rank runs every hook, so guard rank-0-only side effects with pimm.utils.comm.is_main_process() and keep any collective operations consistent across ranks.

after_epoch()

after_step()

after_train()

before_epoch()

before_step()

before_train()

modify_config(cfg)

Called after hooks are registered but before writer is built. Hooks can modify the config here (e.g., wandb_run_name).

trainer = None