Gradient of loss (that depends on gradient of network) with respect to parameters

bwehlin: <p>I’m trying to compute the gradient of my loss function with respect to my model parameters in PyTorch.</p>
<p>That is, let <code>u(x; θ)</code> be the model, where <code>x</code> is the input (in <code>R^n</code>) and <code>θ</code> are the model parameters. I’m trying to compute <code>du/dθ</code>.</p>
<p>For a “simple” loss function, this is not a problem, but my loss function depends on the gradient of the model with respect to its inputs (i.e., <code>du/dx</code>). When I attempt to do this, I’m met with the following error message: <code>One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code></p>
<p>Here is a minimal example to illustrate the issue:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.autograd import grad

model = nn.Sequential(nn.Linear(1, 10), nn.Tanh(), nn.Linear(10, 1))

def loss1(x, u):
    return torch.mean(u)

def loss2(x, u):
    d_u_x = grad(u, x, torch.ones_like(u), retain_graph=True, create_graph=True)[0]
    return torch.mean(d_u_x)

x = torch.randn(10, 1)
x.requires_grad_()
u = model(x)

loss = loss2(x, u)
d_loss_params = grad(loss, model.parameters(), retain_graph=True)
</code></pre>
<p>If I change the second to last line to read <code>loss = loss1(x, u)</code> things work as expected.</p>
<p>Update: it appears to be working if I set <code>bias=False</code> for the <code>nn.Linear</code>s. OK, that makes some sense since the bias is not trainable. But that begs the question, how do I extract only the trainable parameters to use in the gradient computation?</p>

KFrank: <p>Hi bwehlin!</p>
<aside class="quote no-group" data-username="bwehlin" data-post="1" data-topic="217275">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/bwehlin/48/75921_2.png" class="avatar"> bwehlin:</div>
<blockquote>
<p><code>One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code></p>
</blockquote>
</aside>
<p>Roughly speaking, as I don’t have the precise details on this, this is the basic behavior:</p>
<p>When a constant shows up in the computation graph, autograd is smart enough<br>
to no longer track it.  This is why you get a tensor that “appears not to have been<br>
used in the graph.”</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">model = nn.Sequential(nn.Linear(1, 10), nn.Tanh(), nn.Linear(10, 1))
...
def loss2(x, u):
    d_u_x = grad(u, x, torch.ones_like(u), retain_graph=True, create_graph=True)[0]
    return torch.mean(d_u_x)
...
u = model(x)

loss = loss2(x, u)
d_loss_params = grad(loss, model.parameters(), retain_graph=True)
</code></pre>
</blockquote>
</aside>
<p>The second <code>Linear</code> in your <code>Sequential</code> appears at the end.  Therefore that <code>Linear</code>’s<br>
<code>bias</code> appears only as an additive constant in the output of <code>u = model (x)</code>  <code>u</code> has a<br>
non-zero gradient with respect to the second <code>bias</code>, but it’s a constant, so autograd<br>
doesn’t include it in the graph created for <code>d_u_x</code>.</p>
<p>For example, if you change:</p>
<pre data-code-wrap="python"><code class="lang-python">model = nn.Sequential(nn.Linear(1, 10), nn.Tanh(), nn.Linear(10, 1))
</code></pre>
<p>to:</p>
<pre data-code-wrap="python"><code class="lang-python">model = nn.Sequential(nn.Linear(1, 10), nn.Tanh(), nn.Linear(10, 1), nn.Tanh())
</code></pre>
<p>then the output of <code>model</code> will include the second <code>bias</code> passed through the non-linear<br>
function <code>Tanh()</code>, so the gradient with respect to the second <code>bias</code> won’t be a constant<br>
and won’t get removed from the computation graph.  This changes <code>model</code>, of course,<br>
but the specific error will go away.</p>
<p>You could instead change:</p>
<pre data-code-wrap="python"><code class="lang-python">d_loss_params = grad(loss, model.parameters(), retain_graph=True)
</code></pre>
<p>to:</p>
<pre data-code-wrap="python"><code class="lang-python">d_loss_params = grad(loss, model.parameters(), retain_graph=True, allow_unused = True)
</code></pre>
<p>autograd still removes the constant branch from the computation graph, but now you’re<br>
telling the call to <code>grad()</code> that you expect to be trying to differentiate an “unused” tensor<br>
and that it’s not an error on your part.  So autograd proceeds without raising the error.</p>
<aside class="quote no-group">
<blockquote>
<p>Update: it appears to be working if I set <code>bias=False</code> for the <code>nn.Linear</code>s.</p>
</blockquote>
</aside>
<p>Yes, as noted above, this is because <code>bias</code> of the second <code>Linear</code> only contributes a<br>
constant to the output of <code>model</code>.  (You should be able to leave <code>bias = True</code> for the<br>
first <code>Linear</code>, only setting <code>bias = False</code> for the second <code>Linear</code>, and still get rid of the<br>
error.)  Of course, setting <code>bias = False</code> does change the model.</p>
<aside class="quote no-group">
<blockquote>
<p>OK, that makes some sense since the bias is not trainable.</p>
</blockquote>
</aside>
<p>Just to be clear, the <code>bias</code> of a <code>Linear</code> is, in general, trainable.  For example, in your<br>
case, the <code>bias</code> of the first <code>Linear</code> is trainable and would be changed by optimizer<br>
<code>step()</code>s.</p>
<p>If by “not trainable” you mean that the second derivative of the output of <code>model</code> with<br>
respect to the second <code>bias</code> is zero, then yes, you could phrase things that way.</p>
<aside class="quote no-group">
<blockquote>
<p>But that begs the question, how do I extract only the trainable parameters to use in the gradient computation?</p>
</blockquote>
</aside>
<p>As a practical matter, rather than arrange things so that the second <code>bias</code> still contributes<br>
its non-trivial constant value to the output of <code>model</code>, but have it not participate in any<br>
computation graph, just go ahead and use <code>allow_unused = True</code>. autograd will then do<br>
the correct thing (for your use case) – not including (nor computing) the derivative for the<br>
constant tensor in the back propagation and computing an explicit zero for that derivative<br>
both give the same final result.  (It’s just that leaving it out is a little cheaper.)</p>
<p>Best.</p>
<p>K. Frank</p>

bwehlin: <p>Hi KFrank,</p>
<p>Thank you for your answer.</p>
<p>The first thing I tried was actually <code>allow_unused=True</code>, but the problem was that I then passed the result of this on to <code>nn.utils.parameters_to_vector</code> which gives <code>AttributeError: 'NoneType' object has no attribute 'device'</code>.</p>
<p>I had somehow assumed <code>allow_unused=True</code> made all my gradients <code>None</code> but now inspecting further, in fact only certain gradients are <code>None</code>.</p>
<p>I’m trying to implement Inverse Dirichlet Weighting [1] and this requires computing the standard deviation of all gradients wrt the model’s parameters. I took another peek at the documentation of <code>autograd.grad</code> just now, and there is an additional option <code>materialize_grads</code> that when turned on sets <code>None</code> gradients to zero in the result. This gives exactly what I want!</p>
<p>So the solution is:</p>
<p><code>d_loss_params = grad(loss, model.parameters(), retain_graph=True, allow_unused=True, materialize_grads=True)</code></p>
<p>Thanks</p>
<p>References:<br>
[1] Maddu, S., Sturm, D., Müller, C. L., &amp; Sbalzarini, I. F. (2022). Inverse Dirichlet weighting enables reliable training of physics informed neural networks. <em>Machine Learning: Science and Technology</em> , <em>3</em> (1), 015026.</p>