torch-fftconv 1.0.0

Implementation of 1D, 2D, and 3D FFT convolutions in PyTorch. Much faster than direct convolutions for large kernel sizes

FFT Conv PyTorch

This is a fork of fft-conv-pytorch. I made some modifications to support dilated, strided, and transposed convolution, so it can be a drop-in-replacement of original PyTorch Conv*d modules and conv*d functions, with the same behaviour and interface.

Install

pip install torch-fftconv

Example Usage

import torch
from torch_fftconv import fft_conv1d, FFTConv1d

signal = torch.randn(3, 3, 1024 * 1024)
kernel = torch.randn(2, 3, 128)
bias = torch.randn(2)

# Functional execution.  (Easiest for generic use cases.)
out = fft_conv1d(signal, kernel, bias=bias)

# Object-oriented execution.
fft_conv = FFTConv1d(3, 2, 128, bias=True)
out = fft_conv(signal)

# Suppose you want to replace all the existing convolution layers in a neural network with FFTConv, you can do it like this:
m: torch.nn.Module = ...  # some existing model
from torch_fftconv.utils import convert_fft_conv
m = convert_fft_conv(m)  # this will replace all Conv/TransposedConv layers in m with FFTConv layers, and copy the weights and bias.

Benchmarks

The best situation to use FFTConv is when using large size kernel. The following image shows that when the size of input is fixed, the fft method remains an almost constant cost among all size of kernel, regardless.

For details and benchmarks on other parameters, check this notebook.

Features

• [ ] Jittability.
• Dilated Convolution.
• Transposed Convolution.
• Support complex input and kernel.