Basic model test

Author: qubvel

pyproject.toml

@@ -55,3 +55,18 @@ version = {attr = 'segmentation_models_pytorch.__version__.__version__'}
 
 [tool.setuptools.packages.find]
 include = ['segmentation_models_pytorch*']
+
+[tool.pytest.ini_options]
+markers = [
+    "deeplabv3",
+    "deeplabv3plus",
+    "fpn",
+    "linknet",
+    "manet",
+    "pan",
+    "psp",
+    "segformer",
+    "unet",
+    "unetplusplus",
+    "upernet",
+]

tests/encoders/test_timm_universal.py

@@ -3,6 +3,7 @@
 
 class TestTimmUniversalEncoder(base.BaseEncoderTester):
     encoder_names = [
+        "tu-test_resnet.r160_in1k",
         "tu-resnet18",  # for timm universal traditional-like encoder
         "tu-convnext_atto",  # for timm universal transformer-like encoder
         "tu-darknet17",  # for timm universal vgg-like encoder

tests/models/__init__.py

@@ -0,0 +1,56 @@
+import unittest
+from functools import lru_cache
+
+import torch
+import segmentation_models_pytorch as smp
+
+
+class BaseModelTester(unittest.TestCase):
+    test_encoder_name = "tu-test_resnet.r160_in1k"
+
+    # should be overriden
+    test_model_type = None
+
+    # test sample configuration
+    default_batch_size = 1
+    default_num_channels = 3
+    default_height = 64
+    default_width = 64
+
+    @property
+    def model_type(self):
+        if self.test_model_type is None:
+            raise ValueError("test_model_type is not set")
+        return self.test_model_type
+
+    @lru_cache
+    def _get_sample(self, batch_size=1, num_channels=3, height=32, width=32):
+        return torch.rand(batch_size, num_channels, height, width)
+
+    def test_forward_backward(self):
+        sample = self._get_sample(
+            batch_size=self.default_batch_size,
+            num_channels=self.default_num_channels,
+            height=self.default_height,
+            width=self.default_width,
+        )
+        model = smp.create_model(arch=self.model_type)
+
+        # check default in_channels=3
+        output = model(sample)
+
+        # check default output number of classes = 1
+        expected_number_of_classes = 1
+        result_number_of_classes = output.shape[1]
+        self.assertEqual(
+            result_number_of_classes,
+            expected_number_of_classes,
+            f"Default output number of classes should be {expected_number_of_classes}, but got {result_number_of_classes}",
+        )
+
+        # check backward pass
+        output.mean().backward()
+
+    def test_encoder_params_are_set(self):
+        model = smp.create_model(arch=self.model_type)
+        self.assertEqual(model.encoder.name, self.test_encoder_name)

tests/models/base.py

@@ -0,0 +1,16 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.deeplabv3
+class TestDeeplabV3Model(base.BaseModelTester):
+    test_model_type = "deeplabv3"
+
+    default_batch_size = 2
+
+
+@pytest.mark.deeplabv3plus
+class TestDeeplabV3PlusModel(base.BaseModelTester):
+    test_model_type = "deeplabv3plus"
+
+    default_batch_size = 2

tests/models/test_deeplab.py

@@ -0,0 +1,7 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.fpn
+class TestFpnModel(base.BaseModelTester):
+    test_model_type = "fpn"

tests/models/test_fpn.py

@@ -0,0 +1,7 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.linknet
+class TestLinknetModel(base.BaseModelTester):
+    test_model_type = "linknet"

tests/models/test_linknet.py

@@ -0,0 +1,7 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.manet
+class TestManetModel(base.BaseModelTester):
+    test_model_type = "manet"

tests/models/test_manet.py

@@ -0,0 +1,12 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.pan
+class TestPanModel(base.BaseModelTester):
+    test_model_type = "pan"
+    test_encoder_name = "resnet-18"
+
+    default_batch_size = 2
+    default_height = 128
+    default_width = 128

tests/models/test_pan.py

@@ -0,0 +1,9 @@
+import pytest
+from tests.models import base
+
+
+@pytest.mark.psp
+class TestPspModel(base.BaseModelTester):
+    test_model_type = "pspnet"
+
+    default_batch_size = 2