Hook

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

Bases: object

Provides and manages forward hook and prehook functionality.

Hook provides functionality to register and deregister itself as forward hook with a Module object. This is performed using register_forward_hook() to register itself as a forward hook and it manages the returned RemovableHandle to deregister itself.

Parameters:

• prehook (Callable | None, optional) – function to call before hooked module’s forward(). Defaults to None.

• posthook (Callable | None, optional) – function to call after hooked module’s forward(). 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.

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.

Raises:

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

deregister() → None

Deregisters the hook as a forward hook and/or prehook.

If the Hook is not registered, this is still safe to call.

property evalexec: bool

If the hook is called when the module passed in is in evaluation mode.

Parameters:

value (bool) – if the hook should be called when the module is evaluating.

Returns:

if the hook is called when the module is evaluating.

Return type:

bool

register(module: Module) → None

Registers the hook as a forward hook and/or prehook.

Parameters:

module (Module) – PyTorch module to which the forward hook will be registered.

Raises:

TypeError – parameter module must be an instance of Module.

Warns:

RuntimeWarning – each Hook can only be registered to one Module and will ignore register() if already registered.

property registered: bool

If there is a module to which this hook is registered

Returns:

if a module to which this hook is registred.

Return type:

bool

property trainexec: bool

If the hook is called when the module passed in is in training mode.

Parameters:

value (bool) – if the hook should be called when the module is training.

Returns:

if the hook is called when the module is training.

Return type:

bool