Deprecate use_batchnorm in favor of generalized use_norm parameter #1095

GuillaumeErhard · 2025-03-19T23:56:10Z

Hello,

a bit late to implement my suggestion in #983 but here it is. Open to any suggestion to make it clearer / simpler or any wanted test.

Closes #983

codecov · 2025-03-20T07:42:44Z

Codecov Report

Attention: Patch coverage is 60.34483% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
segmentation_models_pytorch/base/modules.py	46.51%	23 Missing ⚠️

Files with missing lines	Coverage Δ
...ntation_models_pytorch/decoders/linknet/decoder.py	`100.00% <100.00%> (ø)`
...mentation_models_pytorch/decoders/linknet/model.py	`94.73% <100.00%> (ø)`
...mentation_models_pytorch/decoders/manet/decoder.py	`97.75% <100.00%> (ø)`
...egmentation_models_pytorch/decoders/manet/model.py	`100.00% <100.00%> (ø)`
...gmentation_models_pytorch/decoders/unet/decoder.py	`91.37% <100.00%> (ø)`
segmentation_models_pytorch/decoders/unet/model.py	`100.00% <100.00%> (ø)`
...on_models_pytorch/decoders/unetplusplus/decoder.py	`92.85% <100.00%> (ø)`
...tion_models_pytorch/decoders/unetplusplus/model.py	`95.00% <100.00%> (ø)`
segmentation_models_pytorch/base/modules.py	`50.45% <46.51%> (-4.95%)`	⬇️

🚀 New features to boost your workflow:

❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

qubvel

Hi @GuillaumeErhard, huge thanks for taking time and working on the PR! Here a few comments

qubvel · 2025-03-20T12:20:30Z

segmentation_models_pytorch/base/modules.py

+        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")
+
+        if use_norm["type"] == "inplace" and InPlaceABN is None:
            raise RuntimeError(
-                "In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. "
-                + "To install see: https://github.com/mapillary/inplace_abn"
+                "In order to use `use_batchnorm='inplace'` or `use_norm='inplace'` the inplace_abn package must be installed. "
+                "To install see: https://github.com/mapillary/inplace_abn"
            )


Let's have a separate function get_norm_layer which will validate input params and return norm layer

Good catch. Much simpler. Done the changes

qubvel · 2025-03-20T12:21:05Z

segmentation_models_pytorch/base/modules.py

+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}))


Should be removed, instead, it would be nice to add a test

Yeap, added some test around Conv2dReLu

qubvel · 2025-03-20T12:22:40Z

segmentation_models_pytorch/decoders/linknet/decoder.py

+        use_batchnorm: Union[bool, str, None] = True,
+        use_norm: Union[bool, str, Dict[str, Any]] = True,


We should have only use_norm here, use_bathcnorm should be replaced on top level (e.g. Unet model)

Made a proposition

qubvel · 2025-03-20T12:23:19Z

segmentation_models_pytorch/decoders/linknet/decoder.py

+    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,
+    ):


same here, let's move it to the model, and leave only use_norm

Made a proposition

GuillaumeErhard · 2025-03-21T00:16:47Z

I am missing some test to check before and after equivalency when creating model using decoder_use_batchnorm and decoder_use_norm. I will do this tomorrow. ( tests/encoders/test_batchnorm_deprecation.py )

But I have hit an error that I did not anticipate with the test. (Thank god for test driven dev). GroupNorm uses num_groups, num_channels param instead of using a simple out_channels As of right now I don't see a good way to make it work so I am thinking of dropping the support for it. If you have any idea, I am open for it

Fix issue in Linknet Handle norm and relu on itsd own Align pspnet Add use_norm in upernet Remove groupnorm possibility Update doc

GuillaumeErhard · 2025-03-24T23:46:42Z

Added the test so that we get the same model before and after the changes. Which showed that Linknet additional changes which are now done.

I respected that all element in model/decoder can still work use_norm can be str / bool / dict for the moment in the case people use them directly. For that I need to sanitize at the model level and on Conv2DRelu. Tell me if you want me to change to dict only in decoder

I also saw that pspnet had also batchnorm usage but through psp_use_batchnorm and aligned it with use_norm.
By the way maybe on line 20 of pspnet/decoder.py it can be changed for a norm instead of identity ?

Removed groupnorm for the moment. Taking any suggestion to make it possible. See problem described above

Added support in upernet to also use_norm.
Is the behavior on l 35 expected ? That is always batchnorm ? Or I can pass use_norm also ?

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

qubvel

Thanks for addressing comments, please see the review 🤗

qubvel · 2025-03-26T14:33:46Z

segmentation_models_pytorch/base/modules.py

