Does torch.fft.irfft2 break gradient tracking in PyTorch when wrapped into a nn.Module?

spacemeerkat: <p>I have a PyTorch model with a custom forward pass that involves applying <code>torch.fft.irfft2</code> to the real component of a complex input tensor.</p>
<p>I’m wondering whether this operation breaks the gradient tracking through the network during training. Below I have a simple example where when I print <code>output.grad</code>, I’m consistently getting a gradient value of <code>None</code>.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.optim as optim

# Model definition
class DummyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder_conv = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
        self.encoder_relu = nn.ReLU()
        self.decoder_conv = nn.Conv2d(16, 1, kernel_size=3, stride=1, padding=1)
        self.decoder_relu = nn.ReLU()

    def forward(self, x):
        x = self.encoder_conv(x)
        x = self.encoder_relu(x)
        x = torch.fft.irfft2(x, s=(x.shape[-2], x.shape[-1]))
        x = self.decoder_conv(x)
        x = self.decoder_relu(x)
        return x

# Create an instance of the model
model = DummyModel()

# Set up loss function and optimizer
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Create a single input and target
input_image = torch.randn(1, 1, 64, 64)  # Real component of input image
target_image = torch.randn(1, 1, 64, 64)  # Real component of target image

# Training loop
for epoch in range(10):
    optimizer.zero_grad()
    output_image = model(input_image)
    loss = criterion(output_image, target_image)
    loss.backward()
    optimizer.step()

    # Check the gradient of the output
    print(output_image.grad)
</code></pre>
<p>I believe the use of <code>torch.fft</code> is supposed to work with autograd and GPU acceleration but I think I must be handling this in a really stupid way. Any help would be greatly appreciated!</p>

ptrblck: <p>You are receiving a warning, which you shouldn’t ignore:</p>
<pre><code class="lang-python">UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
</code></pre>
<p>Calling:</p>
<pre><code class="lang-python">output_image = model(input_image)
output_image.retain_grad()
</code></pre>
<p>properly shows the gradient in this non-leaf tensor.</p>

spacemeerkat: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> as always!</p>
<p>Just to tie up my understanding… is the fact that using <code>torch.fft.irfft2</code> is different to using an <code>nn.something</code> operation the root cause of pytorch interpreting everything beyond <code>torch.fft.irfft2</code> as no longer a leaf Tensor?</p>
<p>Or have I got that completely wrong? :')</p>

ptrblck: <p>No, non-leaf tensors are tensors which were created by an operation and not explicitly without an Autograd history by the user or as e.g. parameters.<br>
Any operation will raise the same warning:</p>
<pre><code class="lang-python">x = torch.randn(1, 1, requires_grad=True)
print(x.is_leaf)
# True

out = F.relu(x)
print(out.is_leaf)
# False

loss = out.mean()
loss.backward()

print(out.grad)
# UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.
</code></pre>
<p>You don’t need to use <code>nn.Module</code>s and can use any differentiable operation, as also modules are using the same API internally.</p>

spacemeerkat: <p>Ah, I understand! Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> !</p>