Clarifications on register_full_backward_hook and register_forward_pre_hook

username: AlphaBetaGamma96
cooked: <p>After reading the documents for <code>register_full_backward_hook</code> <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=register_full_backward_hook#torch.nn.Module.register_full_backward_hook" rel="noopener nofollow ugc">here</a> and <code>register_forward_pre_hook</code> <a href="https://pytorch.org/docs/stable/generated/torch.nn.Module.html?highlight=register_forward_pre_hook#torch.nn.Module.register_forward_pre_hook" rel="noopener nofollow ugc">here</a>. I was wondering if I could ask a few clarification questions that aren’t clear within the documents?</p>
<p>From my understanding, <code>register_full_backward_hook</code> returns the <code>grad_output</code> and <code>grad_input</code> variables which are the gradient of the loss with respect to the <code>output</code> and <code>input</code> respectively. And, for <code>register_forward_pre_hook</code> returns the <code>input</code> of a given layer.</p>
<p>For example, I use a <code>nn.Linear</code> module like <code>Linear(in_features=2, out_features=32, bias=True)</code>. where the input data has shape <code>[B, A, 2]</code> where <code>B</code> is the batch and <code>A</code> is another dimension to the input shape. When applying a <code>forward_pre_hook</code> the <code>input</code> shape is the same as my input which makes sense (it’s shape <code>[B, A, 2]</code>). When applying a <code>backward_full_hook</code> the <code>grad_output</code> shape is <code>[B, 32, A]</code> which makes sense. Is it right to assume that this Tensor is the individial gradients of the loss with respect to the layer output for <strong>all</strong> inputs? I.e. each element here represent the gradient of the loss for a given input before any kind of reduction operation like <code>torch.mean</code>?</p>
<p>Also, is it possible to reconstruct the gradient of the layer from its <code>grad_output</code> variables? For example, with a <code>nn.Linear</code> layer, its gradient of the loss can be decomposed into the gradient of the loss with respect to the output of the layer (aka <code>grad_output</code>) and the gradient of the output of the layer with respect to the weights of the layer (which for a Linear layer is just the input features, which are returned by <code>forward_pre_hook</code>). I tried doing the following to reconstruct the gradients but it’s not always the same as getting <code>nn.Linear.weight.grad</code>,</p>
<pre><code class="lang-auto">grad_output  #has shape [B,32,A]
input #has shape [B,2,A]

grad_from_hooks = torch.einsum("bia,boa-&gt;io", grad_output, input)
grad_from_module = nn.Linear.weight.grad

ratio = grad_from_hooks / grad_from_modules #should be 1 if equal. 
</code></pre>
<p>In my case this ratio value is near 1 for nearly all elements of the Tensor (in range of 0.95 to 1.05). Although, some elements differ by a factor of 2 or 3, and I was wondering why exactly doesn’t this formula work when reconstructing the gradients from the <code>grad_ouput</code> and <code>input</code> Tensors? The expression I saw was the outer-product of the <code>grad_output</code> and <code>input</code> should give the gradient (which is what I’ve done within the <code>torch.einsum</code> expression).</p>
<p>Is it at all possible to reconstruct the gradient of the layer by using the layers <code>grad_output</code> and <code>input</code> values from <code>register_backward_full_hook</code> and <code>register_forward_pre_hook</code>?</p>
<p>Apologizes for this being quite an in-depth question but I’d appreciate any clarification on this issue!</p>
<p>Thank you! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>