Added striding for FFT convolution operators.

Author: fkodom

benchmark.py

@@ -28,19 +28,26 @@ def benchmark(fn: Callable, *args, num_iterations: int = 10, **kwargs) -> Benchm
     return Benchmark(np.mean(times[1:]).item(), np.std(times[1:]).item())
 
 
-def benchmark_conv(signal: Tensor, kernel: Tensor, bias: Tensor, padding: int = 0):
+def benchmark_conv(
+    signal: Tensor, kernel: Tensor, bias: Tensor, padding: int = 0, stride: int = 1
+):
     print(f"Signal size: {signal.shape}")
     print(f"Kernel size: {kernel.shape}")
 
     torch_conv = {1: f.conv1d, 2: f.conv2d, 3: f.conv3d}[signal.ndim - 2]
-    direct_time = benchmark(torch_conv, signal, kernel, bias=bias, padding=padding)
-    fourier_time = benchmark(fft_conv, signal, kernel, bias=bias, padding=padding)
+    direct_time = benchmark(
+        torch_conv, signal, kernel, bias=bias, padding=padding, stride=stride
+    )
+    fourier_time = benchmark(
+        fft_conv, signal, kernel, bias=bias, padding=padding, stride=stride
+    )
     print(f"Direct time: {direct_time}")
     print(f"Fourier time: {fourier_time}")
 
-    y0 = torch_conv(signal, kernel, bias=bias, padding=padding)
-    y1 = fft_conv(signal, kernel, bias=bias, padding=padding)
+    y0 = torch_conv(signal, kernel, bias=bias, padding=padding, stride=stride)
+    y1 = fft_conv(signal, kernel, bias=bias, padding=padding, stride=stride)
     abs_error = torch.abs(y0 - y1)
+    print(f"Output size: {y0.size()}")
     print(f"Abs Error Mean: {abs_error.mean():.3E}")
     print(f"Abs Error Std Dev: {abs_error.std():.3E}")
 

fft_conv.py

@@ -25,7 +25,11 @@ def complex_matmul(a: Tensor, b: Tensor) -> Tensor:
 
 
 def fft_conv(
-    signal: Tensor, kernel: Tensor, bias: Tensor = None, padding: int = 0,
+    signal: Tensor,
+    kernel: Tensor,
+    bias: Tensor = None,
+    padding: int = 0,
+    stride: int = 1,
 ) -> Tensor:
     """Performs N-d convolution of Tensors using a fast fourier transform, which
     is very fast for large kernel sizes. Also, optionally adds a bias Tensor after
@@ -36,6 +40,7 @@ def fft_conv(
         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.
+        stride: (int) Stride size for computing output values.
 
     Returns:
         (Tensor) Convolved tensor
@@ -44,7 +49,8 @@ def fft_conv(
     signal_padding = (signal.ndim - 2) * [padding, padding]
     signal = f.pad(signal, signal_padding)
     kernel_padding = [
-        pad for i in reversed(range(2, signal.ndim))
+        pad
+        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)
@@ -58,8 +64,9 @@ def fft_conv(
     output = irfftn(output_fr, dim=tuple(range(2, signal.ndim)))
 
     # Remove extra padded values
-    crop_slices = [slice(0, output.shape[0]), slice(0, output.shape[1])] + [
-        slice(0, (signal.size(i) - kernel.size(i) + 1)) for i in range(2, signal.ndim)
+    crop_slices = [slice(0, output.size(0)), slice(0, output.size(1))] + [
+        slice(0, (signal.size(i) - kernel.size(i) + 1), stride)
+        for i in range(2, signal.ndim)
     ]
     output = output[crop_slices].contiguous()
 
@@ -80,6 +87,7 @@ def __init__(
         out_channels: int,
         kernel_size: int,
         padding: int = 0,
+        stride: int = 1,
         bias: bool = True,
     ):
         """
@@ -88,13 +96,15 @@ def __init__(
             out_channels: (int) Number of channels in output tensors
             kernel_size: (int) Square radius of the convolution kernel
             padding: (int) Amount of zero-padding to add to the input tensor
+            stride: (int) Stride size for computing output values
             bias: (bool) If True, includes bias, which is added after convolution
         """
         super().__init__()
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
         self.padding = padding
+        self.stride = stride
         self.use_bias = bias
 
         self.weight = torch.empty(0)
@@ -106,6 +116,7 @@ def forward(self, signal):
             self.weight,
             bias=self.bias,
             padding=self.padding,
+            stride=self.stride,
         )
 
 
@@ -116,14 +127,11 @@ def __init__(
         out_channels: int,
         kernel_size: int,
         padding: int = 0,
+        stride: int = 1,
         bias: bool = True,
     ):
         super().__init__(
-            in_channels,
-            out_channels,
-            kernel_size,
-            padding=padding,
-            bias=bias,
+            in_channels, out_channels, kernel_size, padding=padding, bias=bias,
         )
         self.weight = nn.Parameter(torch.randn(out_channels, in_channels, kernel_size))
 
@@ -135,13 +143,15 @@ def __init__(
         out_channels: int,
         kernel_size: int,
         padding: int = 0,
+        stride: int = 1,
         bias: bool = True,
     ):
         super().__init__(
             in_channels,
             out_channels,
             kernel_size,
             padding=padding,
+            stride=stride,
             bias=bias,
         )
         self.weight = nn.Parameter(
@@ -150,20 +160,21 @@ def __init__(
 
 
 class FFTConv3d(_FFTConv):
-
     def __init__(
         self,
         in_channels: int,
         out_channels: int,
         kernel_size: int,
         padding: int = 0,
+        stride: int = 1,
         bias: bool = True,
     ):
         super().__init__(
             in_channels,
             out_channels,
             kernel_size,
             padding=padding,
+            stride=stride,
             bias=bias,
         )
         self.weight = nn.Parameter(