+def normalize_use_norm(decoder_use_norm: Union[bool, str, Dict[str, Any]]) -> Dict[str, Any]:
+    if isinstance(decoder_use_norm, str):
+        norm_str = decoder_use_norm.lower()
+        if norm_str == "inplace":
+            decoder_use_norm = {
+                "type": "inplace",
+                "activation": "leaky_relu",
+                "activation_param": 0.0,
+            }
+        elif norm_str in (
+                "batchnorm",
+                "identity",
+                "layernorm",
+                "groupnorm",
+                "instancenorm",
+        ):
+            decoder_use_norm = {"type": norm_str}
+        else:
+            raise ValueError("Unrecognized normalization type string provided")
+    elif isinstance(decoder_use_norm, bool):
+        decoder_use_norm = {"type": "batchnorm" if decoder_use_norm else "identity"}
+    elif not isinstance(decoder_use_norm, dict):
+        raise ValueError("use_norm must be a dictionary, boolean, or string")
+
+    return decoder_use_norm
+
+def normalize_decoder_norm(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,
+            stacklevel=2
+        )
+        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")
+
+    decoder_use_norm = normalize_use_norm(decoder_use_norm)
+    return decoder_use_norm
+
+
+def get_norm_layer(use_norm: Dict[str, Any], out_channels: int) -> 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)
+    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 == "instancenorm":
+        norm = nn.InstanceNorm2d(out_channels, **extra_kwargs)
+    else:
+        raise ValueError(f"Unrecognized normalization type: {norm_type}")
+
+    return norm


I would rather put everything into a single simple function as follows:

if use_norm is True: params = {"type": "batchnorm"} elif use_norm is False: params = {"type": "identity"} elif use_norm == "inplace": params = {"type": "inplace", "activation": "leaky_relu", "activation_param": 0.0} if not isinstance(params, dict) -> raise if not "type" in params -> raise if not params["type"] in supproted_norms -> raise <dispatch to norm layer> return norm

So no str for default usage like LayerNorm etc... ??? I kept it for the moment. Tell me if I have to remove.

Did changed accordingly and tried to match the proposed flow as closely as possible

qubvel · 2025-03-26T14:40:59Z

segmentation_models_pytorch/base/modules.py

+        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'` inplace_abn package must be installed. "
-                + "To install see: https://github.com/mapillary/inplace_abn"
+                "In order to use `use_batchnorm='inplace'` or `use_norm='inplace'` the inplace_abn package must be installed. "
+                "To install see: https://github.com/mapillary/inplace_abn"


Also not needed here, should be under get_norm_layer function

Removed and placed in get_norm_layer

qubvel · 2025-03-26T14:42:02Z

segmentation_models_pytorch/base/modules.py

@@ -29,21 +101,16 @@ def __init__(
            kernel_size,
            stride=stride,
            padding=padding,
-            bias=not (use_batchnorm),
+            bias=use_norm["type"] != "inplace",


We can initialize norm first and use a separate varaible

is_inplace_batchnorm = norm.__name__ == "InPlaceABN"

yeap simpler

qubvel · 2025-03-26T14:43:15Z

segmentation_models_pytorch/base/modules.py


-        if use_batchnorm == "inplace":
-            bn = InPlaceABN(out_channels, activation="leaky_relu", activation_param=0.0)
+        if use_norm["type"] == "inplace":


qubvel · 2025-03-26T14:44:00Z

segmentation_models_pytorch/decoders/linknet/model.py

@@ -60,7 +79,8 @@ def __init__(
        encoder_name: str = "resnet34",
        encoder_depth: int = 5,
        encoder_weights: Optional[str] = "imagenet",
-        decoder_use_batchnorm: bool = True,
+        decoder_use_batchnorm: Union[bool, str, None] = None,


let's removedecoder_use_batchnorm from signature and pop from krawgs later

removed decoder_use_batchnorm and used kwargs. I like this approach

qubvel · 2025-03-26T14:45:11Z

segmentation_models_pytorch/decoders/linknet/model.py

@@ -82,11 +102,12 @@ def __init__(
            **kwargs,
        )

+        decoder_use_norm = normalize_decoder_norm(decoder_use_batchnorm, decoder_use_norm)


here we just resolve the name

Tell me if this is what you had in mind

qubvel · 2025-03-26T14:47:11Z

segmentation_models_pytorch/decoders/manet/model.py

+        decoder_use_batchnorm: Union[bool, str, None] = None,
+        decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",


Let's keep default

decoder_use_norm: Union[bool, str, Dict[str, Any]] = True

and remove decoder_use_batchnorm entirely from signature, it will be passed in kwargs and poped from there

I am not really that sold with True as default. "batchnorm" seems more meaningful and the bool more of an historic usage.
Changed it to True on model. Tell me if I have to propagate in all the decoder

removed decoder_use_batchnorm and used kwargs. I like this approach

"batchnorm" seems more meaningful and the bool more of an historic usage.

Ok, sounds good!

qubvel · 2025-03-26T14:50:02Z

tests/base/test_modules.py

+def test_conv2drelu_batchnorm():
+    module = Conv2dReLU(3, 16, kernel_size=3, padding=1, use_norm="batchnorm")
+
+    expected = ('Conv2dReLU(\n  (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))'
+                '\n  (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)'
+                '\n  (2): ReLU(inplace=True)\n)')
+    assert repr(module) == expected


Not super robust, formatting can changed in any version, lets use

assert isinstance(module[1], nn.BatchNorm2d)

and below please

Changed asserts

qubvel · 2025-03-26T14:51:02Z

tests/encoders/test_batchnorm_deprecation.py

+@pytest.mark.parametrize("decoder_option", [True, False, "inplace"])
+def test_pspnet_before_after_use_norm(decoder_option):
+    torch.manual_seed(42)
+    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)
+
+    check_two_models_strictly_equal(model_decoder_batchnorm, model_decoder_norm)


qubvel · 2025-03-26T14:51:41Z

tests/utils.py

+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 == v2).all(), f"Tensor mismatch at key '{k1}':\n{v1} !=\n{v2}"


