Second order derivatives and inplace gradient "zeroing"

username: dpernes
cooked: <p>I am developing a model where, for some technical reasons, my minimization objective is actually the squared norm of the gradient of some loss function:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/6/6f9a4c7ee092ad6331447f09b090e88aeab43de2.png" alt="image" data-base62-sha1="fVhA8ERbjMirZqDa64Eip141rO2" width="137" height="49"></p>
<p>In the following toy example, I am implementing this for a very simple case (least-squares linear regression from the 2-dimensional space to the 1-dimensional space):</p>
<pre><code class="lang-auto">import torch
from torch import Tensor
from torch.autograd import Variable
from torch import nn

# some toy data
x = Variable(Tensor([4., 2.]), requires_grad=False)
y = Variable(Tensor([1.]), requires_grad=False)

# linear model and squared difference loss
model = nn.Linear(2, 1)
loss = torch.sum((y - model(x))**2)

# first backward pass
model.zero_grad()
# set create_graph=True to allow the computation of higher order derivatives afterwards
loss.backward(create_graph=True)

# compute the squared norm of the loss gradient
gn2 = Variable(Tensor([0.]))
for param in model.parameters():
  gn2 = gn2 + param.grad.norm()**2

# we do not want to accumulate the previous gradients with the new ones,
# so we have to call zero_grad() again
model.zero_grad()
gn2.backward()

</code></pre>
<p>However, PyTorch complains about backpropagating through variables that were modified by inplace operations:</p>
<pre><code class="lang-auto">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
</code></pre>
<p>After some debug, I found out that the second call to model.zero_grad() is responsible for that error, because it sets gradients to zero using inplace operations. An obvious (but maybe inelegant?) workaround is manually “zeroing” the gradients in the for loop and removing the second call to model.zero_grad(), that is, doing:</p>
<pre><code class="lang-auto">...
# compute the squared norm of the loss gradient
gn2 = Variable(Tensor([0.]))
for param in model.parameters():
  gn2 = gn2 + param.grad.norm()**2
  param.grad = torch.zeros_like(param)

gn2.backward()
</code></pre>
<p>Everything would be cleaner if the function  zero_grad() had an inplace argument, that we could set to False in situations like this…</p>
<p>What do you think? Which is the correct way to deal with this situation?</p>
<p>Thanks in advance.</p>

username: tom
cooked: <p>The usual way is to use torch.autograd.grad instead of backward for the derivative you want to include in your loss.</p>
<p>Best regards</p>
<p>Thomas</p>

username: dpernes
cooked: <p>Thank you very much, Thomas!</p>
<p>This is my code following your suggestion, in case anyone else comes across this problem and finds it useful:</p>
<pre><code class="lang-auto">import torch
from torch import Tensor
from torch.autograd import Variable
from torch.autograd import grad
from torch import nn

# some toy data
x = Variable(Tensor([4., 2.]), requires_grad=False)
y = Variable(Tensor([1.]), requires_grad=False)

# linear model and squared difference loss
model = nn.Linear(2, 1)
loss = torch.sum((y - model(x))**2)

# instead of using loss.backward(), use torch.autograd.grad() to compute gradients
loss_grads = grad(loss, model.parameters(), create_graph=True)

# compute the squared norm of the loss gradient
gn2 = sum([grd.norm()**2 for grd in loss_grads])

model.zero_grad()
gn2.backward()
</code></pre>
<p>Regards,<br>
Diogo</p>

username: Robin_Ranabhat
cooked: <p>Following to this question,</p>
<p>I am hopelessly trying to figure out how the computational graph would look like with second derivatives involved. I understand how auto diff works in 1st order derivative case and can work my way with drawing the computational graph. But for higher-order derivatives, I simply cannot do so.</p>
<p>I can only find the code on how this can be done in the internet but not an exact figurative answer. <a href="https://lucainiaoge.github.io/download/PyTorch-create_graph-is-true_Tutorial_and_Example.pdf" rel="noopener nofollow ugc">This</a> is the closest answer I could find, but still, didn’t quite get it.</p>
<p>And help would be really really appreciated</p>

