Cannot learn through single FFT call

username: F10w10
cooked: <p>Hello,</p>
<p>while playing around with a model that will feature calls to the <code>fft</code> functions, I have noticed something odd about the behavior of the gradient. Basically, I cannot do a basic gradient descent when I have exact target data. My starting point is some volumetric data in the shape <code>[1, size, size, size]</code>, so three dimensional, with an additional dimension for batch size.<br>
Now if I start with a zeroed volume, then do a <code>rfftn</code> call and then calculate an error based on the FFT of my target volume, I end up with no learning whatsoever. Here is my sample code:</p>
<pre><code class="lang-auto">import torch
import torch.fft
import torch.optim

import numpy as np

cuda0 = torch.device('cuda:0')

size = 178

#Target volume
tensorImage = torch.zeros([1, size, size, size], dtype=torch.float32)
for x in range(178):
   # Circle
   if pow(x - size/2,2) &gt; 10:
       continue
   for y in range(178):
       if pow(y - size/2,2) &gt; 10:
           continue
       for z in range(178):
           if pow(z - size/2,2) &gt; 10:
               continue
           tensorImage[0,z,y,x] = 1

           
tensorImage = tensorImage.to(cuda0)
tensorFFT = torch.fft.rfftn(tensorImage, dim=(1,2,3))

#Starting volume, with only 0s
tensorZeroImage = torch.zeros(tensorImage.size(), requires_grad=True, device=cuda0)
optim = torch.optim.SGD([tensorZeroImage], lr=0.01)

for i in range(1000):
   optim.zero_grad()                     
   tensorFFTZero = torch.fft.rfftn(tensorZeroImage, dim=(1,2,3))
   tensorFFTDiff = tensorFFTZero - tensorFFT
   tensorFFTDiffAbs = tensorFFTDiff.abs()
   tensorFFTDiffAbsSqrd = tensorFFTDiffAbs
   error = tensorFFTDiffAbsSqrd.sum()

   error.backward()
   optim.step()
   if (i+1)%100==0:
       print(error)
</code></pre>
<p>Output is</p>
<pre><code class="lang-auto">tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
tensor(8.2181e+10, device='cuda:0', grad_fn=&lt;SumBackward0&gt;)
</code></pre>
<p>So no learning whatsoever. There seems to be something odd about the backpropagation through the <code>rfftn</code> call (or <code>fftn</code>).</p>
<p>If I add an <code>ifftn</code> call and then compare the two volumes directly, I get correct learning as expected.</p>
<p>Is there any explanation, why the code above with an error calculation on the fourier transform does not work as expected?</p>

username: ptrblck
cooked: <p>I don’t think the backpropagation is wrong, but your loss is just extremely high, which would also create large gradients and could thus break the training.<br>
If you use a <code>mean()</code> as the error, you’ll see that (some) training is happening:</p>
<pre><code class="lang-python">tensor(5.0114, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(4.5654, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(4.2215, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.9387, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.7002, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.4965, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.3204, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.1662, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(3.0293, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
tensor(2.9056, device='cuda:0', grad_fn=&lt;MeanBackward0&gt;)
</code></pre>

username: F10w10
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
thank you for your quick reply. You are correct with the loss and therefore gradient being too large for any good training. With enough iterations, it actually recovers the original image. There is still something odd I am noticing about the exact gradient in the first iteration. I have an extended sample code below:</p>
<pre><code class="lang-auto">import torch
import torch.fft
import torch.optim

import numpy as np

cuda0 = torch.device('cuda:0')

size = 178

#Target volume
tensorImage = torch.zeros([1, size, size, size], dtype=torch.float32)
for x in range(178):
    # Circle
    if pow(x - size/2,2) &gt; 10:
        continue
    for y in range(178):
        if pow(y - size/2,2) &gt; 10:
            continue
        for z in range(178):
            if pow(z - size/2,2) &gt; 10:
                continue
            tensorImage[0,z,y,x] = 1

            
tensorImage = tensorImage.to(cuda0)
tensorFFT = torch.fft.rfftn(tensorImage, dim=(1,2,3))

#Starting volume, with only 0s
tensorZeroImage = torch.zeros(tensorImage.size(), requires_grad=True, device=cuda0)
               
tensorFFTZero = torch.fft.rfftn(tensorZeroImage, dim=(1,2,3))
tensorFFTDiff = tensorFFTZero - tensorFFT
tensorFFTDiffAbs = tensorFFTDiff.abs()
tensorFFTDiffAbsSqrd = tensorFFTDiffAbs.pow(2)
error = tensorFFTDiffAbsSqrd.mean()

error.backward()
gradientFromFFT = tensorZeroImage.grad
np_gradientFromFFT = gradientFromFFT.cpu().detach().numpy()

tensorZeroImage = torch.zeros(tensorImage.size(), requires_grad=True, device=cuda0)
               
tensorFFTZero = torch.fft.rfftn(tensorZeroImage, dim=(1,2,3))
tensorIFFTZero = torch.fft.irfftn(tensorFFTZero, dim=(1,2,3))
tensorIFFTDiff = tensorIFFTZero - tensorImage
tensorIFFTDiffSqrd = tensorIFFTDiff.pow(2)
error = tensorIFFTDiffSqrd.mean()

error.backward()
gradientFromIFFT = tensorZeroImage.grad
np_gradientFromIFFT = gradientFromIFFT.cpu().detach().numpy()

import matplotlib.pyplot as plt
plt.figure()
plt.imshow(np_gradientFromFFT[0,89,:,:])
plt.figure()
plt.imshow(np_gradientFromIFFT[0,89,:,:])
plt.show()
</code></pre>
<p>This is the result:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/1/118a2266bbe3ae6191282c150b0ad0a54fcb516f.png" data-download-href="https://discuss.pytorch.org/uploads/default/118a2266bbe3ae6191282c150b0ad0a54fcb516f" title="2021-06-28 10_17_58-Untitled - Jupyter Notebook"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/118a2266bbe3ae6191282c150b0ad0a54fcb516f_2_148x500.png" alt="2021-06-28 10_17_58-Untitled - Jupyter Notebook" data-base62-sha1="2va4ftuJ4kLOc7gwqJCMpzmLYnB" width="148" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/118a2266bbe3ae6191282c150b0ad0a54fcb516f_2_148x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/118a2266bbe3ae6191282c150b0ad0a54fcb516f_2_222x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/118a2266bbe3ae6191282c150b0ad0a54fcb516f_2_296x1000.png 2x" data-dominant-color="C4AF8E"><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">2021-06-28 10_17_58-Untitled - Jupyter Notebook</span><span class="informations">325×1094 46.3 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>There is an odd looking artifact when calculating the loss based on the difference in fourier space compared to adding an additional <code>irfftn</code> call.</p>