Equivalent to torch.rfft() in newest PyTorch version

Viewed 1132

I want to estimate the fourier transform for a given image of size BxCxWxH

In previous torch version the following did the job:

fft_im = torch.rfft(img, signal_ndim=2, onesided=False)

and the output was of size:

BxCxWxHx2

However, with the new version of rfft :

fft_im = torch.fft.rfft2(img, dim=2, norm=None)

I do not get the same results. Do I miss something?

1 Answers

A few issues

  1. The dim argument you provided is an invalid type, it should be a tuple of two numbers or should be omitted. Really PyTorch should raise an exception. I would argue that the fact this ran without exception is a bug in PyTorch (I opened a ticket stating as much).
  2. PyTorch now supports complex tensor types, so FFT functions return those instead of adding a new dimension for the real/imaginary parts. You can use torch.view_as_real to convert to the old representation. Also worth pointing out that view_as_real doesn't copy data since it returns a view so shouldn't slow things down in any noticeable way.
  3. PyTorch no longer gives the option of disabling one-sided calculation in RFFT. Probably because disabling one-sided makes the result identical to torch.fft.fft2, which is in conflict with the 13th aphorism of PEP 20. The whole point of providing a special real-valued version of the FFT is that you need only compute half the values for each dimension, since the rest can be inferred via the Hermition symmetric property.

So from all that you should be able to use

fft_im = torch.view_as_real(torch.fft.fft2(img))

Important If you're going to pass fft_im to other functions in torch.fft (like fft.ifft or fft.fftshift) then you'll need to convert back to the complex representation using torch.view_as_complex so those functions don't interpret the last dimension as a signal dimension.

Related