Hooks — Debug, Inject, Extract¶

Overview¶

Hooks give you a seam into a module's forward and backward passes without modifying its forward. They're the Swiss-army knife for: activation/feature extraction, gradient debugging, weight patching, and profiling.

register_forward_hook(fn(module, input, output)) — runs after forward.
register_forward_pre_hook(fn(module, input)) — runs before forward.
register_full_backward_hook(fn(module, grad_input, grad_output)) — runs in backward.
register_parameter_hook / Tensor.register_hook — parameter & tensor-level grads.

💡 Forward hooks with with torch.no_grad() capture activations for analysis without wasting graph memory.

Forward Hooks¶

import torch, torch.nn as nn

def capture(name):
    def hook(module, inp, out):
        print(f"[{name}] in={tuple(i.shape for i in inp)} out={tuple(out.shape)}")
    return hook

model = nn.Sequential(
    nn.Linear(8, 16),
    nn.ReLU(),
    nn.Linear(16, 4),
)
model[0].register_forward_hook(capture("lin1"))
model[2].register_forward_hook(capture("lin3"))
model(torch.randn(2, 8))

⚠️ Hooks are invoked per forward; in training, they also fire during any internal extra forwards (e.g., BatchNorm stats, torch.compile graph replays) — guard with a flag if needed.

Feature Extraction (the "penultimate layer" trick)¶

activations = {}
def make_hook(name):
    def hook(m, inp, out):
        activations[name] = out.detach()   # detach: no graph, just values
    return hook

model = nn.Sequential(nn.Linear(8, 32), nn.ReLU(), nn.Linear(32, 2))
model[0].register_forward_hook(make_hook("feats"))
model[2].register_forward_hook(make_hook("logits"))

x = torch.randn(4, 8)
model(x)
print("feats shape:", activations["feats"].shape)   # (4, 32)

Use case: embeddings for retrieval, visualization, distillation — no need to split your model.

Patching Behavior with Pre-Hooks¶

def inject_noise(module, inp):
    (x,) = inp
    return (x + torch.randn_like(x) * 0.01,)

model[0].register_forward_pre_hook(inject_noise)   # noise before every forward

✅ Pre-hooks may return a replacement input tuple — a clean way to quantize inputs or swap dtype at a single seam.

Backward Hooks (grad inspection)¶

def grad_watch(name):
    def hook(module, grad_input, grad_output):
        print(f"[{name}] grad_in={[None if g is None else tuple(g.shape) for g in grad_input]}")
        print(f"[{name}] grad_out={[None if g is None else tuple(g.shape) for g in grad_output]}")
    return hook

model = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 4))
model[2].register_full_backward_hook(grad_watch("head"))
loss = model(torch.randn(2, 8)).square().mean()
loss.backward()

💡 register_full_backward_hook — use the full variant (the old register_backward_hook is deprecated).

Tensor-Level `register_hook` — nab a leaf's grad¶

w = nn.Parameter(torch.randn(4, 4))
history = []
w.register_hook(lambda g: history.append(g.detach().clone()))

y = (w @ torch.randn(4, 1)).sum()
y.backward()
print("grad captured:", history[0].shape)

Weight Tying & Hooks in `state_dict`¶

Hooks are not serialized. If you attach feature-extraction hooks, re-register them after load_state_dict or keep a builder function.

def attach_feature_models(model):
    def fhook(m, inp, out): return out.detach()
    model.feat_hooks = [m.register_forward_hook(fhook) for m in model.modules()]
    return model

Removing Hooks & Leaks¶

h = model[0].register_forward_hook(lambda *_: None)
model(torch.randn(2, 8))     # hook fires
h.remove()                   # ALWAYS remove in long loops / libraries

⚠️ Libraries that register hooks (like torchinfo, pytorch-lightning summary) can leave hooks unless removed — check model._forward_hooks.

Use Cases Matrix¶

Task	Hook type	Note
Feature/embedding capture	forward	detach output
Weight decay on units only	pre-forward	patch inputs
Gradient clipping debug	full_backward	inspect grad_norm per layer
NaN tracing	forward on loss	`torch.isnan(out).any()`
Distillation student/teacher	forward	run teacher through hook

Key Takeaways¶

Forward hooks = inspect/patch activations; backward hooks = inspect gradients.
Always detach() captured activations; always remove() hooks you own.
Pre-hooks can replace inputs — perfect for injection seams.
Hooks are not serialized: re-attach after load.
For profiling use torch.profiler or torchviz; hooks are for logic, not timing.