I would also add forward pass here and compare logits

Add input_data param to check that

qubvel

Thanks, just a few minor comments and we are good to merge!

qubvel · 2025-03-28T14:18:55Z

segmentation_models_pytorch/base/modules.py

+    elif isinstance(use_norm, dict):
+        norm_params = use_norm
+    else:
+        raise ValueError("use_norm must be a dictionary, boolean, or string. Please refer to the documentation.")


Let's have a more descriptive error here, I mean specify what kind of string and dict structure it should be.

qubvel · 2025-03-28T14:19:04Z

segmentation_models_pytorch/base/modules.py

+        ):
+            norm_params = {"type": norm_str}
+        else:
+            raise ValueError(f"Unrecognized normalization type string provided: {use_norm}. Should be in {supported_norms}")


qubvel · 2025-03-28T14:22:29Z

segmentation_models_pytorch/decoders/linknet/model.py

+        decoder_use_norm:     Specifies normalization between Conv2D and activation.
+            Accepts the following types:
+            - **True**: Defaults to `"batchnorm"`.
+            - **False**: No normalization (`nn.Identity`).
+            - **str**: Specifies normalization type using default parameters. Available values:
+              `"batchnorm"`, `"identity"`, `"layernorm"`, `"instancenorm"`, `"inplace"`.
+            - **dict**: Fully customizable normalization settings. Structure:
+              ```python
+              {"type": <norm_type>, **kwargs}
+              ```
+              where `norm_name` corresponds to normalization type (see above), and `kwargs` are passed directly to the normalization layer as defined in PyTorch documentation.
+
+            **Example**:
+            ```python
+            use_norm={"type": "layernorm", "eps": 1e-2}
+            ```


Thanks for the detailed docstring, really appretiate it!

qubvel · 2025-03-28T14:24:34Z

segmentation_models_pytorch/decoders/pspnet/model.py

+
+            **Example**:
+            ```python
+            use_norm={"type": "layernorm", "eps": 1e-2}


Suggested change

use_norm={"type": "layernorm", "eps": 1e-2}

decoder_use_norm={"type": "layernorm", "eps": 1e-2}

GuillaumeErhard added 2 commits March 20, 2025 00:52

Deprecate use_batchnorm in favor of generalized use_norm parameter

e26adcd

Merge remote-tracking branch 'upstream/main' into new-norm

d65001b

qubvel reviewed Mar 20, 2025

View reviewed changes

First fix following review

1b16b25

Add test to verify before / after functionality

10d496a

Fix issue in Linknet Handle norm and relu on itsd own Align pspnet Add use_norm in upernet Remove groupnorm possibility Update doc

Set use_batchnorm default to None so that default use_norm

467057a

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

qubvel reviewed Mar 26, 2025

View reviewed changes

Changes following review

1ae11c3

qubvel approved these changes Mar 28, 2025

View reviewed changes

		use_batchnorm: Union[bool, str, None] = True,
		use_norm: Union[bool, str, Dict[str, Any]] = True,

		decoder_use_batchnorm: Union[bool, str, None] = None,
		decoder_use_norm: Union[bool, str, Dict[str, Any]] = "batchnorm",

	use_norm={"type": "layernorm", "eps": 1e-2}
	decoder_use_norm={"type": "layernorm", "eps": 1e-2}

Deprecate use_batchnorm in favor of generalized use_norm parameter #1095

Are you sure you want to change the base?

Deprecate use_batchnorm in favor of generalized use_norm parameter #1095

Conversation

GuillaumeErhard commented Mar 19, 2025 • edited Loading

codecov bot commented Mar 20, 2025

Codecov Report

qubvel left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GuillaumeErhard commented Mar 21, 2025 • edited Loading

GuillaumeErhard commented Mar 24, 2025

qubvel left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

GuillaumeErhard commented Mar 19, 2025 •

edited

Loading

GuillaumeErhard commented Mar 21, 2025 •

edited

Loading