Search code examples
pytorchgradient

Is layer_activation (register_forward_hook) the same as gradient?

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?

Solution

  • 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: