diff --git a/neuralcompression/metrics/_dists.py b/neuralcompression/metrics/_dists.py index 7c9b6ed9..330bd6c0 100644 --- a/neuralcompression/metrics/_dists.py +++ b/neuralcompression/metrics/_dists.py @@ -13,7 +13,7 @@ class NoTrainDists(_DISTS): - def train(self, mode: bool) -> "NoTrainDists": + def train(self, mode: bool = True) -> "NoTrainDists": """keep network in evaluation mode.""" return super().train(False) diff --git a/pyproject.toml b/pyproject.toml index 57baedec..e37c0e34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,6 +8,3 @@ requires = [ [tool.setuptools_scm] write_to = "neuralcompression/version.py" - -[isort] -profile = "black" diff --git a/tests/cached_data/create_cached_tensorflow_data.py b/tests/cached_data/create_cached_tensorflow_data.py index 7d55694f..9f81441d 100644 --- a/tests/cached_data/create_cached_tensorflow_data.py +++ b/tests/cached_data/create_cached_tensorflow_data.py @@ -12,7 +12,7 @@ def create_input(shape, offset: int = 0): - x = np.arange(np.product(shape)).reshape(shape) + offset + x = np.arange(np.prod(shape)).reshape(shape) + offset return torch.from_numpy(x).to(torch.get_default_dtype()) diff --git a/tests/conftest.py b/tests/conftest.py index 76ad8697..f3f8c7ba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,38 @@ import pytest import torch +import torch.nn as nn +from torch import Tensor + + +class MockBackbone(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + + self.model = nn.Sequential( + nn.Conv2d(3, 2048, kernel_size=3, padding=1), + nn.AdaptiveAvgPool2d(output_size=(1, 1)), + ) + for param in self.parameters(): + param.requires_grad_(False) + + self.eval() + + def train(self, mode: bool = True) -> "MockBackbone": + """keep network in evaluation mode.""" + return super().train(False) + + def forward(self, image: Tensor) -> Tensor: + return self.model(image.float() / 255.0).flatten(1) + + +class MockDiffBackbone(nn.Module): + def __init__(self, *args, **kwargs): + super().__init__() + self.model = MockBackbone() + + def forward(self, image1: Tensor, image2: Tensor) -> Tensor: + return self.model(image1) - self.model(image2) @pytest.fixture( diff --git a/tests/metrics/test_dists.py b/tests/metrics/test_dists.py index af761802..dca977ba 100644 --- a/tests/metrics/test_dists.py +++ b/tests/metrics/test_dists.py @@ -5,16 +5,21 @@ import pytest import torch +from conftest import MockDiffBackbone from torch import Tensor +import neuralcompression.metrics._dists from neuralcompression.metrics import DeepImageStructureTextureSimilarity @pytest.mark.parametrize("num_samples", [5]) -def test_dists(num_samples: int, arange_4d_image: Tensor): +def test_dists(num_samples: int, arange_4d_image: Tensor, monkeypatch): if arange_4d_image.shape[1] != 3: return + monkeypatch.setattr( + neuralcompression.metrics._dists, "NoTrainDists", MockDiffBackbone + ) metric = DeepImageStructureTextureSimilarity() for _ in range(num_samples): diff --git a/tests/metrics/test_fid.py b/tests/metrics/test_fid.py index 7bbf39e1..cddc0f3b 100644 --- a/tests/metrics/test_fid.py +++ b/tests/metrics/test_fid.py @@ -5,20 +5,25 @@ import pytest import torch +from conftest import MockBackbone from torch import Tensor import neuralcompression.functional as ncF +import neuralcompression.metrics._fid from neuralcompression.metrics import FrechetInceptionDistance @pytest.mark.parametrize("num_samples", [5]) -def test_dists(num_samples: int, arange_4d_image: Tensor): +def test_fid(num_samples: int, arange_4d_image: Tensor, monkeypatch): if arange_4d_image.shape[1] != 3: return rng = torch.Generator() rng.manual_seed(55) + monkeypatch.setattr( + neuralcompression.metrics._fid, "NoTrainInceptionV3", MockBackbone + ) metric = FrechetInceptionDistance() for _ in range(num_samples): diff --git a/tests/metrics/test_kid.py b/tests/metrics/test_kid.py index 4384423b..f69266bc 100644 --- a/tests/metrics/test_kid.py +++ b/tests/metrics/test_kid.py @@ -5,17 +5,22 @@ import pytest import torch +from conftest import MockBackbone from torch import Tensor import neuralcompression.functional as ncF +import neuralcompression.metrics._kid from neuralcompression.metrics import KernelInceptionDistance @pytest.mark.parametrize("num_samples", [5]) -def test_dists(num_samples: int, arange_4d_image: Tensor): +def test_kid(num_samples: int, arange_4d_image: Tensor, monkeypatch): if arange_4d_image.shape[1] != 3: return + monkeypatch.setattr( + neuralcompression.metrics._kid, "NoTrainInceptionV3", MockBackbone + ) metric = KernelInceptionDistance() for _ in range(num_samples): diff --git a/tests/metrics/test_pickle_size.py b/tests/metrics/test_pickle_size.py index caddcd41..a6a22a02 100644 --- a/tests/metrics/test_pickle_size.py +++ b/tests/metrics/test_pickle_size.py @@ -15,13 +15,13 @@ @pytest.mark.parametrize("arr_size", [64, 128, 256, (64, 64)]) def test_pickle_size(arr_size, tmp_path: Path): - x = np.reshape(np.arange(np.product(arr_size)), arr_size) + x = np.reshape(np.arange(np.prod(arr_size)), arr_size) obj = {"thearr": x} mem_size = pickle_size_of(obj) - tmp_file = f"pickle_size_of_{np.product(arr_size)}.pkl" + tmp_file = f"pickle_size_of_{np.prod(arr_size)}.pkl" with open(tmp_path / tmp_file, "wb") as f: pickle.dump(obj, f) diff --git a/tests/metrics/test_swav_fid.py b/tests/metrics/test_swav_fid.py index 242b70b4..71c0d05e 100644 --- a/tests/metrics/test_swav_fid.py +++ b/tests/metrics/test_swav_fid.py @@ -5,20 +5,25 @@ import pytest import torch +from conftest import MockBackbone from torch import Tensor import neuralcompression.functional as ncF +import neuralcompression.metrics._fid_swav from neuralcompression.metrics import FrechetInceptionDistanceSwAV @pytest.mark.parametrize("num_samples", [5]) -def test_dists(num_samples: int, arange_4d_image: Tensor): +def test_dists(num_samples: int, arange_4d_image: Tensor, monkeypatch): if arange_4d_image.shape[1] != 3: return rng = torch.Generator() rng.manual_seed(60) + monkeypatch.setattr( + neuralcompression.metrics._fid_swav, "NoTrainSwAV", MockBackbone + ) metric = FrechetInceptionDistanceSwAV() for _ in range(num_samples): diff --git a/tests/metrics/test_update_patch_fid.py b/tests/metrics/test_update_patch_fid.py index d2817437..5702b11d 100644 --- a/tests/metrics/test_update_patch_fid.py +++ b/tests/metrics/test_update_patch_fid.py @@ -5,8 +5,12 @@ import pytest import torch +from conftest import MockBackbone from torch import Tensor +import neuralcompression.metrics._fid +import neuralcompression.metrics._fid_swav +import neuralcompression.metrics._kid from neuralcompression.metrics import ( FrechetInceptionDistance, FrechetInceptionDistanceSwAV, @@ -16,10 +20,20 @@ @pytest.mark.parametrize("num_samples", [5]) -def test_dists(num_samples: int, arange_4d_image: Tensor): +def test_dists(num_samples: int, arange_4d_image: Tensor, monkeypatch): if arange_4d_image.shape[1] != 3: return + monkeypatch.setattr( + neuralcompression.metrics._fid, "NoTrainInceptionV3", MockBackbone + ) + monkeypatch.setattr( + neuralcompression.metrics._fid_swav, "NoTrainSwAV", MockBackbone + ) + monkeypatch.setattr( + neuralcompression.metrics._kid, "NoTrainInceptionV3", MockBackbone + ) + fid_metric = FrechetInceptionDistance() fid_swav_metric = FrechetInceptionDistanceSwAV() kid_metric = KernelInceptionDistance() diff --git a/tests/utils.py b/tests/utils.py index d8792e0a..d2a953e9 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -13,7 +13,7 @@ def create_input(shape, offset: int = 0): - x = np.arange(np.product(shape)).reshape(shape) + offset + x = np.arange(np.prod(shape)).reshape(shape) + offset return torch.from_numpy(x).to(torch.get_default_dtype())