Skip to content

← nn.Module Mechanics step 6 of 7

Medium Primitives

Read the middle of a model

Jump to the task

A forward hook lets you inspect a layer's output without changing the model's return value. Register a callback on that layer, run the model, then remove the callback when you are finished:

import torch
import torch.nn as nn

model = nn.Sequential(nn.ReLU(), nn.Identity())
captured = {}

def hook(module, inputs, output):
    captured["value"] = output.detach().clone()

handle = model[0].register_forward_hook(hook)
try:
    result = model(torch.tensor([-2.0, 3.0]))
finally:
    handle.remove()
print(captured["value"].tolist())  # [0.0, 3.0]
print(result.tolist())           # [0.0, 3.0]

With the default hook options, the callback receives the module, a tuple of positional inputs, and its output. It runs after the module's forward returns. Call the module normally, as model(x); calling its forward method directly bypasses hooks attached to that module.

Decide what to keep

For logging, detach removes the captured tensor's connection to autograd. Clone also makes an independent copy, so a later in-place change to the original output does not alter your snapshot. This example assumes a tensor output; a layer returning a tuple or dictionary needs corresponding handling.

Keeping attached activations can retain graph state. Appending one per batch can grow memory use, whereas replacing a single dictionary entry does not by itself imply unbounded growth. Detached copies still occupy memory too. If you need gradients through the captured activation, detaching would remove the connection you need, so choose based on the purpose of the hook.

Remove hooks and avoid changing the output accidentally

Hooks stay registered until removed. Registering another one every iteration without cleanup makes future calls run all of them. Use finally as above so cleanup happens even when the model call raises.

A forward hook that returns None leaves the output unchanged. Returning a non-None value replaces it. An ordinary Python function does not return its final expression automatically; an explicit return, or a lambda whose expression produces a value, can replace the output.

Hooks are useful for temporary inspection. An explicit forward return can be clearer when intermediate features are a permanent part of your model's interface. Forward pre-hooks inspect or replace inputs; backward hooks serve a different purpose, observing gradients rather than forward activations.

Your task

def capture_middle(x_values: list) -> dict

Use the supplied Linear, ReLU, Linear model. For input [a, b], its first layer produces [a, -b], ReLU replaces negative values with zero, and the last layer puts the sum of those two values in both output positions.

Attach a forward hook to the ReLU, save a detached independent copy of its output, run the model, and remove the hook in a finally block. Return:

  • captured: the ReLU output as a list, rounded to four decimals.
  • output: the final output as a list, rounded to four decimals.
  • hooks_left: the number of forward hooks remaining on the ReLU, which must be zero.

For [2.0, 3.0], return {"captured": [2.0, 0.0], "output": [2.0, 2.0], "hooks_left": 0}. The starter uses the private _forward_hooks mapping to inspect the count; use the public handle.remove() method to remove your hook.