PyTorch中的傅立叶卷积：通过FFT有效计算大核卷积的数学原理和代码实现

 def fft_conv(
     signal: Tensor, kernel: Tensor, bias: Tensor = None, padding: int = 0,
 ) -> Tensor:
     # 1. Pad the input signal & kernel tensors
     # 2. Compute FFT for both signal & kernel
     # 3. Multiply the transformed Tensors together
     # 4. Compute inverse FFT
     # 5. Add bias and return

 # 1. Pad the input signal & kernel tensors
 signal = f.pad(signal, [padding, padding])
 kernel_padding = [0, signal.size(-1) - kernel.size(-1)]
 padded_kernel = f.pad(kernel, kernel_padding)

 # 2. Perform fourier convolution
 signal_fr = rfftn(signal, dim=-1)
 kernel_fr = rfftn(padded_kernel, dim=-1)

 # 3. Multiply the transformed matrices
 
 def complex_matmul(a: Tensor, b: Tensor) -> Tensor:
     """Multiplies two complex-valued tensors."""
     # Scalar matrix multiplication of two tensors, over only the first two dimensions.
     # Dimensions 3 and higher will have the same shape after multiplication.
     scalar_matmul = partial(torch.einsum, "ab..., cb... -> ac...")
 
     # Compute the real and imaginary parts independently, then manually insert them
     # into the output Tensor.  This is fairly hacky but necessary for PyTorch 1.7.0,
     # because Autograd is not enabled for complex matrix operations yet.  Not exactly
     # idiomatic PyTorch code, but it should work for all future versions (>= 1.7.0).
     real = scalar_matmul(a.real, b.real) - scalar_matmul(a.imag, b.imag)
     imag = scalar_matmul(a.imag, b.real) + scalar_matmul(a.real, b.imag)
     c = torch.zeros(real.shape, dtype=torch.complex64)
     c.real, c.imag = real, imag
 
     return c
 
 
 # Conjugate the kernel for cross-correlation
 kernel_fr.imag *= -1
 output_fr = complex_matmul(signal_fr, kernel_fr)

 # 4. Compute inverse FFT, and remove extra padded values
 output = irfftn(output_fr, dim=-1)
 output = output[:, :, :signal.size(-1) - kernel.size(-1) + 1]

 # 5. Optionally, add a bias term before returning.
 if bias is not None:
     output += bias.view(1, -1, 1)

 def fft_conv_1d(
     signal: Tensor, kernel: Tensor, bias: Tensor = None, padding: int = 0,
 ) -> Tensor:
     """
     Args:
         signal: (Tensor) Input tensor to be convolved with the kernel.
         kernel: (Tensor) Convolution kernel.
         bias: (Optional, Tensor) Bias tensor to add to the output.
         padding: (int) Number of zero samples to pad the input on the last dimension.
     Returns:
         (Tensor) Convolved tensor
     """
     # 1. Pad the input signal & kernel tensors
     signal = f.pad(signal, [padding, padding])
     kernel_padding = [0, signal.size(-1) - kernel.size(-1)]
     padded_kernel = f.pad(kernel, kernel_padding)
 
     # 2. Perform fourier convolution
     signal_fr = rfftn(signal, dim=-1)
     kernel_fr = rfftn(padded_kernel, dim=-1)
 
     # 3. Multiply the transformed matrices
     kernel_fr.imag *= -1
     output_fr = complex_matmul(signal_fr, kernel_fr)
 
     # 4. Compute inverse FFT, and remove extra padded values
     output = irfftn(output_fr, dim=-1)
     output = output[:, :, :signal.size(-1) - kernel.size(-1) + 1]
 
     # 5. Optionally, add a bias term before returning.
     if bias is not None:
         output += bias.view(1, -1, 1)
 
     return output

 import torch
 import torch.nn.functional as f
 
 
 torch.manual_seed(1234)
 kernel = torch.randn(2, 3, 1025)
 signal = torch.randn(3, 3, 4096)
 bias = torch.randn(2)
 
 y0 = f.conv1d(signal, kernel, bias=bias, padding=512)
 y1 = fft_conv_1d(signal, kernel, bias=bias, padding=512)
 
 abs_error = torch.abs(y0 - y1)
 print(f'\nAbs Error Mean: {abs_error.mean():.3E}')
 print(f'Abs Error Std Dev: {abs_error.std():.3E}')
 
 # Abs Error Mean: 1.272E-05
 # Abs Error Std Dev: 9.937E-06

from timeit import timeit


direct_time = timeit(
    "f.conv1d(signal, kernel, bias=bias, padding=512)", 
    globals=locals(), 
    number=100
) / 100
fourier_time = timeit(
    "fft_conv_1d(signal, kernel, bias=bias, padding=512)", 
    globals=locals(), 
    number=100
) / 100
print(f"Direct time: {direct_time:.3E} s")
print(f"Fourier time: {fourier_time:.3E} s")

# Direct time: 1.523E-02 s
# Fourier time: 1.149E-03 s