Add base params test

Author: qubvel

tests/models/base.py

@@ -1,4 +1,5 @@
 import unittest
+import inspect
 from functools import lru_cache
 
 import torch
@@ -23,6 +24,18 @@ def model_type(self):
             raise ValueError("test_model_type is not set")
         return self.test_model_type
 
+    @property
+    def model_class(self):
+        return smp.MODEL_ARCHITECTURES_MAPPING[self.model_type]
+
+    @property
+    def decoder_channels(self):
+        signature = inspect.signature(self.model_class)
+        # check if decoder_channels is in the signature
+        if "decoder_channels" in signature.parameters:
+            return signature.parameters["decoder_channels"].default
+        return None
+
     @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)
@@ -50,3 +63,29 @@ def test_forward_backward(self):
 
         # check backward pass
         output.mean().backward()
+
+    def test_base_params_are_set(self, in_channels=1, depth=3, classes=7):
+        kwargs = {}
+
+        if self.model_type in ["unet", "unetplusplus", "manet"]:
+            kwargs = {"decoder_channels": self.decoder_channels[:depth]}
+
+        model = smp.create_model(
+            arch=self.model_type,
+            encoder_depth=depth,
+            in_channels=in_channels,
+            classes=classes,
+            **kwargs,
+        )
+        sample = self._get_sample(
+            batch_size=self.default_batch_size,
+            num_channels=in_channels,
+            height=self.default_height,
+            width=self.default_width,
+        )
+
+        # check in channels correctly set
+        with torch.no_grad():
+            output = model(sample)
+
+        self.assertEqual(output.shape[1], classes)