PyTorch Tensor Hooks

A tensor hook is a function that runs while a gradient travels backward through one tensor. The hook receives the gradient and may return a replacement for it.

A hook watches the gradient pass

The hook prints the gradient as it passes h. It returns nothing, so the gradient of x stays untouched.

import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)
h = x * 3


def watch(grad):
    print("Gradient passing h:", grad.tolist())


h.register_hook(watch)
(h ** 2).sum().backward()
print("Gradient of x:", x.grad.tolist())

A returned tensor replaces the gradient

This hook returns a row of zeros instead. The gradient of x prints zero, because the replacement flows onward.

x = torch.tensor([1.0, 2.0], requires_grad=True)
h = x * 3
h.register_hook(lambda grad: torch.zeros_like(grad))
(h ** 2).sum().backward()
print("Gradient of x:", x.grad.tolist())

The handle removes the hook

The hook counts its own calls on a leaf tensor. After removal, a second backward adds no new call, and the gradient still doubles.

calls = []


def count_hook(grad):
    calls.append(1)


x = torch.tensor([2.0], requires_grad=True)
handle = x.register_hook(count_hook)
(x ** 2).backward()
handle.remove()
(x ** 2).backward()
print("Hook calls:", len(calls))
print("Gradient of x:", x.grad.tolist())

Example

import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)
h = x * 3


def clip_and_log(grad):
    print("Hook saw gradient:", grad.tolist())
    print("Hook saw norm:", round(grad.norm().item(), 3))
    return grad.clamp(-5.0, 5.0)


handle = h.register_hook(clip_and_log)
loss = (h ** 2).sum()
loss.backward()
print("Loss:", loss.item())
print("Gradient of x with hook:", x.grad.tolist())
handle.remove()

plain = torch.tensor([1.0, 2.0], requires_grad=True)
plain_h = plain * 3
(plain_h ** 2).sum().backward()
print("Gradient of x without hook:", plain.grad.tolist())

Key points

Use register_hook on a tensor to attach a hook to that tensor
A hook that returns nothing leaves the gradient unchanged
A hook that returns a tensor replaces the gradient with that tensor
register_hook returns a handle, and handle.remove() detaches the hook
Gradient norm: the length of a gradient row, read as one number
In short, a tensor hook sits on one tensor and can watch or rewrite its gradient

In this example, h is three times x, and the loss sums the squares of h. clip_and_log is registered on h with register_hook. During backward, the hook prints the gradient passing h, which is 6 and 12. It prints the norm of that gradient as 13.416. It then returns the gradient clipped to the range from minus five to five. The next step back multiplies that clipped row by three, so x receives 15 and 15. The loss prints 45, from nine plus thirty-six. A fresh copy without the hook gives x the unclipped 18 and 36. handle.remove() then detaches the hook before the plain copy runs.

Note

Module hooks watch whole layers, while a tensor hook watches one tensor

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import torch

x = torch.tensor([1.0, 2.0], requires_grad=True)
h = x * 3


def clip_and_log(grad):
    print("Hook saw gradient:", grad.tolist())
    print("Hook saw norm:", round(grad.norm().item(), 3))
    return grad.clamp(-5.0, 5.0)


handle = h.register_hook(clip_and_log)
loss = (h ** 2).sum()
loss.backward()
print("Loss:", loss.item())
print("Gradient of x with hook:", x.grad.tolist())
handle.remove()

plain = torch.tensor([1.0, 2.0], requires_grad=True)
plain_h = plain * 3
(plain_h ** 2).sum().backward()
print("Gradient of x without hook:", plain.grad.tolist())

Exercise

What happens when a tensor hook returns a tensor?

The returned tensor replaces the gradient The hook runs twice Backward stops at once The tensor loses its value