Remove remote model downloads from CI tests

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/conftest.py

@@ -6,6 +6,25 @@
 
 import pytest
 import torch
+import torch.nn as nn
+
+
+class MockBackbone(nn.Module):
+    def __init__(self):
+        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)
 
 
 @pytest.fixture(
@@ -24,3 +43,8 @@ def arange_4d_image_odd(request):
     x = torch.arange(torch.prod(torch.tensor(request.param))).reshape(request.param)
 
     return x.to(torch.get_default_dtype())
+
+
+@pytest.fixture(scope="session")
+def mock_backbone():
+    return MockBackbone()

tests/metrics/test_dists.py

@@ -7,14 +7,16 @@
 import torch
 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, mock_backbone):
     if arange_4d_image.shape[1] != 3:
         return
 
+    monkeypatch.setattr(neuralcompression.metrics._dists, "NoTrainDists", mock_backbone)
     metric = DeepImageStructureTextureSimilarity()
 
     for _ in range(num_samples):

tests/metrics/test_fid.py

@@ -8,17 +8,21 @@
 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, mock_backbone):
     if arange_4d_image.shape[1] != 3:
         return
 
     rng = torch.Generator()
     rng.manual_seed(55)
 
+    monkeypatch.setattr(
+        neuralcompression.metrics._fid, "NoTrainInceptionV3", mock_backbone
+    )
     metric = FrechetInceptionDistance()
 
     for _ in range(num_samples):

tests/metrics/test_kid.py

@@ -8,14 +8,18 @@
 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, mock_backbone):
     if arange_4d_image.shape[1] != 3:
         return
 
+    monkeypatch.setattr(
+        neuralcompression.metrics._kid, "NoTrainInceptionV3", mock_backbone
+    )
     metric = KernelInceptionDistance()
 
     for _ in range(num_samples):

tests/metrics/test_swav_fid.py

@@ -8,17 +8,21 @@
 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, mock_backbone):
     if arange_4d_image.shape[1] != 3:
         return
 
     rng = torch.Generator()
     rng.manual_seed(60)
 
+    monkeypatch.setattr(
+        neuralcompression.metrics._fid_swav, "NoTrainSwAV", mock_backbone
+    )
     metric = FrechetInceptionDistanceSwAV()
 
     for _ in range(num_samples):

tests/metrics/test_update_patch_fid.py

@@ -7,6 +7,9 @@
 import torch
 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 +19,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, mock_backbone):
     if arange_4d_image.shape[1] != 3:
         return
 
+    monkeypatch.setattr(
+        neuralcompression.metrics._fid, "NoTrainInceptionV3", mock_backbone
+    )
+    monkeypatch.setattr(
+        neuralcompression.metrics._fid_swav, "NoTrainSwAV", mock_backbone
+    )
+    monkeypatch.setattr(
+        neuralcompression.metrics._kid, "NoTrainInceptionV3", mock_backbone
+    )
+
     fid_metric = FrechetInceptionDistance()
     fid_swav_metric = FrechetInceptionDistanceSwAV()
     kid_metric = KernelInceptionDistance()