First fix following review

Author: GuillaumeErhard

segmentation_models_pytorch/base/modules.py

@@ -1,3 +1,4 @@
+from typing import Any, Dict, Tuple, Union
 import warnings
 
 import torch
@@ -8,6 +9,76 @@
 except ImportError:
     InPlaceABN = None
 
+def handle_decoder_use_batchnorm_deprecation(decoder_use_batchnorm: Union[bool, str, None], decoder_use_norm: Union[bool, str, Dict[str, Any]]) -> Dict[str, Any]:
+    if decoder_use_batchnorm is not None:
+        warnings.warn(
+            "The usage of use_batchnorm is deprecated. Please modify your code for use_norm",
+            DeprecationWarning,
+        )
+        if decoder_use_batchnorm is True:
+            decoder_use_norm = {"type": "batchnorm"}
+        elif decoder_use_batchnorm is False:
+            decoder_use_norm = {"type": "identity"}
+        elif decoder_use_batchnorm == "inplace":
+            decoder_use_norm = {
+                "type": "inplace",
+                "activation": "leaky_relu",
+                "activation_param": 0.0,
+            }
+        else:
+            raise ValueError("Unrecognized value for use_batchnorm")
+
+    return decoder_use_norm
+
+
+def normalize_use_norm(use_norm: Union[bool, str, Dict[str, Any]]) -> Dict[str, Any]:
+    if isinstance(use_norm, str):
+        norm_str = use_norm.lower()
+        if norm_str == "inplace":
+            use_norm = {
+                "type": "inplace",
+                "activation": "leaky_relu",
+                "activation_param": 0.0,
+            }
+        elif norm_str in (
+                "batchnorm",
+                "identity",
+                "layernorm",
+                "groupnorm",
+                "instancenorm",
+        ):
+            use_norm = {"type": norm_str}
+        else:
+            raise ValueError("Unrecognized normalization type string provided")
+    elif isinstance(use_norm, bool):
+        use_norm = {"type": "batchnorm" if use_norm else "identity"}
+    elif not isinstance(use_norm, dict):
+        raise ValueError("use_norm must be a dictionary, boolean, or string")
+
+    return use_norm
+
+def get_norm_layer(use_norm: Dict[str, Any], relu: nn.Module, out_channels: int) -> Tuple[nn.Module, nn.Module]:
+    norm_type = use_norm["type"]
+    extra_kwargs = {k: v for k, v in use_norm.items() if k != "type"}
+
+    if norm_type == "inplace":
+        norm = InPlaceABN(out_channels, **extra_kwargs)
+        relu = nn.Identity()
+    elif norm_type == "batchnorm":
+        norm = nn.BatchNorm2d(out_channels, **extra_kwargs)
+    elif norm_type == "identity":
+        norm = nn.Identity()
+    elif norm_type == "layernorm":
+        norm = nn.LayerNorm(out_channels, **extra_kwargs)
+    elif norm_type == "groupnorm":
+        norm = nn.GroupNorm(out_channels, **extra_kwargs)
+    elif norm_type == "instancenorm":
+        norm = nn.InstanceNorm2d(out_channels, **extra_kwargs)
+    else:
+        raise ValueError(f"Unrecognized normalization type: {norm_type}")
+
+    return norm, relu
+
 
 class Conv2dReLU(nn.Sequential):
     def __init__(
@@ -17,50 +88,9 @@ def __init__(
         kernel_size,
         padding=0,
         stride=1,
-        use_batchnorm=True,
         use_norm="batchnorm",
     ):
-        if use_batchnorm is not None:
-            warnings.warn(
-                "The usage of use_batchnorm is deprecated. Please modify your code for use_norm",
-                DeprecationWarning,
-            )
-            if use_batchnorm is True:
-                use_norm = {"type": "batchnorm"}
-            elif use_batchnorm is False:
-                use_norm = {"type": "identity"}
-            elif use_batchnorm == "inplace":
-                use_norm = {
-                    "type": "inplace",
-                    "activation": "leaky_relu",
-                    "activation_param": 0.0,
-                }
-            else:
-                raise ValueError("Unrecognized value for use_batchnorm")
-
-        if isinstance(use_norm, str):
-            norm_str = use_norm.lower()
-            if norm_str == "inplace":
-                use_norm = {
-                    "type": "inplace",
-                    "activation": "leaky_relu",
-                    "activation_param": 0.0,
-                }
-            elif norm_str in (
-                "batchnorm",
-                "identity",
-                "layernorm",
-                "groupnorm",
-                "instancenorm",
-            ):
-                use_norm = {"type": norm_str}
-            else:
-                raise ValueError("Unrecognized normalization type string provided")
-        elif isinstance(use_norm, bool):
-            use_norm = {"type": "batchnorm" if use_norm else "identity"}
-        elif not isinstance(use_norm, dict):
-            raise ValueError("use_norm must be a dictionary, boolean, or string")
-
+        use_norm = normalize_use_norm(use_norm)
         if use_norm["type"] == "inplace" and InPlaceABN is None:
             raise RuntimeError(
                 "In order to use `use_batchnorm='inplace'` or `use_norm='inplace'` the inplace_abn package must be installed. "
@@ -77,24 +107,7 @@ def __init__(
         )
         relu = nn.ReLU(inplace=True)
 
-        norm_type = use_norm["type"]
-        extra_kwargs = {k: v for k, v in use_norm.items() if k != "type"}
-
-        if norm_type == "inplace":
-            norm = InPlaceABN(out_channels, **extra_kwargs)
-            relu = nn.Identity()
-        elif norm_type == "batchnorm":
-            norm = nn.BatchNorm2d(out_channels, **extra_kwargs)
-        elif norm_type == "identity":
-            norm = nn.Identity()
-        elif norm_type == "layernorm":
-            norm = nn.LayerNorm(out_channels, **extra_kwargs)
-        elif norm_type == "groupnorm":
-            norm = nn.GroupNorm(out_channels, **extra_kwargs)
-        elif norm_type == "instancenorm":
-            norm = nn.InstanceNorm2d(out_channels, **extra_kwargs)
-        else:
-            raise ValueError(f"Unrecognized normalization type: {norm_type}")
+        norm, relu = get_norm_layer(use_norm, relu, out_channels)
 
         super(Conv2dReLU, self).__init__(conv, norm, relu)
 
@@ -180,9 +193,3 @@ def __init__(self, name, **params):
 
     def forward(self, x):
         return self.attention(x)
-
-
-if __name__ == "__main__":
-    print(Conv2dReLU(3, 12, 4))
-    print(Conv2dReLU(3, 12, 4, use_norm={"type": "batchnorm"}))
-    print(Conv2dReLU(3, 12, 4, use_norm={"type": "layernorm", "eps": 1e-3}))

segmentation_models_pytorch/decoders/linknet/decoder.py

@@ -10,7 +10,6 @@ def __init__(
         self,
         in_channels: int,
         out_channels: int,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
     ):
         super().__init__()
@@ -21,7 +20,7 @@ def __init__(
             nn.ReLU(inplace=True),
         ]
 
-        if use_batchnorm or use_norm:
+        if use_norm:
             layers.insert(1, nn.BatchNorm2d(out_channels))
 
         super().__init__(*layers)
@@ -32,7 +31,6 @@ def __init__(
         self,
         in_channels: int,
         out_channels: int,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
     ):
         super().__init__()
@@ -42,17 +40,15 @@ def __init__(
                 in_channels,
                 in_channels // 4,
                 kernel_size=1,
-                use_batchnorm=use_batchnorm,
                 use_norm=use_norm,
             ),
             TransposeX2(
-                in_channels // 4, in_channels // 4, use_batchnorm=use_batchnorm
+                in_channels // 4, in_channels // 4, use_norm=use_norm
             ),
             modules.Conv2dReLU(
                 in_channels // 4,
                 out_channels,
                 kernel_size=1,
-                use_batchnorm=use_batchnorm,
                 use_norm=use_norm,
             ),
         )
@@ -72,7 +68,6 @@ def __init__(
         encoder_channels: List[int],
         prefinal_channels: int = 32,
         n_blocks: int = 5,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
     ):
         super().__init__()
@@ -89,7 +84,6 @@ def __init__(
                 DecoderBlock(
                     channels[i],
                     channels[i + 1],
-                    use_batchnorm=use_batchnorm,
                     use_norm=use_norm,
                 )
                 for i in range(n_blocks)

segmentation_models_pytorch/decoders/linknet/model.py

@@ -7,6 +7,7 @@
 )
 from segmentation_models_pytorch.encoders import get_encoder
 from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
+from segmentation_models_pytorch.base.modules import handle_decoder_use_batchnorm_deprecation
 
 from .decoder import LinknetDecoder
 
@@ -101,11 +102,11 @@ def __init__(
             **kwargs,
         )
 
+        decoder_use_norm = handle_decoder_use_batchnorm_deprecation(decoder_use_batchnorm, decoder_use_norm)
         self.decoder = LinknetDecoder(
             encoder_channels=self.encoder.out_channels,
             n_blocks=encoder_depth,
             prefinal_channels=32,
-            use_batchnorm=decoder_use_batchnorm,
             use_norm=decoder_use_norm,
         )
 

segmentation_models_pytorch/decoders/manet/decoder.py

@@ -49,7 +49,6 @@ def __init__(
         in_channels: int,
         skip_channels: int,
         out_channels: int,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
         reduction: int = 16,
     ):
@@ -61,14 +60,12 @@ def __init__(
                 in_channels,
                 kernel_size=3,
                 padding=1,
-                use_batchnorm=use_batchnorm,
                 use_norm=use_norm,
             ),
             md.Conv2dReLU(
                 in_channels,
                 skip_channels,
                 kernel_size=1,
-                use_batchnorm=use_batchnorm,
                 use_norm=use_norm,
             ),
         )