username: tom
cooked: <p>Ha! I could spend hours with this (in fact I did record an entire course on autograd, but I don’t want to spam the forum with advertisement), but so the short story is two things on this:</p>
<ul>
<li>As mentioned above, I strongly advocate to <em>never</em> use <code>create_graph</code> with <code>t.backward</code> but only ever use it with <code>torch.autograd.grad</code>. <code>t.backward</code> is a convenience function that computes gradients and accumulates them in <code>.grad</code>. You almost certainly don’t want that to happen with the bits you take second derivatives of, and it creates all sorts of issues (e.g. with circular references).</li>
<li>Remember that inside the <code>backward</code> of an autograd function, you are using normal PyTorch operations. In this sense an oversimplified explanation of higher order derivatives is that <em>all that <code>create_graph</code> does is to enable <code>requires_grad=True</code> in the backward</em>. This should help you to draw the graph (adding to the computational graph you had in the forward while you are computing the gradient from the initial forward).</li>
</ul>
<p>Best regards</p>
<p>Thomas</p>

username: Robin_Ranabhat
cooked: <p>If I am not asking too much, could you post a rough hand-drawn picture of how the computational graph for this code below would look like. And please do share the link of your auto-grad course. : <strong>robinnarsingha123 at gmail dot com</strong></p>
<pre><code class="lang-auto">import torch

theta = torch.nn.Parameter(data=torch.Tensor([5.,]), requires_grad=True)
x = torch.tensor(1., requires_grad = False)

output = torch.multiply( x , torch.square(theta) )
output_grad = torch.autograd.grad(output, theta, create_graph=True)

final_loss = torch.add( output, output_grad[0] )

final_loss.backward()
</code></pre>

username: tom
cooked: <p>The ptgraph function from my publicly <a href="https://github.com/t-vi/lod2021/" rel="noopener nofollow ugc">available notes (for a general PyTorch tutorial)</a> gives<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/0/b014d47ed5f28f4381b425b5c35fe1f4adef666a.png" alt="image" data-base62-sha1="p7GF6yShA6ZMSAIW6Mc5PiKJF9w" width="467" height="443"></p>
<p>My interactive tool (that I wanted to open source but didn’t get around to just yet) from the autograd course can expand the grad line to this:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/4/d46f362aecbe4119694b6618907983e56b5f0bbc.png" data-download-href="https://discuss.pytorch.org/uploads/default/d46f362aecbe4119694b6618907983e56b5f0bbc" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/4/d46f362aecbe4119694b6618907983e56b5f0bbc.png" alt="image" data-base62-sha1="ujhuAUtNULaAVrYtTnQeZpFiDZW" width="690" height="497" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">946×682 35.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Best regards</p>
<p>Thomas</p>

username: nobody
cooked: <p>I’m sorry to open such an old question, but I did have some difficulties. I need to add gradient L1 regularization to my loss, and I find that only the following two times backward implementation can make its update result different from the result without gradient regularization. If, as you suggested, torch.autograd.grad is used to replace the first backward, the updated result is the same as the result without gradient regularization, I don’t know why this happens. And whether the two backward implementation would actually work as I hoped.</p>
<p>the code of toy:</p>
<pre><code class="lang-auto">import torch
x = torch.randn(3,4)
fc1 = torch.nn.Linear(4,3)
fc2 = torch.nn.Linear(4,3)
fc3 = torch.nn.Linear(4,3)
fc2.weight.data.copy_(fc1.weight.data)
fc2.bias.data.copy_(fc1.bias.data)
fc3.weight.data.copy_(fc1.weight.data)
fc3.bias.data.copy_(fc1.bias.data)
opt1 = torch.optim.SGD(fc1.parameters(), lr=0.01)
opt2 = torch.optim.SGD(fc2.parameters(), lr=0.01)
opt3 = torch.optim.SGD(fc3.parameters(), lr=0.01)

# without gradient regularization
y1 = fc1(x)
loss1 = (1-y1).sum()
loss1.backward()
opt1.step()

# with gradient regularization, first implementation
y2 = fc2(x)
l2 = (1-y2).sum()
l2.backward(create_graph=True)
np2 = sum([p.grad.norm(p=1) for p in fc2.parameters()])
loss2 = l2 + np2
loss2.backward()
opt2.step()

# with gradient regularization, second implementation
y3 = fc3(x)
l3 = (1-y3).sum()
np3 = sum([g.norm(p=1) for g in  torch.autograd.grad(l3, fc3.parameters(), create_graph=True) ])
loss3 = l3 + np3
loss3.backward()
opt3.step()

# print weight after update
print(fc1.weight)
print(fc2.weight)  # different to fc1.weight
print(fc3.weight)  # same to fc1.weight
</code></pre>

username: nobody
cooked: <p>I seem to know where I went wrong. The second derivative of the linear function is 0, which makes torch.autograd.grad same to without gradients norm, but not if I make it be nonlinear function.</p>
<p>Looks like the torch.autograd.grad implementation is the right one? If anyone can help confirm would be greatly appreciated!</p>