🎓 A1Learner
🔍

PyTorch Forward Hooks

A forward hook is a function PyTorch calls each time a layer finishes a forward pass. The hook receives the layer, its inputs, and its output.

One hook on one layer

The hook prints two facts about the output it just saw. The layer itself needs no change.

torch.manual_seed(0)
layer = nn.Linear(4, 3)
def show(module, inputs, output):
    print("min", round(output.min().item(), 2), "max", round(output.max().item(), 2))
handle = layer.register_forward_hook(show)
out = layer(torch.randn(2, 4))

A hook on ReLU

ReLU clips negative values to zero. The hook proves it from the output side, by reading the minimum.

torch.manual_seed(0)
relu = nn.ReLU()
def show(module, inputs, output):
    print("ReLU output min:", output.min().item())
handle = relu.register_forward_hook(show)
out = relu(torch.tensor([-2.0, 0.5, -1.0, 3.0]))

Remove the hook

The handle controls the hook's lifetime. After removal, a second pass adds no new call.

torch.manual_seed(0)
model = nn.Linear(3, 2)
calls = []
def record(module, inputs, output):
    calls.append(tuple(output.shape))
handle = model.register_forward_hook(record)
model(torch.randn(4, 3))
handle.remove()
model(torch.randn(4, 3))
print("Hook calls:", len(calls))

Example

import torch
from torch import nn
from collections import OrderedDict

torch.manual_seed(0)
model = nn.Sequential(OrderedDict([
    ("linear1", nn.Linear(4, 8)),
    ("relu", nn.ReLU()),
    ("linear2", nn.Linear(8, 3)),
]))


def range_hook(name):
    def hook(module, inputs, output):
        print(f"{name}: shape {tuple(output.shape)} "
              f"min {output.min().item():.2f} "
              f"max {output.max().item():.2f} "
              f"mean {output.mean().item():.2f}")
    return hook


handles = []
for name, layer in model.named_children():
    handles.append(layer.register_forward_hook(range_hook(name)))

x = torch.randn(5, 4)
output = model(x)
print("Final output shape:", tuple(output.shape))

for handle in handles:
    handle.remove()
print("Hooks removed:", len(handles))

Key points

  • Hook: a function registered on a layer and called after each of its forward passes
  • Use layer.register_forward_hook(fn) to attach a hook to a layer
  • register_forward_hook returns a handle, and handle.remove() detaches the hook
  • A hook can read a layer's output without changing the model's code
  • Hooks watch a pass from outside the model, one layer at a time
  • A hook runs on every forward pass until it is removed

In this example, hooks watch every layer of a small network during one real pass. The model is a Sequential with three named layers. A small factory builds one hook per layer, and each hook closes over its layer's name. Each hook prints the output shape, minimum, maximum, and mean, all rounded to two decimals. Five random rows then run through the model. The ReLU line shows a minimum of 0.00, because ReLU clipped every negative value. The last layer prints the shape (5, 3), which matches the final output shape. All three handles are removed at the end, and the count prints 3.

Note

A hook that returns a value replaces the layer's output, so a watching hook must return nothing

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does handle.remove() do to a registered forward hook?