Merge pull request #16 from fkodom/bug-fix/unit-test-utils

Bug Fix and Torch Compatibility

Author: fkodom

.github/workflows/test.yaml

@@ -13,7 +13,7 @@ jobs:
     strategy:
       matrix:
         python: ["3.7", "3.8", "3.9"]
-        torch: ["1.7", "1.8", "1.9", "1.10"]
+        torch: ["1.8", "1.9", "1.10"]
 
     steps:
       - name: Checkout
@@ -26,7 +26,7 @@ jobs:
 
       - name: Install Package
         run: |
-          pip install torch==${{ matrix.torch }}
+          pip install torch~=${{ matrix.torch }}.0
           pip install .[test]
 
       - name: Test

doc/scripts/generate_benchmark_plot.py

@@ -1,6 +1,5 @@
 from functools import lru_cache, partial
-from timeit import Timer
-from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Sequence, Union
+from typing import Dict, Iterable, List, Optional, Sequence, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -9,25 +8,7 @@
 from tqdm import tqdm
 
 from fft_conv_pytorch.fft_conv import fft_conv, to_ntuple
-
-
-class Benchmark(NamedTuple):
-    mean: float
-    std: float
-
-    def __repr__(self):
-        return f"BenchmarkResult(mean: {self.mean:.3e}, std: {self.std:.3e})"
-
-    def __str__(self):
-        return f"({self.mean:.3e} \u00B1 {self.std:.3e}) s"
-
-
-def benchmark(fn: Callable, *args, num_iterations: int = 10, **kwargs) -> Benchmark:
-    timer = Timer(
-        "fn(*args, **kwargs)", globals={"fn": fn, "args": args, "kwargs": kwargs},
-    )
-    times = timer.repeat(number=1, repeat=num_iterations + 1)
-    return Benchmark(np.mean(times[1:]).item(), np.std(times[1:]).item())
+from fft_conv_pytorch.utils import Benchmark, benchmark
 
 
 @lru_cache(maxsize=1)

fft_conv_pytorch/utils.py

@@ -0,0 +1,39 @@
+from timeit import Timer
+from typing import Callable, NamedTuple
+
+import numpy as np
+import torch
+from torch import Tensor
+
+
+class Benchmark(NamedTuple):
+    mean: float
+    std: float
+
+    def __repr__(self):
+        return f"BenchmarkResult(mean: {self.mean:.3e}, std: {self.std:.3e})"
+
+    def __str__(self):
+        return f"({self.mean:.3e} \u00B1 {self.std:.3e}) s"
+
+
+def benchmark(fn: Callable, *args, num_iterations: int = 10, **kwargs) -> Benchmark:
+    timer = Timer(
+        "fn(*args, **kwargs)",
+        globals={"fn": fn, "args": args, "kwargs": kwargs},
+    )
+    times = timer.repeat(number=1, repeat=num_iterations + 1)
+    return Benchmark(np.mean(times[1:]).item(), np.std(times[1:]).item())
+
+
+def _assert_almost_equal(x: Tensor, y: Tensor) -> bool:
+    abs_error = torch.abs(x - y)
+    assert abs_error.mean().item() < 5e-5
+    assert abs_error.max().item() < 1e-4
+    return True
+
+
+def _gcd(x: int, y: int) -> int:
+    while y:
+        x, y = y, x % y
+    return x

setup.py

@@ -24,7 +24,7 @@ def get_version_tag() -> str:
     description="Implementation of 1D, 2D, and 3D FFT convolutions in PyTorch.",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",
-    install_requires=["numpy", "torch>=1.7"],
+    install_requires=["numpy", "torch>=1.8"],
     extras_require={"test": ["black", "flake8", "isort", "pytest", "pytest-cov"]},
     classifiers=[
         "Programming Language :: Python :: 3",

tests/__init__.py

@@ -5,7 +5,7 @@
 import torch.nn.functional as f
 
 from fft_conv_pytorch.fft_conv import fft_conv, to_ntuple
-from tests.utils import _assert_almost_equal, _gcd
+from fft_conv_pytorch.utils import _assert_almost_equal, _gcd
 
 
 @pytest.mark.parametrize("in_channels", [2, 3])
@@ -53,7 +53,12 @@ def test_fft_conv_functional(
     b0 = torch.randn(out_channels, requires_grad=True) if bias else None
     b1 = b0.detach().clone().requires_grad_() if bias else None
 
-    kwargs = dict(padding=padding, stride=stride, dilation=dilation, groups=groups,)
+    kwargs = dict(
+        padding=padding,
+        stride=stride,
+        dilation=dilation,
+        groups=groups,
+    )
 
     y0 = fft_conv(signal, w0, bias=b0, **kwargs)
     y1 = torch_conv(signal, w1, bias=b1, **kwargs)
@@ -99,7 +104,12 @@ def test_fft_conv_backward_functional(
     b0 = torch.randn(out_channels, requires_grad=True) if bias else None
     b1 = b0.detach().clone().requires_grad_() if bias else None
 
-    kwargs = dict(padding=padding, stride=stride, dilation=dilation, groups=groups,)
+    kwargs = dict(
+        padding=padding,
+        stride=stride,
+        dilation=dilation,
+        groups=groups,
+    )
 
     y0 = fft_conv(signal, w0, bias=b0, **kwargs)
     y1 = torch_conv(signal, w1, bias=b1, **kwargs)

tests/test_functional.py

@@ -5,7 +5,7 @@
 import torch.nn.functional as f
 
 from fft_conv_pytorch.fft_conv import _FFTConv
-from tests.utils import _assert_almost_equal, _gcd
+from fft_conv_pytorch.utils import _assert_almost_equal, _gcd
 
 
 @pytest.mark.parametrize("in_channels", [2, 3])
@@ -50,7 +50,12 @@ def test_fft_conv_module(
     weight = fft_conv_layer.weight
     bias = fft_conv_layer.bias
 
-    kwargs = dict(padding=padding, stride=stride, dilation=dilation, groups=groups,)
+    kwargs = dict(
+        padding=padding,
+        stride=stride,
+        dilation=dilation,
+        groups=groups,
+    )
 
     y0 = fft_conv_layer(signal)
     y1 = torch_conv(signal, weight, bias=bias, **kwargs)
@@ -102,7 +107,12 @@ def test_fft_conv_backward_module(
     b0 = fft_conv_layer.bias
     b1 = b0.detach().clone().requires_grad_() if bias else None
 
-    kwargs = dict(padding=padding, stride=stride, dilation=dilation, groups=groups,)
+    kwargs = dict(
+        padding=padding,
+        stride=stride,
+        dilation=dilation,
+        groups=groups,
+    )
 
     y0 = fft_conv_layer(signal)
     y1 = torch_conv(signal, w1, bias=b1, **kwargs)