ContextualHook

class ContextualHook(prehook: str | None = None, posthook: str | None = None, *, prehook_kwargs: dict[str, Any] | None = None, posthook_kwargs: dict[str, Any] | None = None, train_update: bool = True, eval_update: bool = True)[source]

Bases: Hook

Provides forward hook and prehook functionality for subclasses.

This is used to manage references to the ContextualHook in a safe way for the garbage collector (i.e. without cyclic references).

Parameters:
  • prehook (str | None, optional) – name of the prehook method, if any, to execute, no prehook when None. Defaults to None.

  • posthook (str | None, optional) – name of the posthook method, if any, to execute, no posthook when None. Defaults to None.

  • prehook_kwargs (dict[str, Any] | None, optional) – keyword arguments passed to register_forward_pre_hook(). Defaults to None.

  • posthook_kwargs (dict[str, Any] | None, optional) – keyword arguments passed to register_forward_hook(). Defaults to None.

  • train_update (bool, optional) – if the hooks should be run when hooked module is in train mode. Defaults to True.

  • eval_update (bool, optional) – if the hooks should be run when hooked module is in eval mode. Defaults to True.

Raises:

RuntimeError – at least one of prehook and posthook must not be None.

Note

If not None, the signature of the prehook must be of the following form.

hook(module, args) -> None or modified input

Or, if with_kwargs is passed as a keyword argument.

hook(module, args, kwargs) -> None or modified input

See register_forward_pre_hook() for further information.

Note

If not None, the signature of the posthook must be of the following form.

hook(module, args, output) -> None or modified output

Or, if with_kwargs is passed as a keyword argument.

hook(module, args, kwargs, output) -> None or modified output

See register_forward_hook() for further information.