Pytorch Module Hook

• register_forward_hook()

• register_full_backward_hook()

• register_forward_pre_hook()

Forward (post)

from transformers import AutoModel
model = AutoModel.from_pretrained("model-name")
def get_activations(module, input, output):
    print(output)
hook = model.transformer.layer[-1].register_forward_hook(get_activations)
outputs = model(**inputs)
hook.remove()

Forward pre

def __call__(self, module: nn.Module, inputs: tuple[Tensor]) -> Tensor:
  residual: Tensor = inputs[0]

Backward

Modules — PyTorch 2.2 documentation

Building blocks of stateful computation. PyTorch provides a robust library of modules and makes it simple to define new custom modules, allowing for easy construction of elaborate, multi-layer neural networks.

https://pytorch.org/docs/stable/notes/modules.html#module-hooks

torch.nn.modules.module.register_module_forward_hook — PyTorch 2.2 documentation

This adds global state to the nn.module module and it is only intended for debugging/profiling purposes.

https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_forward_hook.html