A package that provides a PyTorch C extension for performing batches of 2D CuFFT transformations, by Eric Wong
Update: FFT functionality is now officially in PyTorch 0.4, see the documentation here. This repository is only useful for older versions of PyTorch, and will no longer be updated.
This package is on PyPi. Install with pip install pytorch-fft.
- From the
pytorch_fft.fftmodule, you can use the following to do foward and backward FFT transformations (complex to complex)fftandifftfor 1D transformationsfft2andifft2for 2D transformationsfft3andifft3for 3D transformations
- From the same module, you can also use the following for
real to complex / complex to real FFT transformations
rfftandirfftfor 1D transformationsrfft2andirfft2for 2D transformationsrfft3andirfft3for 3D transformations
- For an
d-D transformation, the input tensors are required to have >= (d+1) dimensions (n1 x ... x nk x m1 x ... x md) wheren1 x ... x nkis the batch of FFT transformations, andm1 x ... x mdare the dimensions of thed-D transformation.dmust be a number from 1 to 3. - Finally, the module contains the following helper functions you may find
useful
reverse(X, group_size=1)reverses the elements of a tensor and returns the result in a new tensor. Note that PyTorch does not current support negative slicing, see this issue. If a group size is supplied, the elements will be reversed in groups of that size.expand(X, imag=False, odd=True)takes a tensor output of a real 2D or 3D FFT and expands it with its redundant entries to match the output of a complex FFT.
- For autograd support, use the following functions in the
pytorch_fft.fft.autogradmodule:FftandIfftfor 1D transformationsFft2dandIfft2dfor 2D transformationsFft3dandIfft3dfor 3D transformations
# Example that does a batch of three 2D transformations of size 4 by 5.
import torch
import pytorch_fft.fft as fft
A_real, A_imag = torch.randn(3,4,5).cuda(), torch.zeros(3,4,5).cuda()
B_real, B_imag = fft.fft2(A_real, A_imag)
fft.ifft2(B_real, B_imag) # equals (A, zeros)
B_real, B_imag = fft.rfft2(A) # is a truncated version which omits
# redundant entries
reverse(torch.arange(0,6)) # outputs [5,4,3,2,1,0]
reverse(torch.arange(0,6), 2) # outputs [4,5,2,3,0,1]
expand(B_real) # is equivalent to fft.fft2(A, zeros)[0]
expand(B_imag, imag=True) # is equivalent to fft.fft2(A, zeros)[1]# Example that uses the autograd for 2D fft:
import torch
from torch.autograd import Variable
import pytorch_fft.fft.autograd as fft
import numpy as np
f = fft.Fft2d()
invf= fft.Ifft2d()
fx, fy = (Variable(torch.arange(0,100).view((1,1,10,10)).cuda(), requires_grad=True),
Variable(torch.zeros(1, 1, 10, 10).cuda(),requires_grad=True))
k1,k2 = f(fx,fy)
z = k1.sum() + k2.sum()
z.backward()
print(fx.grad, fy.grad)- This follows NumPy semantics and behavior, so
ifft2(fft2(x)) = x. Note that CuFFT semantics for inverse FFT only flip the sign of the transform, but it is not a true inverse. - Similarly, the real to complex / complex to real variants also follow NumPy
semantics and behavior. In the 1D case, this means that for an input of size
N, it returns an output of sizeN//2+1(it omits redundant entries, see the Numpy docs) - The functions in the
pytorch_fft.fftmodule do not implement the PyTorch autogradFunction, and are semantically and functionally like their numpy equivalents. - Autograd functionality is in the
pytorch_fft.fft.autogradmodule.
- pytorch_fft/src: C source code
- pytorch_fft/fft: Python convenience wrapper
- build.py: compilation file
- test.py: tests against NumPy FFTs and Autograd checks
If you have any issues or feature requests, file an issue or send in a PR.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent