Registering a hook for a module parameter forward pass instead of a backward pass

autograd
1

I have a list of select parameters from various parts of a model (type: nn.Module).

I want to get the inputs to the parameters as a forward pass is called on the model.

Is this possible? I only found hooks for the backward pass of a tensor: torch.Tensor.register_hook — PyTorch 2.0 documentation

2

Yes, you could use forward hooks via nn.Module.register_forward_hook.

3

I’m looking for a hook for parameters rather than modules. Something like torch.Tensor.register_forward_hook if it exists.

4

I don’t think there is one. Why can’t you access the parameters via the module?

5

Some modules already have other hooks for other reasons. Is it possible to register multiple hooks for the same module?

6

Yes, it is possible to use multiple hooks:

lin = nn.Linear(1, 1)
lin.register_forward_hook(lambda m, input, output: print("hook1 for module {}".format(m)))
lin.register_forward_hook(lambda m, input, output: print("hook2 for module {}".format(m)))

x = torch.randn(1, 1)
out = lin(x)
# hook1 for module Linear(in_features=1, out_features=1, bias=True)
# hook2 for module Linear(in_features=1, out_features=1, bias=True)