Strange behavior of Torch.Angle()'s AngleBackward

MinkyuChoi: <p>I am using <code>Torch.Angle()</code> as a part of my model. And when I train my model, I always faces Nan issue on gradients. The errors are as follows.</p>
<blockquote>
<p><strong>[W python_anomaly_mode.cpp:104] Warning: Error detected in AngleBackward. Traceback of forward call that caused the error:</strong><br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/threading.py”, line 930, in _bootstrap<br>
self._bootstrap_inner()<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/threading.py”, line 973, in _bootstrap_inner<br>
self.run()<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/threading.py”, line 910, in run<br>
self._target(*self._args, **self._kwargs)<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py”, line 61, in _worker<br>
output = module(*input, **kwargs)<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1051, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/train.py”, line 106, in forward<br>
pred, _ = self.cls_net(img_s)<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1051, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/model_fourier1_magPhaseCorr_noBN_noSA_f11l_fixPhaseMag_noConjPad_scale.py”, line 137, in forward<br>
out1_i, f1_res = self.conv1_2(out1_res)<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1051, in _call_impl<br>
return forward_call(*input, **kwargs)<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/model_fourier1_magPhaseCorr_noBN_noSA_f11l_fixPhaseMag_noConjPad_scale.py”, line 64, in forward<br>
f_conv_i = f_scale.angle() #(self.sigmoid(self.conv_i(f_pad_pos_concat))-0.5)<em>2</em>3.141592<br>
(function _print_stack)<br>
Traceback (most recent call last):<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/train.py”, line 292, in <br>
main()<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/train.py”, line 83, in main<br>
train(args, train_loader, test_loader)<br>
File “/home/choi574/research_mk/fourier_cnn/cnn_fourier1_magPhaseCorr_noBN_noSA_f11_fixPhaseMag_noConjPad_scale/train.py”, line 230, in train<br>
loss_b.backward()<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/_tensor.py”, line 255, in backward<br>
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)<br>
File “/home/choi574/.conda/envs/pytorch_mk_190/lib/python3.9/site-packages/torch/autograd/<strong>init</strong>.py”, line 147, in backward<br>
Variable._execution_engine.run_backward(<br>
<strong>RuntimeError: Function ‘AngleBackward’ returned nan values in its 0th output.</strong></p>
</blockquote>
<p>Always the AngleBackward is causing an issue of Nan. So, I simplified my code as follows. In the code, there is <code>Torch.Angle()</code> between the <code>FFT</code> and <code>iFFT</code>. Basically, the block the <code>Torch.Angle()</code> is included does not add any new information. It just divides the frequency into magnitude and phase, and then merge them into frequency domain (complex number) again. So, commenting the block including <code>Torch.Angle()</code> does not make output different (please see the code below for details).</p>
<pre><code class="lang-auto">    class Model(torch.nn.Module):
      def __init__(self):
    
        self.conv = nn.Conv2d(3, 64)
        self.relu = nn.ReLU()
  
      def forward(self, x):
        '''
        x: image
        '''
    
        out = self.relu(self.conv(x))
    
        ## FFT
        fft2 = torch.fft.fftshift(torch.fft.fft2(x), [-2, -1])
    

        ######### This block does nothing in overall #########
        ## When this block is commented, there is no Nan issue in training. 
        ## When this block is not commented, there is Nan issue in Angle(). 
        a = fft2.angle()
        m = fft2.abs()
        
        dum = torch.zeros_like(a).unsqueeze(-1)
        aj = torch.view_as_complex(torch.cat((dum, a.unsqueeze(-1)), -1))
        aj_exp = torch.exp(aj)
        x_res = m * aj_exp
        ########################################################
    
        ## iFFT
        ifft2 = torch.fft.ifft2(torch.fft.ifftshift(x_res, [-2, -1]))
    
        return ifft2
</code></pre>
<p>Seeing the fact that the Nan issue always occurs from <code>Torch.Angle()</code>, I tried to search if there is similar case in the Internet. And there is a previous discussion with the same topic (<a href="https://discuss.pytorch.org/t/anglebackward-returned-nan-values/122336/3" class="inline-onebox">'AngleBackward' returned nan values - #3 by eagomez</a>). Basically, the solution from the discussion is to eliminate the (near) zero values from the input to the <code>Torch.Angle()</code>, which is <code>fft2</code> in my code. I think it is due to avoid dividing by zero inside <code>Torch.Angle()</code>. When I use the suggested method from the discussion, my Nan issue is gone.</p>
<p>Therefore, it seems the <code>Torch.Angle()</code> returns Nan value for some cases, but this issue is gone when the suggested solution is applied to the input of the <code>Torch.Angle()</code>.</p>
<p>I wonder why <code>Torch.Angle()</code> returns Nan for some cases and why the suggested solution solves the Nan issue. The code above is used as a part of a simple classification model using ImageNet. I tried to manually make <code>AngleBackward</code> to return Nan value, but I could not. It only occurs when the input is meets certain conditions.<br>
Could anyone help me to understand why and how <code>AngleBackward</code> returns Nan?</p>

miccio-dk: <p>I am also stumbling into this issue:<br>
<code>Function 'AngleBackward0' returned nan values in its 0th output.</code></p>
<p>I believe that the source of NaNs is to be found <a href="https://github.com/pytorch/pytorch/blob/v1.12.1/torch/csrc/autograd/FunctionsManual.cpp#L373" rel="noopener nofollow ugc">here</a> but I can’t quite determine what combination of inputs (that are not already nans) would cause this.</p>