TopolearnHooks power feature extraction, activation statistics, gradient debugging and interpretability tools.
A hook is a callback PyTorch runs at a specific point. Registering one returns a handle; call handle.remove() when done, or the hook keeps running (and keeps references alive).
module.register_forward_hook(fn)callsfn(module, inputs, output)afterforward. Returning a value replaces the output. Use it to capture intermediate activations.module.register_forward_pre_hook(fn)runs beforeforwardand can modify inputs.module.register_full_backward_hook(fn)callsfn(module, grad_input, grad_output)during backward.tensor.register_hook(fn)callsfn(grad)when that tensor's gradient is computed; returning a tensor replaces the gradient. Handy for gradient clipping on one tensor or for inspecting non-leaf gradients.
Typical uses:
- Extract embeddings from a pretrained network's penultimate layer.
- Log activation means and standard deviations per layer to diagnose dead ReLUs or exploding values.
- Build saliency and Grad-CAM style visualisations from activations plus gradients.
Store output.detach() in forward hooks unless you need gradients through it; otherwise you keep the graph alive. Hooks only fire when the module is called via module(x), not module.forward(x). For torchvision models, torchvision.models.feature_extraction.create_feature_extractor is a cleaner alternative for pulling out named layers.
Example
import torch
from torch import nn
model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
acts = {}
def save_output(name):
def hook(module, inputs, output):
acts[name] = output.detach()
return hook
handle = model[1].register_forward_hook(save_output("relu"))
model(torch.randn(3, 4))
print(acts["relu"].shape) # torch.Size([3, 8])
print((acts["relu"] == 0).float().mean()) # fraction of dead units in this batch
handle.remove() # stop capturing
x = torch.randn(3, 4, requires_grad=True)
x.register_hook(lambda g: print("input grad shape", g.shape))
model(x).sum().backward() # input grad shape torch.Size([3, 4])Remember
- Forward hooks see (module, inputs, output); returning a value replaces output
tensor.register_hooksees and can rewrite gradients- Keep the handle and
remove()it
Watch out for
- Saving
outputwithout.detach(), which keeps the whole autograd graph in memory.