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