Skip to content

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.

-

  • [Custom Layers](/06-pytorch/02-module-and-layer-engineering/(01-custom-layers-advanced-containers/)
  • Autograd
  • Custom Autograd