Remove remote model downloads from CI tests (#228)

* Remove remote model downloads from CI tests

* Fix class mocking

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)
 

pyproject.toml

@@ -8,6 +8,3 @@ requires = [
 
 [tool.setuptools_scm]
 write_to = "neuralcompression/version.py"
-
-[isort]
-profile = "black"

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())
 

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(

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):

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):

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):

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)
 

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):

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()

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())