Gradient of bias in final layer "None" after backward through torch.autograd.grad

ArneT: <p>Hi,</p>
<p>I am trying to use the gradient of the output of my neural network w.r.t. the inputs as part of my loss calculation. However, I noticed that this causes the bias in my final layer to receive no gradient (stuck at “None”).</p>
<p>I created the following example code to demonstrate this behavior:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class test_nn(nn.Module):
    def __init__(
        self,
    ):
        super(test_nn, self).__init__()
        hidden_layer_size = 6
        n_hidden_layers = 4
        self.in_fcn = nn.Linear(16, hidden_layer_size)
        self.hidden_layers = nn.ModuleList(
            modules=[
                nn.Linear(hidden_layer_size, hidden_layer_size)
                for _ in range(n_hidden_layers)
            ]
        )
        self.out_fcn = nn.Linear(hidden_layer_size, 1)
        self.activation = nn.Softplus()

    def forward(self, x):
        x = self.activation(self.in_fcn(x))
        for layer in self.hidden_layers:
            x = self.activation(layer(x))
        x = self.out_fcn(x).squeeze(0)
        return x


if __name__ == "__main__":
    net = test_nn()
    x = torch.rand(32, 16, requires_grad=True)

    out = net(x)
    loss = nn.MSELoss()(out, torch.rand(32, 1))
    loss.backward()

    print("Grads of normal output")
    print("out bias grad: ", net.out_fcn.bias.grad)
    print("in bias grad: ", net.in_fcn.bias.grad)

    net.zero_grad()

    out = net(x)
    grad = torch.autograd.grad(
        out,
        x,
        grad_outputs=torch.ones_like(out),
        create_graph=True,
        retain_graph=True,
    )[0]
    print(grad.shape)
    loss = nn.MSELoss()(grad, torch.rand(32, 16))
    loss.backward()

    print("Grads of normal output")
    print("out bias grad: ", net.out_fcn.bias.grad)
    print("in bias grad: ", net.in_fcn.bias.grad)
</code></pre>
<p>Which results in the following output (obviously the actual values of the gradients are variable between runs):</p>
<pre><code class="lang-auto">Grads of normal loss
out bias grad:  tensor([0.4833])
in bias grad:  tensor([-7.7400e-05,  3.6427e-04,  2.8540e-04, -9.1523e-05, -8.3902e-04,
         3.4275e-06])
torch.Size([32, 16])
---------------------
Grads of grad loss
out bias grad:  None
in bias grad:  tensor([-7.0170e-07,  6.5941e-06, -1.0075e-05, -5.6637e-06,  2.2724e-05,
         2.1236e-06])
</code></pre>
<p>Does anybody know if this is expected behavior that I am overlooking? Or is there something else going wrong?</p>

BavoLesy: <p>Weirdly I have had this same issue. I have not found a solution yet <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>

KFrank: <p>Hi Arne!</p>
<aside class="quote no-group" data-username="ArneT" data-post="1" data-topic="207820" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/f475e1/48.png" class="avatar"> ArneT:</div>
<blockquote>
<p>the bias in my final layer to receive no gradient (stuck at “None”).</p>
</blockquote>
</aside>
<p>The gradient of <code>out</code> with respect to <code>x</code> does not depend on<br>
<code>out_fcn.bias</code>.  This is because <code>out_fcn.bias</code> is just added on<br>
to <code>out</code> (without any further mixing).  Therefore the gradient of the<br>
gradient of <code>out</code> with respect to <code>out_fcn.bias</code> is zero.</p>
<p>I’ve have seen this before where pytorch, recognizing that such a<br>
gradient will always be zero, sets such a gradient to <code>None</code> (rather<br>
than to <code>0.0</code>).  I don’t know the details of how or when this happens,<br>
but this is what is happening to you.</p>
<p>I’ve tweaked your code so that the first gradient does depend on<br>
<code>out_fcn.bias</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">    out = net(x)
    out = out * out   # make the gradient of out with respect to x depend on out_fcn.bias
    grad = torch.autograd.grad(
        out,
    ...
</code></pre>
<p>and now get a non-zero, non-<code>None</code> <code>out_fcn.bias.grad</code>:</p>
<pre><code class="lang-plaintext">2.4.0
Grads of normal output
out bias grad:  tensor([-1.8349])
in bias grad:  tensor([-7.3644e-04,  5.1432e-04,  1.7947e-03,  5.6672e-05,  5.3856e-04,
        -6.7751e-04])
torch.Size([32, 16])
Grads of normal output
out bias grad:  tensor([8.5547e-05])
in bias grad:  tensor([-4.5758e-07, -2.1357e-07, -9.9013e-06,  4.3018e-07,  1.6015e-06,
        -4.4209e-06])
</code></pre>
<p>(Note that the <code>print()</code> statements in the code you posted don’t<br>
match the output text that you posted.  I did not correct this in my<br>
tweaked version.)</p>
<p>Best.</p>
<p>K. Frank</p>