Set use_batchnorm default to None so that default use_norm

Make warning visible by changing filter and add a test for it
Fix test before after so that the value is looked and not the shape of tensor

Author: GuillaumeErhard

segmentation_models_pytorch/base/modules.py

@@ -40,6 +40,7 @@ def normalize_decoder_norm(decoder_use_batchnorm: Union[bool, str, None], decode
         warnings.warn(
             "The usage of use_batchnorm is deprecated. Please modify your code for use_norm",
             DeprecationWarning,
+            stacklevel=2
         )
         if decoder_use_batchnorm is True:
             decoder_use_norm = {"type": "batchnorm"}

segmentation_models_pytorch/decoders/linknet/model.py

@@ -79,7 +79,7 @@ def __init__(
         encoder_name: str = "resnet34",
         encoder_depth: int = 5,
         encoder_weights: Optional[str] = "imagenet",
-        decoder_use_batchnorm: Union[bool, str, None] = True,
+        decoder_use_batchnorm: Union[bool, str, None] = None,
         decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",
         in_channels: int = 3,
         classes: int = 1,

segmentation_models_pytorch/decoders/manet/model.py

@@ -82,7 +82,7 @@ def __init__(
         encoder_name: str = "resnet34",
         encoder_depth: int = 5,
         encoder_weights: Optional[str] = "imagenet",
-        decoder_use_batchnorm: Union[bool, str, None] = True,
+        decoder_use_batchnorm: Union[bool, str, None] = None,
         decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",
         decoder_channels: List[int] = (256, 128, 64, 32, 16),
         decoder_pab_channels: int = 64,

segmentation_models_pytorch/decoders/pspnet/model.py

@@ -81,7 +81,7 @@ def __init__(
         encoder_weights: Optional[str] = "imagenet",
         encoder_depth: int = 3,
         psp_out_channels: int = 512,
-        psp_use_batchnorm: Union[bool, str, None] = True,
+        psp_use_batchnorm: Union[bool, str, None] = None,
         decoder_use_norm: Union[bool, str, Dict[str, Any], None] = "batchnorm",
         psp_dropout: float = 0.2,
         in_channels: int = 3,

segmentation_models_pytorch/decoders/unet/model.py

@@ -114,7 +114,7 @@ def __init__(
         encoder_name: str = "resnet34",
         encoder_depth: int = 5,
         encoder_weights: Optional[str] = "imagenet",
-        decoder_use_batchnorm: Union[bool, str, None] = True,
+        decoder_use_batchnorm: Union[bool, str, None] = None,
         decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",
         decoder_channels: Sequence[int] = (256, 128, 64, 32, 16),
         decoder_attention_type: Optional[str] = None,

segmentation_models_pytorch/decoders/unetplusplus/model.py

@@ -83,7 +83,7 @@ def __init__(
         encoder_name: str = "resnet34",
         encoder_depth: int = 5,
         encoder_weights: Optional[str] = "imagenet",
-        decoder_use_batchnorm: Union[bool, str, None] = True,
+        decoder_use_batchnorm: Union[bool, str, None] = None,
         decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",
         decoder_channels: List[int] = (256, 128, 64, 32, 16),
         decoder_attention_type: Optional[str] = None,

tests/encoders/test_batchnorm_deprecation.py

@@ -10,7 +10,13 @@
 @pytest.mark.parametrize("decoder_option", [True, False, "inplace"])
 def test_seg_models_before_after_use_norm(model_name, decoder_option):
     torch.manual_seed(42)
-    model_decoder_batchnorm = create_model(model_name, "mobilenet_v2", None, decoder_use_batchnorm=decoder_option)
+    with pytest.warns(DeprecationWarning):
+        model_decoder_batchnorm = create_model(
+            model_name,
+            "mobilenet_v2",
+            None,
+            decoder_use_batchnorm=decoder_option
+        )
     torch.manual_seed(42)
     model_decoder_norm = create_model(model_name, "mobilenet_v2", None, decoder_use_batchnorm=None, decoder_use_norm=decoder_option)
 
@@ -21,7 +27,13 @@ def test_seg_models_before_after_use_norm(model_name, decoder_option):
 @pytest.mark.parametrize("decoder_option", [True, False, "inplace"])
 def test_pspnet_before_after_use_norm(decoder_option):
     torch.manual_seed(42)
-    model_decoder_batchnorm = create_model("pspnet", "mobilenet_v2", None, psp_use_batchnorm=decoder_option)
+    with pytest.warns(DeprecationWarning):
+        model_decoder_batchnorm = create_model(
+            "pspnet",
+            "mobilenet_v2",
+            None,
+            psp_use_batchnorm=decoder_option
+        )
     torch.manual_seed(42)
     model_decoder_norm = create_model("pspnet", "mobilenet_v2", None, psp_use_batchnorm=None, decoder_use_norm=decoder_option)
 

tests/utils.py

@@ -60,8 +60,8 @@ def check_run_test_on_diff_or_main(filepath_patterns: List[str]):
     return False
 
 
-def check_two_models_strictly_equal(model_a, model_b):
+def check_two_models_strictly_equal(model_a: torch.nn.Module, model_b: torch.nn.Module) -> None:
     for (k1, v1), (k2, v2) in zip(model_a.state_dict().items(),
                                   model_b.state_dict().items()):
         assert k1 == k2, f"Key mismatch: {k1} != {k2}"
-        assert v1.shape == v2.shape, f"Shape mismatch in {k1}: {v1.shape} != {v2.shape}"
+        assert (v1 == v2).all(), f"Tensor mismatch at key '{k1}':\n{v1} !=\n{v2}"