Is layer_activation (register_forward_hook) the same as gradient?

Question

I was wondering if the intermediate layer output initialised by register_forward_hook is the same as the gradient of the intermediate output wrt the image?

Answer 1

You can attach a callback function on a given module with nn.Module.register_full_backward_hook to hook onto the backward pass of that layer. This allows you to access the gradient.

Here is a minimal example, define the hook as you did:

def backward_hook(module, grad_input, grad_output):
    print('grad_output:', grad_output)

Initialize your model and attach the hook on its layers

>>> model = nn.Sequential(nn.Linear(10, 5), nn.Linear(5, 2))
Sequential(
  (0): Linear(in_features=10, out_features=5, bias=True)
  (1): Linear(in_features=5, out_features=2, bias=True)
)

>>> for name, layer in model.named_children():
...     print(f'hook onto {name}')
...     layer.register_full_backward_hook(backward_hook)
hook onto 0
hook onto 1

Perform an inference:

>>> x = torch.rand(5, 10)
>>> y = model(x).mean()

Perform the backward pass:

>>> y.backward()
grad_output: (tensor([[0.1000, 0.1000],
        [0.1000, 0.1000],
        [0.1000, 0.1000],
        [0.1000, 0.1000],
        [0.1000, 0.1000]]),)
grad_output: (tensor([[ 0.0135,  0.0141, -0.0468, -0.0378, -0.0123],
        [ 0.0135,  0.0141, -0.0468, -0.0378, -0.0123],
        [ 0.0135,  0.0141, -0.0468, -0.0378, -0.0123],
        [ 0.0135,  0.0141, -0.0468, -0.0378, -0.0123],
        [ 0.0135,  0.0141, -0.0468, -0.0378, -0.0123]]),)

---

For more examples, you can look at my other answers related to register_full_backward_hook:

• Using a non-full backward hook when the forward contains multiple Autograd nodes

• How to create a PyTorch hook with conditions?

• How to get all the tensors in a graph?