Autograd problem when designing a custom loss

username: Luis
cooked: <p>Hello everybody,</p>
<p>I am having a hard time while I am trying to design a loss function that applies Sobel filter to the batches before computing MSE. I am quite sure that the problem is related to an “autograd computational graph detachment”, but I just cannot solve it.</p>
<p>Here is my code. Does anyone can see what I am missing?</p>
<pre><code>def sobel_MSE(output, target):
dx = (torch.tensor([[1.0, 0.0, -1.0],[2.0, 0.0, -2.0],[1.0, 0.0, -1.0]], requires_grad=True)).float()
dy = (torch.tensor([[1.0, 2.0, 1.0], [0.0, 0.0, 0.0], [-1.0, -2.0, -1.0]], requires_grad=True)).float()

dx = dx.cuda()
dy = dy.cuda()

dx = dx.view((1, 1, 3, 3))
dy = dy.view((1, 1, 3, 3))

doutdx = nn.functional.conv2d(output, dx, padding=1)
doutdy = nn.functional.conv2d(output, dy, padding=1)

dtardx = nn.functional.conv2d(target, dx, padding=1)
dtardy = nn.functional.conv2d(target, dy, padding=1)

dout = torch.sqrt(torch.pow(doutdx, 2) + torch.pow(doutdy, 2))
dtar = torch.sqrt(torch.pow(dtardx, 2) + torch.pow(dtardy, 2))

out = torch.mean(torch.pow(dout-dtar,2))

return out</code></pre>

username: ptrblck
cooked: <p>Would you like to train <code>dx</code> and <code>dy</code> (since you’ve set their <code>requires_grad</code> attribute to <code>True</code>)?<br>
If so, you shouldn’t overwrite them in these lines of code, but would need to use a new variable name and pass the original <code>dx</code> and <code>dy</code> to the optimizer:</p>
<pre><code class="lang-python">dx = dx.cuda()
dy = dy.cuda()

dx = dx.view((1, 1, 3, 3))
dy = dy.view((1, 1, 3, 3))
</code></pre>
<p>On the other hand, if you don’t want to train these tensors, you could remove the <code>requires_grad</code> setting and as long as <code>output</code> is attached to the graph, the code should work.<br>
Are you seeing any errors or unexpected behavior?</p>

username: Luis
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="91404">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>ong as <code>output</code> is attached to the graph, the code should work.</p>
</blockquote>
</aside>
<p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a> for you reply.</p>
<p>I removed the requires_grad setting to make the sobel filter static. The problem is, when I run the code - after the first loss.backward() - the network output becomes nan. I believe it was some problem when computing grad.</p>
<p>If I change the loss for MSE doing:</p>
<p>def my_mse(output, target):<br>
out = torch.mean(torch.pow(output-target,2))</p>
<p>everything works fine. This way I suppose that there is no problem in the rest of the code.</p>
<p>Do you have any idea about what I am missing?</p>
<p>Thanks a lot</p>

username: ptrblck
cooked: <p>Could you rerun the code with <code>torch.autograd.set_detect_anomaly(True)</code> at the beginning of the script and post the stack trace here?<br>
Based on the description I assume the loss does not contain any invalid values?<br>
If that’s the case, could you check all gradients after the first <code>backward()</code> pass for NaN values?</p>

username: Luis
cooked: <p>Thanks again <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<p>By enablig the anomaly detection, I found the following error trace:</p>
<pre><code class="lang-auto">..\torch\csrc\autograd\python_anomaly_mode.cpp:57: UserWarning: Traceback of forward call that caused the error:
 File "C:/Users/Luis/PycharmProjects/unet-googlenet/train.py", line 291, in &lt;module&gt;
   train()
 File "C:/Users/Luis/PycharmProjects/unet-googlenet/train.py", line 214, in train
   loss = sobel_MSE(outputs, labels)
 File "C:/Users/Luis/PycharmProjects/unet-googlenet/train.py", line 142, in sobel_MSE
   dout = torch.sqrt(torch.pow(doutdx, 2) + torch.pow(doutdy, 2))

Traceback (most recent call last):
 File "C:/Users/Luis/PycharmProjects/unet-googlenet/train.py", line 291, in &lt;module&gt;
   train()
 File "C:/Users/Luis/PycharmProjects/unet-googlenet/train.py", line 216, in train
   loss.backward()
 File "C:\Users\Luis\Anaconda3\envs\unet-goolenet\lib\site-packages\torch\tensor.py", line 166, in backward
   torch.autograd.backward(self, gradient, retain_graph, create_graph)
 File "C:\Users\Luis\Anaconda3\envs\unet-goolenet\lib\site-packages\torch\autograd\__init__.py", line 99, in backward
   allow_unreachable=True)  # allow_unreachable flag
RuntimeError: Function 'SqrtBackward' returned nan values in its 0th output.
</code></pre>
<p>Does it means that the gradient calculated by torch.srqt was nan at the following line?</p>
<p>dout = torch.sqrt(torch.pow(doutdx, 2) + torch.pow(doutdy, 2))</p>

username: Luis
cooked: <aside class="quote no-group" data-username="Luis" data-post="5" data-topic="91404">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/luis/48/27409_2.png" class="avatar"> Luis:</div>
<blockquote>
<p><code>h.sqrt(torch.pow(doutdx, 2) + torch.pow(doutdy, 2))</code></p>
</blockquote>
</aside>
<p>Yes!</p>
<p>I replaced</p>
<p><code>dout = torch.sqrt(torch.pow(doutdx, 2) + torch.pow(doutdy, 2))</code></p>
<p>by</p>
<p><code>dout = (torch.pow(doutdx, 2) + torch.pow(doutdy, 2))</code></p>
<p>and now it seems to compute the backprojection fine! By removing the sqrt we do not lose the edge detection semantic, but now I am curious to understand why it is happening.</p>
<p>Do you have any idea?</p>

username: ptrblck
cooked: <p><code>torch.sqrt()</code> would give you <code>NaN</code> gradients for negative inputs, which shouldn’t be possible given your code snippet and an <code>Inf</code> gradient for a zero input, which might be the case for your use case:</p>
<pre><code class="lang-python">x = torch.tensor([0.], requires_grad=True)
y = torch.sqrt(x)
y.backward()
print(x.grad)
&gt; tensor([inf])
</code></pre>
<p>You could add a small <code>eps</code> value to the <code>sqrt</code> op to avoid this. However, note that this might blow up your gradients due to the behavior of the derivative of the <code>sqrt</code>.</p>