Hi @ptrblck
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:
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) > 10:
        continue
    for y in range(178):
        if pow(y - size/2,2) > 10:
            continue
        for z in range(178):
            if pow(z - size/2,2) > 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()

This is the result:

2021-06-28 10_17_58-Untitled - Jupyter Notebook325×1094 46.3 KB

There is an odd looking artifact when calculating the loss based on the difference in fourier space compared to adding an additional irfftn call.