@@ -93,15 +90,13 @@ def __init__(
             out_channels,
             kernel_size=3,
             padding=1,
-            use_batchnorm=use_batchnorm,
             use_norm=use_norm,
         )
         self.conv2 = md.Conv2dReLU(
             out_channels,
             out_channels,
             kernel_size=3,
             padding=1,
-            use_batchnorm=use_batchnorm,
             use_norm=use_norm,
         )
 
@@ -127,7 +122,6 @@ def __init__(
         in_channels: int,
         skip_channels: int,
         out_channels: int,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
     ):
         super().__init__()
@@ -136,15 +130,13 @@ def __init__(
             out_channels,
             kernel_size=3,
             padding=1,
-            use_batchnorm=use_batchnorm,
             use_norm=use_norm,
         )
         self.conv2 = md.Conv2dReLU(
             out_channels,
             out_channels,
             kernel_size=3,
             padding=1,
-            use_batchnorm=use_batchnorm,
             use_norm=use_norm,
         )
 
@@ -166,7 +158,6 @@ def __init__(
         decoder_channels: List[int],
         n_blocks: int = 5,
         reduction: int = 16,
-        use_batchnorm: Union[bool, str, None] = True,
         use_norm: Union[bool, str, Dict[str, Any]] = True,
         pab_channels: int = 64,
     ):
