Repository: Project-MONAI/MONAI

Issue:
Support **kwargs for some pytorch args of networks
For networks like `SegResNet`, we cannot change parameters for normalization layers (such as set `track_running_stats=True` when we use BatchNorm3d). Therefore, the implementation should be more generic. 


Discussion:

Patch:
diff --git a/monai/networks/blocks/acti_norm.py b/monai/networks/blocks/acti_norm.py
--- a/monai/networks/blocks/acti_norm.py
+++ b/monai/networks/blocks/acti_norm.py
@@ -13,8 +13,8 @@
 
 import torch.nn as nn
 
-from monai.networks.layers.factories import Act, Dropout, Norm, split_args
-from monai.utils import has_option
+from monai.networks.blocks.utils import get_act_layer, get_norm_layer
+from monai.networks.layers.factories import Dropout, split_args
 
 
 class ADN(nn.Sequential):
@@ -82,22 +82,15 @@ def __init__(
         op_dict = {"A": None, "D": None, "N": None}
         # define the normalization type and the arguments to the constructor
         if norm is not None:
-            if norm_dim is None and dropout_dim is None:
+            spatial_dims = norm_dim or dropout_dim
+            if isinstance(spatial_dims, int):
+                op_dict["N"] = get_norm_layer(spatial_dims=spatial_dims, channels=in_channels, norm=norm)
+            else:
                 raise ValueError("norm_dim or dropout_dim needs to be specified.")
-            norm_name, norm_args = split_args(norm)
-            norm_type = Norm[norm_name, norm_dim or dropout_dim]
-            kw_args = dict(norm_args)
-            if has_option(norm_type, "num_features") and "num_features" not in kw_args:
-                kw_args["num_features"] = in_channels
-            if has_option(norm_type, "num_channels") and "num_channels" not in kw_args:
-                kw_args["num_channels"] = in_channels
-            op_dict["N"] = norm_type(**kw_args)
 
         # define the activation type and the arguments to the constructor
         if act is not None:
-            act_name, act_args = split_args(act)
-            act_type = Act[act_name]
-            op_dict["A"] = act_type(**act_args)
+            op_dict["A"] = get_act_layer(act)
 
         if dropout is not None:
             # if dropout was specified simply as a p value, use default name and make a keyword map with the value
diff --git a/monai/networks/blocks/dynunet_block.py b/monai/networks/blocks/dynunet_block.py
--- a/monai/networks/blocks/dynunet_block.py
+++ b/monai/networks/blocks/dynunet_block.py
@@ -9,14 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Optional, Sequence, Tuple, Union
+from typing import Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
 import torch.nn as nn
 
 from monai.networks.blocks.convolutions import Convolution
-from monai.networks.layers.factories import Act, Norm, split_args
+from monai.networks.blocks.utils import get_act_layer, get_norm_layer
+from monai.networks.layers.factories import Act, Norm
 
 
 class UnetResBlock(nn.Module):
@@ -31,9 +32,8 @@ class UnetResBlock(nn.Module):
         out_channels: number of output channels.
         kernel_size: convolution kernel size.
         stride: convolution stride.
-        norm_name: [``"batch"``, ``"instance"``, ``"group"``]
-            feature normalization type and arguments. In this module, if using ``"group"``,
-            `in_channels` should be divisible by 16 (default value for ``num_groups``).
+        norm_name: feature normalization type and arguments.
+
     """
 
     def __init__(
@@ -43,7 +43,7 @@ def __init__(
         out_channels: int,
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
-        norm_name: str,
+        norm_name: Union[Tuple, str],
     ):
         super(UnetResBlock, self).__init__()
         self.conv1 = get_conv_layer(
@@ -70,10 +70,10 @@ def __init__(
             stride=stride,
             conv_only=True,
         )
-        self.lrelu = get_acti_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
-        self.norm1 = get_norm_layer(spatial_dims, out_channels, norm_name)
-        self.norm2 = get_norm_layer(spatial_dims, out_channels, norm_name)
-        self.norm3 = get_norm_layer(spatial_dims, out_channels, norm_name)
+        self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
+        self.norm1 = get_norm_layer(spatial_dims, channels=out_channels, norm=norm_name)
+        self.norm2 = get_norm_layer(spatial_dims, channels=out_channels, norm=norm_name)
+        self.norm3 = get_norm_layer(spatial_dims, channels=out_channels, norm=norm_name)
         self.downsample = in_channels != out_channels
         stride_np = np.atleast_1d(stride)
         if not np.all(stride_np == 1):
@@ -106,9 +106,8 @@ class UnetBasicBlock(nn.Module):
         out_channels: number of output channels.
         kernel_size: convolution kernel size.
         stride: convolution stride.
-        norm_name: [``"batch"``, ``"instance"``, ``"group"``]
-            feature normalization type and arguments. In this module, if using ``"group"``,
-            `in_channels` should be divisible by 16 (default value for ``num_groups``).
+        norm_name: feature normalization type and arguments.
+
     """
 
     def __init__(
@@ -118,7 +117,7 @@ def __init__(
         out_channels: int,
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
-        norm_name: str,
+        norm_name: Union[Tuple, str],
     ):
         super(UnetBasicBlock, self).__init__()
         self.conv1 = get_conv_layer(
@@ -137,9 +136,9 @@ def __init__(
             stride=1,
             conv_only=True,
         )
-        self.lrelu = get_acti_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
-        self.norm1 = get_norm_layer(spatial_dims, out_channels, norm_name)
-        self.norm2 = get_norm_layer(spatial_dims, out_channels, norm_name)
+        self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
+        self.norm1 = get_norm_layer(spatial_dims, channels=out_channels, norm=norm_name)
+        self.norm2 = get_norm_layer(spatial_dims, channels=out_channels, norm=norm_name)
 
     def forward(self, inp):
         out = self.conv1(inp)
@@ -164,9 +163,8 @@ class UnetUpBlock(nn.Module):
         kernel_size: convolution kernel size.
         stride: convolution stride.
         upsample_kernel_size: convolution kernel size for transposed convolution layers.
-        norm_name: [``"batch"``, ``"instance"``, ``"group"``]
-            feature normalization type and arguments. In this module, if using ``"group"``,
-            `in_channels` should be divisible by 16 (default value for ``num_groups``).
+        norm_name: feature normalization type and arguments.
+
     """
 
     def __init__(
@@ -177,7 +175,7 @@ def __init__(
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
         upsample_kernel_size: Union[Sequence[int], int],
-        norm_name: str,
+        norm_name: Union[Tuple, str],
     ):
         super(UnetUpBlock, self).__init__()
         upsample_stride = upsample_kernel_size
@@ -219,24 +217,6 @@ def forward(self, inp):
         return out
 
 
-def get_acti_layer(act: Union[Tuple[str, Dict], str]):
-    act_name, act_args = split_args(act)
-    act_type = Act[act_name]
-    return act_type(**act_args)
-
-
-def get_norm_layer(spatial_dims: int, out_channels: int, norm_name: str, num_groups: int = 16):
-    if norm_name not in ["batch", "instance", "group"]:
-        raise ValueError(f"Unsupported normalization mode: {norm_name}")
-    if norm_name == "group":
-        if out_channels % num_groups != 0:
-            raise AssertionError("out_channels should be divisible by num_groups.")
-        norm = Norm[norm_name](num_groups=num_groups, num_channels=out_channels, affine=True)
-    else:
-        norm = Norm[norm_name, spatial_dims](out_channels, affine=True)
-    return norm
-
-
 def get_conv_layer(
     spatial_dims: int,
     in_channels: int,
diff --git a/monai/networks/blocks/segresnet_block.py b/monai/networks/blocks/segresnet_block.py
--- a/monai/networks/blocks/segresnet_block.py
+++ b/monai/networks/blocks/segresnet_block.py
@@ -9,30 +9,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Union
+from typing import Tuple, Union
 
 import torch.nn as nn
 
 from monai.networks.blocks.convolutions import Convolution
 from monai.networks.blocks.upsample import UpSample
-from monai.networks.layers.factories import Act, Norm
+from monai.networks.blocks.utils import get_norm_layer
+from monai.networks.layers.factories import Act
 from monai.utils import InterpolateMode, UpsampleMode
 
 
-def get_norm_layer(spatial_dims: int, in_channels: int, norm_name: str, num_groups: int = 8):
-    if norm_name not in ["batch", "instance", "group"]:
-        raise ValueError(f"Unsupported normalization mode: {norm_name}")
-    if norm_name == "group":
-        norm = Norm[norm_name](num_groups=num_groups, num_channels=in_channels)
-    else:
-        norm = Norm[norm_name, spatial_dims](in_channels)
-    if norm.bias is not None:
-        nn.init.zeros_(norm.bias)
-    if norm.weight is not None:
-        nn.init.ones_(norm.weight)
-    return norm
-
-
 def get_conv_layer(
     spatial_dims: int, in_channels: int, out_channels: int, kernel_size: int = 3, stride: int = 1, bias: bool = False
 ):
@@ -73,33 +60,27 @@ def __init__(
         self,
         spatial_dims: int,
         in_channels: int,
+        norm: Union[Tuple, str],
         kernel_size: int = 3,
-        norm_name: str = "group",
-        num_groups: int = 8,
     ) -> None:
         """
         Args:
             spatial_dims: number of spatial dimensions, could be 1, 2 or 3.
             in_channels: number of input channels.
+            norm: feature normalization type and arguments.
             kernel_size: convolution kernel size, the value should be an odd number. Defaults to 3.
-            norm_name: feature normalization type, this module only supports group norm,
-                batch norm and instance norm. Defaults to ``group``.
-            num_groups: number of groups to separate the channels into, in this module,
-                in_channels should be divisible by num_groups. Defaults to 8.
         """
 
         super().__init__()
 
         if kernel_size % 2 != 1:
             raise AssertionError("kernel_size should be an odd number.")
-        if in_channels % num_groups != 0:
-            raise AssertionError("in_channels should be divisible by num_groups.")
 
-        self.norm1 = get_norm_layer(spatial_dims, in_channels, norm_name, num_groups=num_groups)
-        self.norm2 = get_norm_layer(spatial_dims, in_channels, norm_name, num_groups=num_groups)
+        self.norm1 = get_norm_layer(spatial_dims, channels=in_channels, norm=norm)
+        self.norm2 = get_norm_layer(spatial_dims, channels=in_channels, norm=norm)
         self.relu = Act[Act.RELU](inplace=True)
-        self.conv1 = get_conv_layer(spatial_dims, in_channels, in_channels)
-        self.conv2 = get_conv_layer(spatial_dims, in_channels, in_channels)
+        self.conv1 = get_conv_layer(spatial_dims, in_channels=in_channels, out_channels=in_channels)
+        self.conv2 = get_conv_layer(spatial_dims, in_channels=in_channels, out_channels=in_channels)
 
     def forward(self, x):
 
diff --git a/monai/networks/blocks/utils.py b/monai/networks/blocks/utils.py
new file mode 100644
--- /dev/null
+++ b/monai/networks/blocks/utils.py
@@ -0,0 +1,32 @@
+# Copyright 2020 - 2021 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple, Union
+
+from monai.networks.layers.factories import Act, Norm, split_args
+from monai.utils import has_option
+
+
+def get_norm_layer(spatial_dims: int, norm: Union[Tuple, str], channels: Optional[int] = None):
+    norm_name, norm_args = split_args(norm)
+    norm_type = Norm[norm_name, spatial_dims]
+    kw_args = dict(norm_args)
+    if has_option(norm_type, "num_features") and "num_features" not in kw_args:
+        kw_args["num_features"] = channels
+    if has_option(norm_type, "num_channels") and "num_channels" not in kw_args:
+        kw_args["num_channels"] = channels
+    return norm_type(**kw_args)
+
+
+def get_act_layer(act: Union[Tuple, str]):
+    act_name, act_args = split_args(act)
+    act_type = Act[act_name]
+    return act_type(**act_args)
diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py
--- a/monai/networks/nets/dynunet.py
+++ b/monai/networks/nets/dynunet.py
@@ -10,7 +10,7 @@
 # limitations under the License.
 
 
-from typing import List, Optional, Sequence, Union
+from typing import List, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -79,8 +79,7 @@ class DynUNet(nn.Module):
         kernel_size: convolution kernel size.
         strides: convolution strides for each blocks.
         upsample_kernel_size: convolution kernel size for transposed convolution layers.
-        norm_name: [``"batch"``, ``"instance"``, ``"group"``]
-            feature normalization type and arguments.
+        norm_name: feature normalization type and arguments. Defaults to ``INSTANCE``.
         deep_supervision: whether to add deep supervision head before output. Defaults to ``False``.
             If ``True``, in training mode, the forward function will output not only the last feature
             map, but also the previous feature maps that come from the intermediate up sample layers.
@@ -109,7 +108,7 @@ def __init__(
         kernel_size: Sequence[Union[Sequence[int], int]],
         strides: Sequence[Union[Sequence[int], int]],
         upsample_kernel_size: Sequence[Union[Sequence[int], int]],
-        norm_name: str = "instance",
+        norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}),
         deep_supervision: bool = False,
         deep_supr_num: int = 1,
         res_block: bool = False,
diff --git a/monai/networks/nets/segresnet.py b/monai/networks/nets/segresnet.py
--- a/monai/networks/nets/segresnet.py
+++ b/monai/networks/nets/segresnet.py
@@ -9,15 +9,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Sequence, Union
+from typing import Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from monai.networks.blocks.segresnet_block import ResBlock, get_conv_layer, get_norm_layer, get_upsample_layer
-from monai.networks.layers.factories import Act, Dropout
+from monai.networks.blocks.segresnet_block import ResBlock, get_conv_layer, get_upsample_layer
+from monai.networks.blocks.utils import get_act_layer, get_norm_layer
+from monai.networks.layers.factories import Dropout
 from monai.utils import UpsampleMode
 
 
@@ -34,9 +35,8 @@ class SegResNet(nn.Module):
         in_channels: number of input channels for the network. Defaults to 1.
         out_channels: number of output channels for the network. Defaults to 2.
         dropout_prob: probability of an element to be zero-ed. Defaults to ``None``.
-        norm_name: feature normalization type, this module only supports group norm,
-            batch norm and instance norm. Defaults to ``group``.
-        num_groups: number of groups to separate the channels into. Defaults to 8.
+        act: activation type and arguments. Defaults to ``RELU``.
+        norm: feature normalization type and arguments. Defaults to ``GROUP``.
         use_conv_final: if add a final convolution block to output. Defaults to ``True``.
         blocks_down: number of down sample blocks in each layer. Defaults to ``[1,2,2,4]``.
         blocks_up: number of up sample blocks in each layer. Defaults to ``[1,1,1]``.
@@ -57,8 +57,8 @@ def __init__(
         in_channels: int = 1,
         out_channels: int = 2,
         dropout_prob: Optional[float] = None,
-        norm_name: str = "group",
-        num_groups: int = 8,
+        act: Union[Tuple, str] = ("RELU", {"inplace": True}),
+        norm: Union[Tuple, str] = ("GROUP", {"num_groups": 8}),
         use_conv_final: bool = True,
         blocks_down: tuple = (1, 2, 2, 4),
         blocks_up: tuple = (1, 1, 1),
@@ -71,17 +71,17 @@ def __init__(
 
         self.spatial_dims = spatial_dims
         self.init_filters = init_filters
+        self.in_channels = in_channels
         self.blocks_down = blocks_down
         self.blocks_up = blocks_up
         self.dropout_prob = dropout_prob
-        self.norm_name = norm_name
-        self.num_groups = num_groups
+        self.act = get_act_layer(act)
+        self.norm = norm
         self.upsample_mode = UpsampleMode(upsample_mode)
         self.use_conv_final = use_conv_final
         self.convInit = get_conv_layer(spatial_dims, in_channels, init_filters)
         self.down_layers = self._make_down_layers()
         self.up_layers, self.up_samples = self._make_up_layers()
-        self.relu = Act[Act.RELU](inplace=True)
         self.conv_final = self._make_final_conv(out_channels)
 
         if dropout_prob is not None:
@@ -89,12 +89,11 @@ def __init__(
 
     def _make_down_layers(self):
         down_layers = nn.ModuleList()
-        blocks_down, spatial_dims, filters, norm_name, num_groups = (
+        blocks_down, spatial_dims, filters, norm = (
             self.blocks_down,
             self.spatial_dims,
             self.init_filters,
-            self.norm_name,
-            self.num_groups,
+            self.norm,
         )
         for i in range(len(blocks_down)):
             layer_in_channels = filters * 2 ** i
@@ -105,33 +104,26 @@ def _make_down_layers(self):
             )
             down_layer = nn.Sequential(
                 pre_conv,
-                *[
-                    ResBlock(spatial_dims, layer_in_channels, norm_name=norm_name, num_groups=num_groups)
-                    for _ in range(blocks_down[i])
-                ],
+                *[ResBlock(spatial_dims, layer_in_channels, norm=norm) for _ in range(blocks_down[i])],
             )
             down_layers.append(down_layer)
         return down_layers
 
     def _make_up_layers(self):
         up_layers, up_samples = nn.ModuleList(), nn.ModuleList()
-        upsample_mode, blocks_up, spatial_dims, filters, norm_name, num_groups = (
+        upsample_mode, blocks_up, spatial_dims, filters, norm = (
             self.upsample_mode,
             self.blocks_up,
             self.spatial_dims,
             self.init_filters,
-            self.norm_name,
-            self.num_groups,
+            self.norm,
         )
         n_up = len(blocks_up)
         for i in range(n_up):
             sample_in_channels = filters * 2 ** (n_up - i)
             up_layers.append(
                 nn.Sequential(
-                    *[
-                        ResBlock(spatial_dims, sample_in_channels // 2, norm_name=norm_name, num_groups=num_groups)
-                        for _ in range(blocks_up[i])
-                    ]
+                    *[ResBlock(spatial_dims, sample_in_channels // 2, norm=norm) for _ in range(blocks_up[i])]
                 )
             )
             up_samples.append(
@@ -146,9 +138,9 @@ def _make_up_layers(self):
 
     def _make_final_conv(self, out_channels: int):
         return nn.Sequential(
-            get_norm_layer(self.spatial_dims, self.init_filters, norm_name=self.norm_name, num_groups=self.num_groups),
-            self.relu,
-            get_conv_layer(self.spatial_dims, self.init_filters, out_channels=out_channels, kernel_size=1, bias=True),
+            get_norm_layer(self.spatial_dims, channels=self.init_filters, norm=self.norm),
+            self.act,
+            get_conv_layer(self.spatial_dims, self.init_filters, out_channels, kernel_size=1, bias=True),
         )
 
     def forward(self, x):
@@ -192,15 +184,14 @@ class SegResNetVAE(SegResNet):
         in_channels: number of input channels for the network. Defaults to 1.
         out_channels: number of output channels for the network. Defaults to 2.
         dropout_prob: probability of an element to be zero-ed. Defaults to ``None``.
-        norm_name: feature normalization type, this module only supports group norm,
-            batch norm and instance norm. Defaults to ``group``.
-        num_groups: number of groups to separate the channels into. Defaults to 8.
+        act: activation type and arguments. Defaults to ``RELU``.
+        norm: feature normalization type and arguments. Defaults to ``GROUP``.
         use_conv_final: if add a final convolution block to output. Defaults to ``True``.
         blocks_down: number of down sample blocks in each layer. Defaults to ``[1,2,2,4]``.
         blocks_up: number of up sample blocks in each layer. Defaults to ``[1,1,1]``.
         upsample_mode: [``"deconv"``, ``"nontrainable"``, ``"pixelshuffle"``]
             The mode of upsampling manipulations.
-            Using the ``nontrainable`` modes cannot guarantee the model's reproducibility. Defaults to `nontrainable`.
+            Using the ``nontrainable`` modes cannot guarantee the model's reproducibility. Defaults to``nontrainable``.
 
             - ``deconv``, uses transposed convolution layers.
             - ``nontrainable``, uses non-trainable `linear` interpolation.
@@ -218,12 +209,12 @@ def __init__(
         in_channels: int = 1,
         out_channels: int = 2,
         dropout_prob: Optional[float] = None,
-        norm_name: str = "group",
-        num_groups: int = 8,
+        act: Union[str, tuple] = ("RELU", {"inplace": True}),
+        norm: Union[Tuple, str] = ("GROUP", {"num_groups": 8}),
         use_conv_final: bool = True,
         blocks_down: tuple = (1, 2, 2, 4),
         blocks_up: tuple = (1, 1, 1),
-        upsample_mode: Union[UpsampleMode, str] = "nontrainable",
+        upsample_mode: Union[UpsampleMode, str] = UpsampleMode.NONTRAINABLE,
     ):
         super(SegResNetVAE, self).__init__(
             spatial_dims=spatial_dims,
@@ -231,8 +222,7 @@ def __init__(
             in_channels=in_channels,
             out_channels=out_channels,
             dropout_prob=dropout_prob,
-            norm_name=norm_name,
-            num_groups=num_groups,
+            norm=norm,
             use_conv_final=use_conv_final,
             blocks_down=blocks_down,
             blocks_up=blocks_up,
@@ -257,13 +247,11 @@ def _prepare_vae_modules(self):
         total_elements = int(self.smallest_filters * np.prod(self.fc_insize))
 
         self.vae_down = nn.Sequential(
-            get_norm_layer(self.spatial_dims, v_filters, norm_name=self.norm_name, num_groups=self.num_groups),
-            self.relu,
+            get_norm_layer(self.spatial_dims, channels=v_filters, norm=self.norm),
+            self.act,
             get_conv_layer(self.spatial_dims, v_filters, self.smallest_filters, stride=2, bias=True),
-            get_norm_layer(
-                self.spatial_dims, self.smallest_filters, norm_name=self.norm_name, num_groups=self.num_groups
-            ),
-            self.relu,
+            get_norm_layer(self.spatial_dims, channels=self.smallest_filters, norm=self.norm),
+            self.act,
         )
         self.vae_fc1 = nn.Linear(total_elements, self.vae_nz)
         self.vae_fc2 = nn.Linear(total_elements, self.vae_nz)
@@ -272,8 +260,8 @@ def _prepare_vae_modules(self):
         self.vae_fc_up_sample = nn.Sequential(
             get_conv_layer(self.spatial_dims, self.smallest_filters, v_filters, kernel_size=1),
             get_upsample_layer(self.spatial_dims, v_filters, upsample_mode=self.upsample_mode),
-            get_norm_layer(self.spatial_dims, v_filters, norm_name=self.norm_name, num_groups=self.num_groups),
-            self.relu,
+            get_norm_layer(self.spatial_dims, channels=v_filters, norm=self.norm),
+            self.act,
         )
 
     def _get_vae_loss(self, net_input: torch.Tensor, vae_input: torch.Tensor):
@@ -302,7 +290,7 @@ def _get_vae_loss(self, net_input: torch.Tensor, vae_input: torch.Tensor):
             x_vae = z_mean + z_sigma * z_mean_rand
 
         x_vae = self.vae_fc3(x_vae)
-        x_vae = self.relu(x_vae)
+        x_vae = self.act(x_vae)
         x_vae = x_vae.view([-1, self.smallest_filters] + self.fc_insize)
         x_vae = self.vae_fc_up_sample(x_vae)
 


Test patch:
diff --git a/tests/test_dynunet.py b/tests/test_dynunet.py
--- a/tests/test_dynunet.py
+++ b/tests/test_dynunet.py
@@ -65,7 +65,7 @@
                 "kernel_size": (3, (1, 1, 3), 3, 3),
                 "strides": ((1, 2, 1), 2, 2, 1),
                 "upsample_kernel_size": (2, 2, 1),
-                "norm_name": "instance",
+                "norm_name": ("INSTANCE", {"affine": True}),
                 "deep_supervision": False,
                 "res_block": res_block,
             },
@@ -88,7 +88,7 @@
                         "kernel_size": [3] * len(strides),
                         "strides": strides,
                         "upsample_kernel_size": strides[1:],
-                        "norm_name": "group",
+                        "norm_name": ("group", {"num_groups": 16}),
                         "deep_supervision": True,
                         "deep_supr_num": deep_supr_num,
                         "res_block": res_block,
diff --git a/tests/test_dynunet_block.py b/tests/test_dynunet_block.py
--- a/tests/test_dynunet_block.py
+++ b/tests/test_dynunet_block.py
@@ -22,7 +22,7 @@
 for spatial_dims in range(2, 4):
     for kernel_size in [1, 3]:
         for stride in [1, 2]:
-            for norm_name in ["group", "batch", "instance"]:
+            for norm_name in [("GROUP", {"num_groups": 16}), ("batch", {"track_running_stats": False}), "instance"]:
                 for in_size in [15, 16]:
                     padding = get_padding(kernel_size, stride)
                     if not isinstance(padding, int):
@@ -76,7 +76,7 @@ def test_shape(self, input_param, input_shape, expected_shape):
                 self.assertEqual(result.shape, expected_shape)
 
     def test_ill_arg(self):
-        with self.assertRaises(ValueError):
+        with self.assertRaises(KeyError):
             UnetBasicBlock(3, 4, 2, kernel_size=3, stride=1, norm_name="norm")
         with self.assertRaises(AssertionError):
             UnetResBlock(3, 4, 2, kernel_size=1, stride=4, norm_name="batch")
diff --git a/tests/test_segresnet.py b/tests/test_segresnet.py
--- a/tests/test_segresnet.py
+++ b/tests/test_segresnet.py
@@ -25,14 +25,14 @@
 for spatial_dims in range(2, 4):
     for init_filters in [8, 16]:
         for dropout_prob in [None, 0.2]:
-            for norm_name in ["group", "batch", "instance"]:
+            for norm in [("GROUP", {"num_groups": 8}), ("batch", {"track_running_stats": False}), "instance"]:
                 for upsample_mode in UpsampleMode:
                     test_case = [
                         {
                             "spatial_dims": spatial_dims,
                             "init_filters": init_filters,
                             "dropout_prob": dropout_prob,
-                            "norm_name": norm_name,
+                            "norm": norm,
                             "upsample_mode": upsample_mode,
                             "use_conv_final": False,
                         },
diff --git a/tests/test_segresnet_block.py b/tests/test_segresnet_block.py
--- a/tests/test_segresnet_block.py
+++ b/tests/test_segresnet_block.py
@@ -21,14 +21,13 @@
 for spatial_dims in range(2, 4):
     for in_channels in range(1, 4):
         for kernel_size in [1, 3]:
-            for norm_name in ["group", "batch", "instance"]:
+            for norm in ["group", "batch", "instance"]:
                 test_case = [
                     {
                         "spatial_dims": spatial_dims,
                         "in_channels": in_channels,
                         "kernel_size": kernel_size,
-                        "norm_name": norm_name,
-                        "num_groups": in_channels,
+                        "norm": norm,
                     },
                     (2, in_channels, *([16] * spatial_dims)),
                     (2, in_channels, *([16] * spatial_dims)),
@@ -46,11 +45,9 @@ def test_shape(self, input_param, input_shape, expected_shape):
 
     def test_ill_arg(self):
         with self.assertRaises(AssertionError):
-            ResBlock(spatial_dims=3, in_channels=8, kernel_size=2, num_groups=8)
-        with self.assertRaises(ValueError):
-            ResBlock(spatial_dims=3, in_channels=8, norm_name="norm", num_groups=8)
-        with self.assertRaises(AssertionError):
-            ResBlock(spatial_dims=3, in_channels=8, num_groups=3)
+            ResBlock(spatial_dims=3, in_channels=8, norm="group", kernel_size=2)
+        with self.assertRaises(KeyError):
+            ResBlock(spatial_dims=3, in_channels=8, norm="norm")
 
 
 if __name__ == "__main__":