How to avoid nan output from atan2 during backward pass?

iansampson: <p>Hi everyone, first post here!</p>
<p>I’m training a model that uses a single atan2 in its final stage. I’ve enabled anomaly detection with <code>torch.autograd.set_detect_anomaly(True)</code>, and after a while I get an error that “[f]unction ‘Atan2Backward0’ returned nan values in its 0th output.” I know that atan2 can produce nan if both the nominator and denominator are 0, so I add a small epsilon value to any denominator values that are zero:</p>
<pre><code class="lang-auto">    # Calculate the phase of a complex spectrogram
    def find_phase(self, subbands: torch.Tensor) -&gt; torch.Tensor:
        numerator = subbands[:, 1:2, :, :]
        denominator = subbands[:, 0:1, :, :]
        # Add epsilon to denominator to avoid nan
        epsilon = 1e-7
        nudge = (denominator == 0) * epsilon
        denominator = denominator + nudge
        return torch.atan2(numerator, denominator)
</code></pre>
<p>Strangely, though, the error still occurs. I wonder if I’m missing something here, and if there’s a better way to ensure that atan2 doesn’t return nan. Thanks!</p>
<p>Also, I’m training on a Nvidia A100 with PyTorch 2.0.0 + CUDA 11.7, if that’s any help.</p>

ptrblck: <p>Your approach sounds valid and also works in this minimal code snippet:</p>
<pre><code class="lang-python">numerator = torch.randn(10, 10)
numerator[0, 0] = 0.
numerator.requires_grad_(True)
denominator = torch.randn(10, 10)
denominator[0, 0] = 0.

epsilon = 1e-7
nudge = (denominator == 0) * epsilon
denominator = denominator + nudge
out =  torch.atan2(numerator, denominator)
out.mean().backward()
print(numerator.grad)
</code></pre>
<p>Could you check which values are set to NaN in the backward pass and check the inputs to this method?</p>

iansampson: <p>Thanks, I really appreciate the help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>So I iterated over all of the model parameters after calling <code>loss.backward()</code> and checked for NaN in each of the gradients and weights.</p>
<pre><code class="lang-auto">        # Check for non-finite parameters
        all_gradients_are_finite = True
        for parameter_name, parameter in self.model.named_parameters():
            if parameter.grad is not None:
                if not parameter.grad.isfinite().all():
                    print("Parameter gradient is non-finite: {}".format(parameter_name))
                    all_gradients_are_finite = False
            elif parameter.data is not None:
                if not parameter.data.isfinite().all():
                    print("Parameter data is non-finite: {}".format(parameter_name))
</code></pre>
<p>None of the weights contain NaN (probably because my training loop avoids calling <code>optimizer.step()</code> if any values are NaN) but <em>all</em> of the gradients do. The first layer with a NaN gradient in the backward pass is the last trainable layer of the model (Conv2d). This makes sense, I suppose, because the atan2 call in question appears after that final layer: it’s part of an output stage where I convert the estimated spectrogram back to a waveform. So if atan2 returns NaN in the backward pass it would propagate to the whole model.</p>
<p>I checked the inputs to the <code>find_phase</code> method and they don’t contain NaN at all during the forward pass. The loss doesn’t contain NaN either (as long as I don’t call <code>optimizer.step()</code> when NaN gradients are detected).</p>
<p>I also tried removing the <code>find_phase</code> method and the NaNs disappeared. So tan2 does seem to be the culprit. I’m a little ignorant about how autograd works, to be honest, but I wonder if adding epsilon here only guarantees a finite output from atan2 during the forward pass, and it’s still possible for the backward pass to feed it a zero-valued numerator and denominator?</p>

ptrblck: <p>Your observation indeed points towards <code>atan2</code> and you can find the gradient definition <a href="https://github.com/pytorch/pytorch/blob/9fd3eba6ceb048cfdcb430e34f9168eda888b4c8/tools/autograd/derivatives.yaml#L331-L333">here</a>.</p>
<aside class="quote no-group" data-username="iansampson" data-post="3" data-topic="176890">
<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/i/b5a626/48.png" class="avatar"> iansampson:</div>
<blockquote>
<p>but I wonder if adding epsilon here only guarantees a finite output from atan2 during the forward pass, and it’s still possible for the backward pass to feed it a zero-valued numerator and denominator?</p>
</blockquote>
</aside>
<p>This shouldn’t happen as my previous code snippet would only fail without the <code>eps</code>, which also matches the gradient formula.</p>
<p>Could you add a check into the <code>forward</code> method checking for zero outputs e.g. via:</p>
<pre data-code-wrap="python"><code class="lang-python">if (out == 0.0).any():
    print(f"zero output detected at {(out==0.0).nonzero()}")
