Add depth validation

Author: qubvel

segmentation_models_pytorch/encoders/densenet.py

@@ -32,7 +32,13 @@
 
 class DenseNetEncoder(DenseNet, EncoderMixin):
     def __init__(self, out_channels, depth=5, output_stride=32, **kwargs):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         super().__init__(**kwargs)
+
         self._depth = depth
         self._in_channels = 3
         self._out_channels = out_channels

segmentation_models_pytorch/encoders/dpn.py

@@ -44,6 +44,11 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         super().__init__(**kwargs)
         self._stage_idxs = stage_idxs
         self._depth = depth

segmentation_models_pytorch/encoders/efficientnet.py

@@ -43,6 +43,11 @@ def __init__(
         depth: int = 5,
         output_stride: int = 32,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         blocks_args, global_params = get_model_params(model_name, override_params=None)
         super().__init__(blocks_args, global_params)
 

segmentation_models_pytorch/encoders/inceptionresnetv2.py

@@ -39,6 +39,11 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/inceptionv4.py

@@ -40,6 +40,10 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/mix_transformer.py

@@ -529,6 +529,10 @@ class MixVisionTransformerEncoder(MixVisionTransformer, EncoderMixin):
     def __init__(
         self, out_channels: List[int], depth: int = 5, output_stride: int = 32, **kwargs
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/mobilenet.py

@@ -34,6 +34,10 @@ class MobileNetV2Encoder(torchvision.models.MobileNetV2, EncoderMixin):
     def __init__(
         self, out_channels: List[int], depth: int = 5, output_stride: int = 32, **kwargs
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/mobileone.py

@@ -319,6 +319,11 @@ def __init__(
         :param use_se: Whether to use SE-ReLU activations.
         :param num_conv_branches: Number of linear conv branches.
         """
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         super().__init__()
 
         assert len(width_multipliers) == 4

segmentation_models_pytorch/encoders/resnet.py

@@ -38,6 +38,10 @@ class ResNetEncoder(ResNet, EncoderMixin):
     def __init__(
         self, out_channels: List[int], depth: int = 5, output_stride: int = 32, **kwargs
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/senet.py

@@ -43,6 +43,10 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/timm_efficientnet.py

@@ -105,6 +105,10 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._stage_idxs = stage_idxs

segmentation_models_pytorch/encoders/timm_sknet.py

@@ -14,6 +14,10 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(**kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/timm_universal.py

@@ -68,6 +68,13 @@ def __init__(
             output_stride (int): Desired output stride (default: 32).
             **kwargs: Additional arguments passed to `timm.create_model`.
         """
+        # At the moment we do not support models with more than 5 stages,
+        # but can be reconfigured in the future.
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
+
         super().__init__()
         self.name = name
 

segmentation_models_pytorch/encoders/vgg.py

@@ -53,6 +53,10 @@ def __init__(
         output_stride: int = 32,
         **kwargs,
     ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(make_layers(config, batch_norm=batch_norm), **kwargs)
 
         self._depth = depth

segmentation_models_pytorch/encoders/xception.py

@@ -1,10 +1,22 @@
+from typing import List
 from pretrainedmodels.models.xception import Xception
 
 from ._base import EncoderMixin
 
 
 class XceptionEncoder(Xception, EncoderMixin):
-    def __init__(self, out_channels, *args, depth=5, output_stride=32, **kwargs):
+    def __init__(
+        self,
+        out_channels: List[int],
+        *args,
+        depth: int = 5,
+        output_stride: int = 32,
+        **kwargs,
+    ):
+        if depth > 5 or depth < 1:
+            raise ValueError(
+                f"{self.__class__.__name__} depth should be in range [1, 5], got {depth}"
+            )
         super().__init__(*args, **kwargs)
 
         self._depth = depth

tests/encoders/base.py

@@ -149,6 +149,12 @@ def test_depth(self):
                     f"Encoder `{encoder_name}` should have {depth + 1} out_channels, but has {len(encoder.out_channels)}",
                 )
 
+    def test_invalid_depth(self):
+        with self.assertRaises(ValueError):
+            smp.encoders.get_encoder(self.encoder_names[0], depth=6)
+        with self.assertRaises(ValueError):
+            smp.encoders.get_encoder(self.encoder_names[0], depth=0)
+
     def test_dilated(self):
         sample = self._get_sample().to(default_device)