Fixed bug for inputs with odd length in the final dimension.

Author: fkodom

benchmark.py

@@ -1,4 +1,4 @@
-from typing import Callable, NamedTuple
+from typing import Callable, NamedTuple, Union, Iterable
 from timeit import Timer
 
 import torch
@@ -29,7 +29,11 @@ def benchmark(fn: Callable, *args, num_iterations: int = 10, **kwargs) -> Benchm
 
 
 def benchmark_conv(
-    signal: Tensor, kernel: Tensor, bias: Tensor, padding: int = 0, stride: int = 1
+    signal: Tensor,
+    kernel: Tensor,
+    bias: Tensor,
+    padding: Union[int, Iterable[int]] = 0,
+    stride: Union[int, Iterable[int]] = 1,
 ):
     print(f"Signal size: {signal.shape}")
     print(f"Kernel size: {kernel.shape}")
@@ -54,24 +58,27 @@ def benchmark_conv(
 
 print("\n--- 1D Convolution ---")
 benchmark_conv(
-    signal=torch.randn(3, 3, 4096),
+    signal=torch.randn(3, 3, 4091),
     kernel=torch.randn(2, 3, 1025),
     bias=torch.randn(2),
     padding=512,
+    stride=3,
 )
 
 print("\n--- 2D Convolution ---")
 benchmark_conv(
-    signal=torch.randn(3, 3, 256, 256),
-    kernel=torch.randn(2, 3, 21, 21),
+    signal=torch.randn(3, 3, 256, 235),
+    kernel=torch.randn(2, 3, 19, 21),
     bias=torch.randn(2),
-    padding=10,
+    padding=(9, 10),
+    stride=(2, 3),
 )
 
 print("\n--- 3D Convolution ---")
 benchmark_conv(
-    signal=torch.randn(3, 3, 64, 64, 64),
-    kernel=torch.randn(2, 3, 9, 9, 9),
+    signal=torch.randn(3, 3, 64, 72, 61),
+    kernel=torch.randn(2, 3, 5, 7, 9),
     bias=torch.randn(2),
-    padding=4,
+    padding=(2, 3, 4),
+    stride=(1, 2, 3)
 )

fft_conv.py

@@ -75,15 +75,23 @@ def fft_conv(
     # Pad the input signal & kernel tensors
     signal_padding = [p for p in padding_[::-1] for _ in range(2)]
     signal = f.pad(signal, signal_padding)
+
+    # Because PyTorch computes a *one-sided* FFT, we need the final dimension to
+    # have *even* length.  Just pad with one more zero if the final dimension is odd.
+    if signal.size(-1) % 2:
+        signal_ = f.pad(signal, [0, 1])
+    else:
+        signal_ = signal
+
     kernel_padding = [
         pad
-        for i in reversed(range(2, signal.ndim))
-        for pad in [0, signal.size(i) - kernel.size(i)]
+        for i in reversed(range(2, signal_.ndim))
+        for pad in [0, signal_.size(i) - kernel.size(i)]
     ]
     padded_kernel = f.pad(kernel, kernel_padding)
 
     # Perform fourier convolution -- FFT, matrix multiply, then IFFT
-    signal_fr = rfftn(signal, dim=tuple(range(2, signal.ndim)))
+    signal_fr = rfftn(signal_, dim=tuple(range(2, signal.ndim)))
     kernel_fr = rfftn(padded_kernel, dim=tuple(range(2, signal.ndim)))
 
     kernel_fr.imag *= -1