fn
register_module_forward_hook
→RemovableHandleregister_module_forward_hook(hook: Callable[..., object], with_kwargs: bool = False, always_call: bool = False)Register a global post-forward hook fired after every Module.__call__.
The hook runs after forward returns and may inspect or replace
the output. Common uses: activation logging / visualisation, output
transformation (e.g. moving outputs to a different device), or
side-effect bookkeeping.
Parameters
hookCallableThe post-hook function. Signature:
hook(module, inputs, output) -> None | new_output(default)hook(module, inputs, kwargs, output) -> None | new_output(with_kwargs=True)
None leaves the output unchanged; returning any
other value replaces the forward output.with_kwargsbool= FalseIf
True, the hook also receives the kwargs dict that was
passed to forward.always_callbool= FalseIf
True, the hook fires even when forward raises an
exception (the output argument is then None). Useful for
cleanup-style hooks.Returns
RemovableHandleHandle for later deregistration.
Notes
Global forward hooks fire for every Module call. When multiple hooks are registered they run in registration order and each may see the previous hook's replacement output — a chain of transforms.
Examples
>>> from lucid.nn.hooks import register_module_forward_hook
>>> activations = {}
>>> def capture(mod, inputs, output):
... activations[type(mod).__name__] = output.detach()
>>> with register_module_forward_hook(capture):
... _ = model(x)
>>> activations.keys()
dict_keys([...])