Understanding backward hooks

Alexander_Soare: <p>I wrote this snippet below to try and understand what’s going on with these hooks.</p>
<pre><code class="lang-python">class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10,5)
        self.fc2 = nn.Linear(5,1)
        self.fc1.register_forward_hook(self._forward_hook)
        self.fc1.register_backward_hook(self._backward_hook)
    
    def forward(self, inp):
        return self.fc2(self.fc1(inp))

    def _forward_hook(self, module, input, output):
        print(type(input))
        print(len(input))
        print(type(output))
        print(input[0].shape)
        print(output.shape)
        print()

    def _backward_hook(self, module, grad_input, grad_output):
        print(type(grad_input))
        print(len(grad_input))
        print(type(grad_output))
        print(len(grad_output))
        print(grad_input[0].shape)
        print(grad_input[1].shape)
        print(grad_output[0].shape)
        print()

model = Model()
out = model(torch.tensor(np.arange(10).reshape(1,1,10), dtype=torch.float32))
out.backward()
</code></pre>
<p>Produces output</p>
<pre><code class="lang-auto">&lt;class 'tuple'&gt;
1
&lt;class 'torch.Tensor'&gt;
torch.Size([1, 1, 10])
torch.Size([1, 1, 5])

&lt;class 'tuple'&gt;
2
&lt;class 'tuple'&gt;
1
torch.Size([1, 1, 5])
torch.Size([5])
torch.Size([1, 1, 5])
</code></pre>
<p>You can also follow the CNN example <a href="https://pytorch.org/tutorials/beginner/former_torchies/nnft_tutorial.html#forward-and-backward-function-hooks" rel="noopener nofollow ugc">here</a>. In fact, it’s needed to understand the rest of my question.</p>
<p>I have a few questions:</p>
<ol>
<li>
<p>I would normally think that <code>grad_input</code> (backward hook) should be the  same shape as <code>output</code> (forward hook) because when we go backwards, the direction is reversed. But the CNN example seems to indicate otherwise. I’m still a bit confused. Which way around is it?</p>
</li>
<li>
<p>Why are <code>grad_input[0]</code> and <code>grad_output[0]</code> the same shape on my <code>Linear</code> layer here? Regardless of the answer to my question 1, at least one of them should be <code>torch.Size([1, 1, 10])</code> right?</p>
</li>
<li>
<p>What’s with the second element of the tuple <code>grad_input</code>? In the CNN case I copy pasted the example and did <code>print(grad_input[1].size())</code> with output <code>torch.Size([20, 10, 5, 5])</code>. So I presume it’s the gradients of the weights. I also ran <code>print(grad_input[2].size())</code> and got <code>torch.Size([20])</code>. So it seemed clear I was looking at the gradients of the biases. But then in my <code>Linear</code> example <code>grad_input</code> is length 2, so I can only access up to <code>grad_input[1]</code>, which seems to be giving me the gradients of the biases. So then where are the gradients of the weights?</p>
</li>
</ol>
<p>In summary, there are two apparent contradictions between the behaviour of the backwards hook in the cases of <code>Conv2d</code> and `Linear’ modules. This has left me totally confused about what to expect with this hook.</p>
<p>Thanks for your help!</p>