Backward pass of sobel filter loss gives NaN

MyNameIsArko: <p>Hi! Thanks in advance for anyone that will take a look at my problem.</p>
<p>Here is my custom loss function:</p>
<pre data-code-wrap="python"><code class="lang-python">def mean_gradient_error(outputs, targets, weight=1e-4):
    print()
    filter_x = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=outputs.dtype, device=outputs.device)
    filter_y = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=outputs.dtype, device=outputs.device)

    filter_x = filter_x.unsqueeze(0).expand(3, -1, -1).unsqueeze(0)
    filter_y = filter_y.unsqueeze(0).expand(3, -1, -1).unsqueeze(0)

    # Output gradient
    output_gradient_x = torch.nn.functional.conv2d(outputs, filter_x, padding='same')
    output_gradient_y = torch.nn.functional.conv2d(outputs, filter_y, padding='same')
    
    sq_output_gradient_x = torch.square(output_gradient_x)
    sq_output_gradient_y = torch.square(output_gradient_y)

    print('sq_y', sq_output_gradient_y.min().item(), sq_output_gradient_y.max().item())
    sq_output_gradient_y.register_hook(lambda y: print('sq_y back', y.min().item(), y.max().item()))

    # Target gradient
    target_gradient_x = torch.square(torch.nn.functional.conv2d(targets, filter_x, padding='same'))
    target_gradient_y = torch.square(torch.nn.functional.conv2d(targets, filter_y, padding='same'))

    # Compute mean gradient error
    output_gradients = torch.sqrt(sq_output_gradient_x + sq_output_gradient_y)

    output_gradients.register_hook(lambda o: print('sqrt_grad back', o.min().item(), o.max().item()))
    print('sqrt_grad', output_gradients.min().item(), output_gradients.max().item())
    
    target_gradients = torch.sqrt(target_gradient_x + target_gradient_y)

    shape = output_gradients.shape[2:4]
    diff = torch.square(output_gradients - target_gradients)

    mge = torch.sum(diff / (shape[0] * shape[1]))

    return mge * weight
</code></pre>
<p>Following <a href="https://discuss.pytorch.org/t/nan-in-backward-pass-for-torch-square/97489">this conversation</a> I added prints of values for every calculation and narrowed down, that error is between backward pass of <code>torch.sqrt(sq_output_gradient_x + sq_output_gradient_y)</code> and <code>torch.square(output_gradient_y)</code>, both of which I left prints in code above.</p>
<p>Here are the prints of last epochs:</p>
<pre><code class="lang-auto">[...]
sq_y 5.115907697472721e-13 1121.9591064453125
sqrt_grad 0.0012385726440697908 33.49565887451172
sqrt_grad back -9.766469588612381e-08 5.78952175089853e-08
sq_y back -1.497764174018812e-06 1.525878978725359e-09

sq_y 5.1159076974727213e-11 1134.2796630859375
sqrt_grad 0.00048042723210528493 33.679073333740234
sqrt_grad back -1.0675010031491183e-07 6.707357869117914e-08
sq_y back -3.1196209420158993e-06 1.525878978725359e-09

sq_y 0.0 919.6836547851562
sqrt_grad 0.0 32.1100959777832
sqrt_grad back -1.3250506469830725e-07 5.5954753719333894e-08
sq_y back -inf 1.525878978725359e-09
</code></pre>
<p>As you can see backward pass of <code>torch.square</code> produces <code>-inf</code>, which I don’t know why.</p>
<p>Here is autograd error and python error:</p>
<pre><code class="lang-auto">[...]
    if await self.run_code(code, result, async_=asy):
  File "/home/adrian/.miniforge3/envs/torch/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3577, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/tmp/ipykernel_34168/3323993857.py", line 15, in &lt;module&gt;
    loss, info = train_step(model, mse_loss, scaler, small_img, img)
  File "/tmp/ipykernel_34168/187903903.py", line 3, in train_step
    _, loss, info = autocast_forward(model, criterion, small_img, img)
  File "/tmp/ipykernel_34168/1175262128.py", line 8, in autocast_forward
    g_loss = mean_gradient_error(output, img)
  File "/tmp/ipykernel_34168/3758660440.py", line 14, in mean_gradient_error
    sq_output_gradient_y = torch.square(output_gradient_y)
 (Triggered internally at /opt/conda/conda-bld/pytorch_1712608847532/work/torch/csrc/autograd/python_anomaly_mode.cpp:111.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
</code></pre>
<pre><code class="lang-auto">[...]
File ~/.miniforge3/envs/torch/lib/python3.12/site-packages/torch/autograd/graph.py:744, in _engine_run_backward(t_outputs, *args, **kwargs)
    742     unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
    743 try:
--&gt; 744     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    745         t_outputs, *args, **kwargs
    746     )  # Calls into the C++ engine to run the backward pass
    747 finally:
    748     if attach_logging_hooks:

RuntimeError: Function 'PowBackward0' returned nan values in its 0th output.
</code></pre>
<p>I already tried gradient clipping to norm <code>1.0</code> but it didn’t change anything. I assume that my loss function is somehow numerically instable, but dunno how to stabilize.</p>

AlphaBetaGamma96: <aside class="quote no-group" data-username="MyNameIsArko" data-post="1" data-topic="207316">
<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/m/8e7dd6/48.png" class="avatar"> MyNameIsArko:</div>
<blockquote>
<pre><code class="lang-auto">RuntimeError: Function 'PowBackward0' returned nan values in its 0th output.
</code></pre>
</blockquote>
</aside>
<p>It seems that your function crashes after raising a tensor to a given power. One thing you can do is run your code within a <code>torch.autograd.set_detect_anomaly</code> context manager (docs here: <a href="https://pytorch.org/docs/stable/autograd.html#debugging-and-anomaly-detection" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.4 documentation</a>) and that will raise the line which is causing the <code>NaN</code> value.</p>

MyNameIsArko: <p>I forgot to add that I get this error logs from enabling anomaly detection. In autograd logs, it points to <code>sq_output_gradient_y = torch.square(output_gradient_y)</code> which maps to prints statements as this is where it produced -infinity.</p>

MyNameIsArko: <p>I don’t know why, but adding <code>1e-16</code> to <code>torch.square</code></p>
<p><code>sq_output_gradient_y = torch.square(output_gradient_y + 1e-16)</code></p>
<p>seems to make it stable. I thought that maybe this would be enough to make <code>-inf</code> value go out of <code>-inf</code>, and it worked. (It even works without any gradient clipping)</p>
<p>I’m not marking question as solved, as I don’t understand why it is needed for numerical stability, if squaring any number shouldn’t make it go to negative infinity.</p>

MyNameIsArko: <p>Unfortunately, this “solution” doesn’t work with mixed precision. No matter the <code>weight</code> value nor value to torch.square, it instantly gives nan in the first forward pass.</p>