</code></pre>

iansampson: <p>I wasn’t sure exactly where to add this check, but I tried directly after <code>atan2</code> in the <code>find_phase</code> method as well as at the output of the model itself. Both checks did return zero outputs now and then, although this is to be expected since the model ultimately returns a waveform, and <code>find_phase</code> returns the phase of a complex spectrogram, both of which may contain zero crossings, silent sections, etc. Could zero outputs cause <code>atan2</code> to return NaN in the backward pass, I wonder?</p>

ptrblck: <aside class="quote no-group" data-username="iansampson" data-post="5" data-topic="176890">
<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/i/b5a626/48.png" class="avatar"> iansampson:</div>
<blockquote>
<p>Could zero outputs cause <code>atan2</code> to return NaN in the backward pass, I wonder?</p>
</blockquote>
</aside>
<p>Based on this code, yes I think so:</p>
<pre data-code-wrap="python"><code class="lang-python">numerator = torch.randn(10, 10)
numerator[0, 0] = 0.
numerator.requires_grad_(True)
denominator = torch.randn(10, 10)
denominator[0, 0] = 0.

# epsilon = 1e-10000
# nudge = (denominator == 0) * epsilon
# denominator = denominator + nudge
out =  torch.atan2(numerator, denominator)
out.mean().backward()
print(numerator.grad)
# tensor([[        nan,  4.0503e-04, -3.5165e-03,  5.8801e-03,  1.3173e-02,
#          -2.1934e-02, -3.9802e-03, -1.4191e-03, -2.0235e-03,  7.3573e-03],

if (out == 0.0).any():
    print(f"zero output detected at {(out==0.0).nonzero()}")
# zero output detected at tensor([[0, 0]])
</code></pre>
<p>but note that both inputs are also set to zero.<br>
Given you are adding a small <code>eps</code> value to the <code>denominator</code> this should not happen (as seen in my previous code), but still something is causing the issue which is why I suggested to now dig a bit into the actual values.</p>

iansampson: <p>I figured it out! By digging into the actual values, as you suggested. My <code>find_phase</code> method above replaces zero values with <code>eps</code>, yes, but other values that are very close to zero (e.g. <code>1e-20</code>) slip through unchanged, and become zero during the backward pass via rounding errors. The solution:</p>
<pre><code class="lang-auto">def find_phase(self, subbands: torch.Tensor) -&gt; torch.Tensor:
    numerator = subbands[:, 1:2, :, :]
    denominator = subbands[:, 0:1, :, :]
    # Add epsilon to denominator to avoid NaN in backward pass
    epsilon = 1e-10
    near_zeros = denominator &lt; epsilon
    denominator = denominator * (near_zeros.logical_not())
    denominator = denominator + (near_zeros * epsilon)
    return torch.atan2(numerator, denominator)
</code></pre>
<p>This replaces any values <em>less than</em> <code>eps</code> with <code>eps</code>. And no more <code>NaN</code> or <code>inf</code>! Autograd no longer detects any anomalies.</p>
<p>Here’s a minimal example that demonstrates the problem:</p>
<pre><code class="lang-auto">numerator = torch.tensor([1e-20], requires_grad=True)
denominator = torch.tensor([1e-20], requires_grad=True)
out =  torch.atan2(numerator, denominator)
out.mean().backward()
print(numerator.grad) # tensor([inf])
print(denominator.grad) # tensor([-inf])
</code></pre>
<p>Anyway, all fixed now. Thanks again for pointing me in the right direction! Been scratching my head for hours over this one and very relieved to see it resolved.</p>

ptrblck: <p>Ah, of course! I have even experimented with tiny values but was blind.<br>
Great to hear you have figured it out. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

arrowhyx: <pre><code class="lang-python">near_zeros = denominator &lt; epsilon
</code></pre>
<p>The code in the previous would be better modified as follows</p>
<pre><code class="lang-python">near_zeros = denominator.abs() &lt; epsilon
</code></pre>