Add test with hub checkpoint

Author: qubvel

misc/generate_test_models.py

@@ -0,0 +1,41 @@
+import os
+import torch
+import tempfile
+import huggingface_hub
+import segmentation_models_pytorch as smp
+
+HUB_REPO = "smp-test-models"
+ENCODER_NAME = "tu-resnet18"
+
+api = huggingface_hub.HfApi(token=os.getenv("HF_TOKEN"))
+
+for model_name, model_class in smp.MODEL_ARCHITECTURES_MAPPING.items():
+    model = model_class(encoder_name=ENCODER_NAME)
+    model = model.eval()
+
+    # generate test sample
+    torch.manual_seed(423553)
+    sample = torch.rand(1, 3, 256, 256)
+
+    with torch.no_grad():
+        output = model(sample)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # save model
+        model.save_pretrained(f"{tmpdir}")
+
+        # save input and output
+        torch.save(sample, f"{tmpdir}/input-tensor.pth")
+        torch.save(output, f"{tmpdir}/output-tensor.pth")
+
+        # create repo
+        repo_id = f"{HUB_REPO}/{model_name}-{ENCODER_NAME}"
+        if not api.repo_exists(repo_id=repo_id):
+            api.create_repo(repo_id=repo_id, repo_type="model")
+
+        # upload to hub
+        api.upload_folder(
+            folder_path=tmpdir,
+            repo_id=f"{HUB_REPO}/{model_name}-{ENCODER_NAME}",
+            repo_type="model",
+        )

tests/encoders/test_pretrainedmodels_encoders.py

@@ -1,5 +1,5 @@
 from tests.encoders import base
-from tests.config import RUN_ALL_ENCODERS
+from tests.utils import RUN_ALL_ENCODERS
 
 
 class TestDenseNetEncoder(base.BaseEncoderTester):

tests/encoders/test_smp_encoders.py

@@ -1,5 +1,5 @@
 from tests.encoders import base
-from tests.config import RUN_ALL_ENCODERS
+from tests.utils import RUN_ALL_ENCODERS
 
 
 class TestMobileoneEncoder(base.BaseEncoderTester):

tests/encoders/test_timm_ported_encoders.py

@@ -1,5 +1,5 @@
 from tests.encoders import base
-from tests.config import RUN_ALL_ENCODERS
+from tests.utils import RUN_ALL_ENCODERS
 
 
 class TestTimmEfficientNetEncoder(base.BaseEncoderTester):

tests/encoders/test_timm_universal.py

@@ -1,5 +1,5 @@
 from tests.encoders import base
-from tests.config import has_timm_test_models
+from tests.utils import has_timm_test_models
 
 # check if timm >= 1.0.12
 timm_encoders = [

tests/encoders/test_torchvision_encoders.py

@@ -1,5 +1,5 @@
 from tests.encoders import base
-from tests.config import RUN_ALL_ENCODERS
+from tests.utils import RUN_ALL_ENCODERS
 
 
 class TestMobileoneEncoder(base.BaseEncoderTester):

tests/models/base.py

@@ -7,7 +7,7 @@
 import torch
 import segmentation_models_pytorch as smp
 
-from tests.config import has_timm_test_models
+from tests.utils import has_timm_test_models, slow_test
 
 
 class BaseModelTester(unittest.TestCase):
@@ -30,6 +30,10 @@ def model_type(self):
             raise ValueError("test_model_type is not set")
         return self.test_model_type
 
+    @property
+    def hub_checkpoint(self):
+        return f"smp-test-models/{self.model_type}-tu-resnet18"
+
     @property
     def model_class(self):
         return smp.MODEL_ARCHITECTURES_MAPPING[self.model_type]
@@ -166,3 +170,27 @@ def test_save_load_with_hub_mixin(self):
         # check dataset and metrics are saved in readme
         self.assertIn("test_dataset", readme)
         self.assertIn("my_awesome_metric", readme)
+
+    @slow_test
+    def test_preserve_forward_output(self):
+        from huggingface_hub import hf_hub_download
+
+        model = smp.from_pretrained(self.hub_checkpoint).eval()
+
+        input_tensor_path = hf_hub_download(
+            repo_id=self.hub_checkpoint, filename="input-tensor.pth"
+        )
+        output_tensor_path = hf_hub_download(
+            repo_id=self.hub_checkpoint, filename="output-tensor.pth"
+        )
+
+        input_tensor = torch.load(input_tensor_path, weights_only=True)
+        output_tensor = torch.load(output_tensor_path, weights_only=True)
+
+        with torch.no_grad():
+            output = model(input_tensor)
+
+        self.assertEqual(output.shape, output_tensor.shape)
+        is_close = torch.allclose(output, output_tensor, atol=1e-3)
+        max_diff = torch.max(torch.abs(output - output_tensor))
+        self.assertTrue(is_close, f"Max diff: {max_diff}")

tests/models/test_segformer.py

@@ -1,7 +1,41 @@
+import torch
 import pytest
+import segmentation_models_pytorch as smp
+
 from tests.models import base
+from tests.utils import slow_test, default_device
 
 
 @pytest.mark.segformer
 class TestSegformerModel(base.BaseModelTester):
     test_model_type = "segformer"
+
+    @slow_test
+    def test_load_pretrained(self):
+        hub_checkpoint = "smp-hub/segformer-b0-512x512-ade-160k"
+
+        model = smp.from_pretrained(hub_checkpoint)
+        model = model.eval().to(default_device)
+
+        sample = torch.ones([1, 3, 512, 512]).to(default_device)
+
+        with torch.no_grad():
+            output = model(sample)
+
+        self.assertEqual(output.shape, (1, 150, 512, 512))
+
+        expected_logits_slice = torch.tensor(
+            [-4.4172, -4.4723, -4.5273, -4.5824, -4.6375, -4.7157]
+        )
+        resulted_logits_slice = output[0, 0, 256, :6].cpu()
+        is_equal = torch.allclose(
+            expected_logits_slice, resulted_logits_slice, atol=1e-2
+        )
+        max_diff = torch.max(torch.abs(expected_logits_slice - resulted_logits_slice))
+        self.assertTrue(
+            is_equal,
+            f"Expected logits slice and resulted logits slice are not equal.\n"
+            f"Max diff: {max_diff}\n"
+            f"Expected: {expected_logits_slice}\n"
+            f"Resulted: {resulted_logits_slice}\n",
+        )

tests/config.py renamed to tests/utils.py

@@ -1,9 +1,13 @@
 import os
 import timm
+import torch
+import unittest
+
 from packaging.version import Version
 
 
 has_timm_test_models = Version(timm.__version__) >= Version("1.0.12")
+default_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def get_commit_message():
@@ -28,3 +32,22 @@ def get_commit_message():
     os.getenv("RUN_SLOW", "false").lower() in ["true", "1", "y", "yes"]
     or "run-slow" in commit_message
 )
+
+
+def slow_test(test_case):
+    """
+    Decorator marking a test as slow.
+
+    Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.
+
+    """
+    return unittest.skipUnless(RUN_SLOW, "test is slow")(test_case)
+
+
+def requires_torch_greater_or_equal(version: str):
+    torch_version = Version(torch.__version__)
+    provided_version = Version(version)
+    return unittest.skipUnless(
+        torch_version >= provided_version,
+        f"torch version {torch_version} is less than {provided_version}",
+    )