Autograd runtime error with operations involving ffts

gbilla: <p>I am working with FFTs with autograd turned on. As I understand, the new <code>pytorch.fft</code> module does not yet support all autograd operations so I am sticking to the old API. I have incorporated some of the suggestions I could find in other related discussion topics but I might have missed something.</p>
<p>I am turning gradients on for an array in Fourier space and doing a bunch of operations including forward and reverse FFTs. The actual code uses 3D FFTs and the update is more complex. Here I show a 1D example that throws no errors and a 2D example throws an error similar to my 3D code:</p>
<blockquote>
<p>RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation:</p>
</blockquote>
<p>1D example (works)</p>
<pre><code class="lang-auto">IFFT_SIZE = (8, )

x_k = torch.rfft(torch.randn(IFFT_SIZE), 1)

# switch on grad and clone
x_k.requires_grad_(True)
x_k_clone = x_k.clone()

# inverse fft
x_r = torch.irfft(x_k_clone, 1, signal_sizes=IFFT_SIZE)

# define some update
dx = 2.0 * torch.rfft(x_r, 1) 

# update
x_k_clone = x_k_clone + dx 

# inverse fft
x_r = torch.irfft(x_k_clone, 1, signal_sizes=IFFT_SIZE)

(x_r * x_r).sum().backward()
</code></pre>
<p>2D example (does not work)</p>
<pre><code class="lang-auto">IFFT_SIZE = (8, 8)

x_r = torch.empty(2, *IFFT_SIZE)
x_k = torch.rfft(torch.randn(2, *IFFT_SIZE), 2)
dx = torch.zeros_like(x_k)

# switch on grad and clone
x_k.requires_grad_(True)
x_k_clone = x_k.clone()

# inverse fft
for i in range(2):
    x_r[i] = torch.irfft(x_k_clone[i], 2, signal_sizes=IFFT_SIZE)

# define some update
dx = 2.0 * torch.stack([torch.rfft(x_r[0], 2), torch.rfft(x_r[1], 2)], dim=0)

# update
x_k_clone = x_k_clone + dx

# inverse fft
for i in range(2):
    x_r[i] = torch.irfft(x_k_clone[i], 2, signal_sizes=IFFT_SIZE)

(x_r * x_r).sum().backward()
</code></pre>
<p>With <code>torch.autograd.set_detect_anomaly(True)</code>, it points to this line as the source of the error<br>
<code>dx = 2.0 * torch.stack([torch.rfft(x_r[0], 2), torch.rfft(x_r[1], 2)], dim=0)</code></p>
<p>I don’t understand how this is an in-place operation. Any suggestions? Thanks!</p>

gbilla: <p>Fixed it. Key takeaways in fixing autograd problems like these:</p>
<ol>
<li>Use <code>torch.autograd.set_detect_anomaly(True)</code> to isolate the line with the problem.</li>
<li>Call <code>.clone()</code> on the variables on the right hand side of the update.</li>
</ol>
<p>had to replace the line:</p>
<pre><code class="lang-auto">dx = 2.0 * torch.stack([torch.rfft(x_r[0], 2), torch.rfft(x_r[1], 2)], dim=0)
</code></pre>
<p>with:</p>
<pre><code class="lang-auto">dx = 2.0 * torch.stack([torch.rfft(x_r[0].clone(), 2), torch.rfft(x_r[1].clone(), 2)], dim=0)
</code></pre>
<p>Found the following threads / blogs useful:</p>
<ul>
<li>
<p><a href="https://www.yongfengli.tk/2018/04/13/inplace-operation-in-pytorch.html" rel="noopener nofollow ugc">https://www.yongfengli.tk/2018/04/13/inplace-operation-in-pytorch.html</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836" class="inline-onebox">Encounter the RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation-torch-floattensor-3-1-which-is-output-0-of-tanhbackward-is-at-version-1-expected-version-0-instead/87630" class="inline-onebox">RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3, 1]], which is output 0 of TanhBackward, is at version 1; expected version 0 instead</a></p>
</li>
</ul>