@@ -195,7 +186,7 @@ def __init__(
 
         # combine decoder keyword arguments
         kwargs = dict(
-            use_batchnorm=use_batchnorm, use_norm=use_norm
+            use_norm=use_norm
         )  # no attention type here
         blocks = [
             MFABBlock(in_ch, skip_ch, out_ch, reduction=reduction, **kwargs)

segmentation_models_pytorch/decoders/manet/model.py

@@ -7,6 +7,7 @@
 )
 from segmentation_models_pytorch.encoders import get_encoder
 from segmentation_models_pytorch.base.hub_mixin import supports_config_loading
+from segmentation_models_pytorch.base.modules import handle_decoder_use_batchnorm_deprecation
 
 from .decoder import MAnetDecoder
 
@@ -101,11 +102,12 @@ def __init__(
             **kwargs,
         )
 
+        decoder_use_norm = handle_decoder_use_batchnorm_deprecation(decoder_use_batchnorm, decoder_use_norm)
+
         self.decoder = MAnetDecoder(
             encoder_channels=self.encoder.out_channels,
             decoder_channels=decoder_channels,
             n_blocks=encoder_depth,
-            use_batchnorm=decoder_use_batchnorm,
             use_norm=decoder_use_norm,
             pab_channels=decoder_pab_channels,
         )