PyTorch Hooks

Hooks are callback functions you attach to PyTorch modules or tensors. They fire automatically at specific points during the forward or backward pass. Hooks let you inspect activations, gradients, and layer outputs without modifying your model's architecture or training code — essential for debugging, visualization, and research.

Why Hooks Are Useful

Without hooks:
  To inspect layer outputs, you must rewrite model.forward()
  To inspect gradients, you add temporary print statements everywhere
  → invasive, messy, easy to forget to remove

With hooks:
  Attach a function to any layer → it fires automatically
  Collect outputs, gradients, statistics
  Remove hook when done → zero impact on model

Three Types of Hooks

1. Forward hook      → fires after a module's forward() runs
                       receives: module, input, output
                       use for: inspecting activations

2. Forward pre-hook  → fires BEFORE a module's forward() runs
                       receives: module, input
                       use for: modifying input before it enters a layer

3. Backward hook     → fires during the backward pass
                       receives: module, grad_input, grad_output
                       use for: inspecting or modifying gradients

Forward Hook — Capturing Layer Outputs

import torch
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(4, 8),
    nn.ReLU(),
    nn.Linear(8, 2)
)

activations = {}

def capture_output(name):
    def hook(module, input, output):
        activations[name] = output.detach()
    return hook

# Attach hook to the first Linear layer
handle = model[0].register_forward_hook(capture_output('layer0'))

x   = torch.rand(3, 4)
out = model(x)

print(activations['layer0'].shape)   # torch.Size([3, 8])
print(activations['layer0'])

handle.remove()   # always remove when done

Capturing All Layer Outputs

layer_outputs = {}

def make_hook(name):
    def hook(module, input, output):
        layer_outputs[name] = output.detach()
    return hook

handles = []
for name, module in model.named_modules():
    if len(name) > 0:   # skip the top-level container
        h = module.register_forward_hook(make_hook(name))
        handles.append(h)

x = torch.rand(2, 4)
_ = model(x)

for name, out in layer_outputs.items():
    print(f"{name}: shape = {out.shape}")

# Remove all hooks
for h in handles:
    h.remove()

Backward Hook — Capturing Gradients

gradients = {}

def capture_grad(name):
    def hook(module, grad_input, grad_output):
        gradients[name] = grad_output[0].detach()
    return hook

handle = model[0].register_full_backward_hook(capture_grad('layer0_grad'))

x      = torch.rand(3, 4)
target = torch.randint(0, 2, (3,))
loss   = nn.CrossEntropyLoss()(model(x), target)
loss.backward()

print(gradients['layer0_grad'].shape)   # torch.Size([3, 8])

handle.remove()

Tensor-Level Hook with register_hook

You can also attach hooks directly to tensors. This fires when the tensor's gradient is computed during backpropagation.

x = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)
y = (x ** 2).sum()

def print_grad(grad):
    print("Gradient of x:", grad)

x.register_hook(print_grad)

y.backward()
# Gradient of x: tensor([2., 4., 6.])

Practical Use Case: Gradient-Weighted Class Activation Map (Grad-CAM)

Grad-CAM shows which parts of an image a CNN focused on when making a classification decision. It uses forward hooks to capture feature maps and backward hooks to capture gradients — then multiplies them to produce a heatmap.

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import ResNet18_Weights

model  = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.eval()

feature_maps = {}
grads        = {}

def fwd_hook(m, i, o): feature_maps['last_conv'] = o
def bwd_hook(m, i, o): grads['last_conv']        = o[0]

fwd_h = model.layer4[-1].register_forward_hook(fwd_hook)
bwd_h = model.layer4[-1].register_full_backward_hook(bwd_hook)

x   = torch.rand(1, 3, 224, 224)
out = model(x)
out[0, out.argmax()].backward()

# Compute Grad-CAM
weights = grads['last_conv'].mean(dim=(2, 3), keepdim=True)
cam     = (weights * feature_maps['last_conv']).sum(dim=1).squeeze()
cam     = torch.relu(cam)   # keep only positive activations
print("CAM shape:", cam.shape)   # torch.Size([7, 7])

fwd_h.remove()
bwd_h.remove()

Forward Pre-Hook — Modify Input

def add_noise(module, input):
    # input is a tuple; return modified tuple
    noisy = input[0] + 0.01 * torch.randn_like(input[0])
    return (noisy,)

handle = model[0].register_forward_pre_hook(add_noise)

x   = torch.rand(2, 4)
out = model(x)   # noise added to input of layer 0 automatically

handle.remove()

Hook Best Practices

Always call handle.remove() when you're done.
  → Hooks stay attached until removed or the model is deleted.
  → Orphaned hooks slow down training and consume memory.

Use .detach() inside hooks when storing tensors.
  → Prevents the stored tensors from being part of the computation graph.

Keep hook functions lightweight.
  → Heavy computation inside a hook adds latency to every forward pass.

Use a list of handles to remove all at once:
  for h in handles: h.remove()

Summary

Hooks are callback functions that fire at specific points during forward or backward passes. Forward hooks capture layer outputs (activations) after each layer runs. Backward hooks capture gradients as they flow backward. Tensor-level hooks fire when a specific tensor's gradient is computed. Register hooks with register_forward_hook, register_full_backward_hook, or tensor's register_hook. Always store the returned handle and call handle.remove() when done. Hooks are the primary tool for activation visualization, gradient analysis, Grad-CAM, and debugging deep networks without touching model code.