Add dropout arg in DynUNet init

monai/networks/blocks/dynunet_block.py

@@ -33,6 +33,7 @@ class UnetResBlock(nn.Module):
         kernel_size: convolution kernel size.
         stride: convolution stride.
         norm_name: feature normalization type and arguments.
+        dropout: dropout probability
 
     """
 
@@ -44,6 +45,7 @@ def __init__(
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
         norm_name: Union[Tuple, str],
+        dropout: Optional[Union[Tuple, str, float]] = None,
     ):
         super(UnetResBlock, self).__init__()
         self.conv1 = get_conv_layer(
@@ -52,6 +54,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv2 = get_conv_layer(
@@ -60,6 +63,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=1,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv3 = get_conv_layer(
@@ -68,6 +72,7 @@ def __init__(
             out_channels,
             kernel_size=1,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
@@ -107,6 +112,7 @@ class UnetBasicBlock(nn.Module):
         kernel_size: convolution kernel size.
         stride: convolution stride.
         norm_name: feature normalization type and arguments.
+        dropout: dropout probability
 
     """
 
@@ -118,6 +124,7 @@ def __init__(
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
         norm_name: Union[Tuple, str],
+        dropout: Optional[Union[Tuple, str, float]] = None,
     ):
         super(UnetBasicBlock, self).__init__()
         self.conv1 = get_conv_layer(
@@ -126,6 +133,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv2 = get_conv_layer(
@@ -134,6 +142,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=1,
+            dropout=dropout,
             conv_only=True,
         )
         self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
@@ -164,6 +173,7 @@ class UnetUpBlock(nn.Module):
         stride: convolution stride.
         upsample_kernel_size: convolution kernel size for transposed convolution layers.
         norm_name: feature normalization type and arguments.
+        dropout: dropout probability
 
     """
 
@@ -176,6 +186,7 @@ def __init__(
         stride: Union[Sequence[int], int],
         upsample_kernel_size: Union[Sequence[int], int],
         norm_name: Union[Tuple, str],
+        dropout: Optional[Union[Tuple, str, float]] = None,
     ):
         super(UnetUpBlock, self).__init__()
         upsample_stride = upsample_kernel_size
@@ -185,6 +196,7 @@ def __init__(
             out_channels,
             kernel_size=upsample_kernel_size,
             stride=upsample_stride,
+            dropout=dropout,
             conv_only=True,
             is_transposed=True,
         )
@@ -194,6 +206,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=1,
+            dropout=dropout,
             norm_name=norm_name,
         )
 
@@ -206,10 +219,12 @@ def forward(self, inp, skip):
 
 
 class UnetOutBlock(nn.Module):
-    def __init__(self, spatial_dims: int, in_channels: int, out_channels: int):
+    def __init__(
+        self, spatial_dims: int, in_channels: int, out_channels: int, dropout: Optional[Union[Tuple, str, float]] = None
+    ):
         super(UnetOutBlock, self).__init__()
         self.conv = get_conv_layer(
-            spatial_dims, in_channels, out_channels, kernel_size=1, stride=1, bias=True, conv_only=True
+            spatial_dims, in_channels, out_channels, kernel_size=1, stride=1, dropout=dropout, bias=True, conv_only=True
         )
 
     def forward(self, inp):
@@ -224,6 +239,7 @@ def get_conv_layer(
     stride: Union[Sequence[int], int] = 1,
     act: Optional[Union[Tuple, str]] = Act.PRELU,
     norm: Union[Tuple, str] = Norm.INSTANCE,
+    dropout: Optional[Union[Tuple, str, float]] = None,
     bias: bool = False,
     conv_only: bool = True,
     is_transposed: bool = False,
@@ -240,6 +256,7 @@ def get_conv_layer(
         kernel_size=kernel_size,
         act=act,
         norm=norm,
+        dropout=dropout,
         bias=bias,
         conv_only=conv_only,
         is_transposed=is_transposed,

monai/networks/blocks/dynunet_block_v1.py

@@ -32,6 +32,7 @@ def __init__(
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
         norm_name: str,
+        dropout: float = 0.0,
     ):
         nn.Module.__init__(self)
         self.conv1 = get_conv_layer(
@@ -40,6 +41,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv2 = get_conv_layer(
@@ -48,6 +50,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=1,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv3 = get_conv_layer(
@@ -56,6 +59,7 @@ def __init__(
             out_channels,
             kernel_size=1,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
@@ -81,6 +85,7 @@ def __init__(
         kernel_size: Union[Sequence[int], int],
         stride: Union[Sequence[int], int],
         norm_name: str,
+        dropout: float = 0.0,
     ):
         nn.Module.__init__(self)
         self.conv1 = get_conv_layer(
@@ -89,6 +94,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=stride,
+            dropout=dropout,
             conv_only=True,
         )
         self.conv2 = get_conv_layer(
@@ -97,6 +103,7 @@ def __init__(
             out_channels,
             kernel_size=kernel_size,
             stride=1,
+            dropout=dropout,
             conv_only=True,
         )
         self.lrelu = get_act_layer(("leakyrelu", {"inplace": True, "negative_slope": 0.01}))
@@ -118,6 +125,7 @@ def __init__(
         stride: Union[Sequence[int], int],
         upsample_kernel_size: Union[Sequence[int], int],
         norm_name: str,
+        dropout: float = 0.0,
     ):
         nn.Module.__init__(self)
         upsample_stride = upsample_kernel_size
@@ -127,6 +135,7 @@ def __init__(
             out_channels,
             kernel_size=upsample_kernel_size,
             stride=upsample_stride,
+            dropout=dropout,
             conv_only=True,
             is_transposed=True,
         )
@@ -137,6 +146,7 @@ def __init__(
             kernel_size=kernel_size,
             stride=1,
             norm_name=norm_name,
+            dropout=dropout,
         )
 
 

monai/networks/nets/dynunet.py

@@ -86,6 +86,7 @@ class DynUNet(nn.Module):
         strides: convolution strides for each blocks.
         upsample_kernel_size: convolution kernel size for transposed convolution layers. The values should
             equal to strides[1:].
+        dropout: dropout ratio. Defaults to no dropout.
         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
@@ -115,6 +116,7 @@ def __init__(
         kernel_size: Sequence[Union[Sequence[int], int]],
         strides: Sequence[Union[Sequence[int], int]],
         upsample_kernel_size: Sequence[Union[Sequence[int], int]],
+        dropout: Optional[Union[Tuple, str, float]] = None,
         norm_name: Union[Tuple, str] = ("INSTANCE", {"affine": True}),
         deep_supervision: bool = False,
         deep_supr_num: int = 1,
@@ -128,6 +130,7 @@ def __init__(
         self.strides = strides
         self.upsample_kernel_size = upsample_kernel_size
         self.norm_name = norm_name
+        self.dropout = dropout
         self.conv_block = UnetResBlock if res_block else UnetBasicBlock
         self.filters = [min(2 ** (5 + i), 320 if spatial_dims == 3 else 512) for i in range(len(strides))]
         self.input_block = self.get_input_block()
@@ -184,7 +187,7 @@ def create_skips(index, downsamples, upsamples, superheads, bottleneck):
     def check_kernel_stride(self):
         kernels, strides = self.kernel_size, self.strides
         error_msg = "length of kernel_size and strides should be the same, and no less than 3."
-        if not (len(kernels) == len(strides) and len(kernels) >= 3):
+        if len(kernels) != len(strides) or len(kernels) < 3:
             raise AssertionError(error_msg)
 
         for idx, k_i in enumerate(kernels):
@@ -225,6 +228,7 @@ def get_input_block(self):
             self.kernel_size[0],
             self.strides[0],
             self.norm_name,
+            dropout=self.dropout,
         )
 
     def get_bottleneck(self):
@@ -235,14 +239,11 @@ def get_bottleneck(self):
             self.kernel_size[-1],
             self.strides[-1],
             self.norm_name,
+            dropout=self.dropout,
         )
 
     def get_output_block(self, idx: int):
-        return UnetOutBlock(
-            self.spatial_dims,
-            self.filters[idx],
-            self.out_channels,
-        )
+        return UnetOutBlock(self.spatial_dims, self.filters[idx], self.out_channels, dropout=self.dropout)
 
     def get_downsamples(self):
         inp, out = self.filters[:-2], self.filters[1:-1]
@@ -276,6 +277,7 @@ def get_module_list(
                     "kernel_size": kernel,
                     "stride": stride,
                     "norm_name": self.norm_name,
+                    "dropout": self.dropout,
                     "upsample_kernel_size": up_kernel,
                 }
                 layer = conv_block(**params)
@@ -289,6 +291,7 @@ def get_module_list(
                     "kernel_size": kernel,
                     "stride": stride,
                     "norm_name": self.norm_name,
+                    "dropout": self.dropout,
                 }
                 layer = conv_block(**params)
                 layers.append(layer)

monai/networks/nets/dynunet_v1.py

@@ -38,6 +38,7 @@ class DynUNetV1(DynUNet):
         kernel_size: convolution kernel size.
         strides: convolution strides for each blocks.
         upsample_kernel_size: convolution kernel size for transposed convolution layers.
+        dropout: dropout ratio. Defaults to no dropout.
         norm_name: [``"batch"``, ``"instance"``, ``"group"``]. Defaults to "instance".
         deep_supervision: whether to add deep supervision head before output. Defaults to ``False``.
         deep_supr_num: number of feature maps that will output during deep supervision head. Defaults to 1.
@@ -57,6 +58,7 @@ def __init__(
         kernel_size: Sequence[Union[Sequence[int], int]],
         strides: Sequence[Union[Sequence[int], int]],
         upsample_kernel_size: Sequence[Union[Sequence[int], int]],
+        dropout: float = 0.0,
         norm_name: str = "instance",
         deep_supervision: bool = False,
         deep_supr_num: int = 1,
@@ -70,6 +72,7 @@ def __init__(
         self.strides = strides
         self.upsample_kernel_size = upsample_kernel_size
         self.norm_name = norm_name
+        self.dropout = dropout
         self.conv_block = _UnetResBlockV1 if res_block else _UnetBasicBlockV1  # type: ignore
         self.filters = [min(2 ** (5 + i), 320 if spatial_dims == 3 else 512) for i in range(len(strides))]
         self.input_block = self.get_input_block()

tests/test_dynunet.py

@@ -26,14 +26,14 @@
 expected_shape: Sequence[Any]
 
 TEST_CASE_DYNUNET_2D = []
+out_channels = 2
+in_size = 64
+spatial_dims = 2
 for kernel_size in [(3, 3, 3, 1), ((3, 1), 1, (3, 3), (1, 1))]:
     for strides in [(1, 1, 1, 1), (2, 2, 2, 1)]:
+        expected_shape = (1, out_channels, *[in_size // strides[0]] * spatial_dims)
         for in_channels in [2, 3]:
             for res_block in [True, False]:
-                out_channels = 2
-                in_size = 64
-                spatial_dims = 2
-                expected_shape = (1, out_channels, *[in_size // strides[0]] * spatial_dims)
                 test_case = [
                     {
                         "spatial_dims": spatial_dims,
@@ -45,18 +45,19 @@
                         "norm_name": "batch",
                         "deep_supervision": False,
                         "res_block": res_block,
+                        "dropout": None,
                     },
                     (1, in_channels, in_size, in_size),
                     expected_shape,
                 ]
                 TEST_CASE_DYNUNET_2D.append(test_case)
 
 TEST_CASE_DYNUNET_3D = []  # in 3d cases, also test anisotropic kernel/strides
+in_channels = 1
+in_size = 64
 for out_channels in [2, 3]:
+    expected_shape = (1, out_channels, 64, 32, 64)
     for res_block in [True, False]:
-        in_channels = 1
-        in_size = 64
-        expected_shape = (1, out_channels, 64, 32, 64)
         test_case = [
             {
                 "spatial_dims": 3,
@@ -68,6 +69,7 @@
                 "norm_name": ("INSTANCE", {"affine": True}),
                 "deep_supervision": False,
                 "res_block": res_block,
+                "dropout": ("alphadropout", {"p": 0.25}),
             },
             (1, in_channels, in_size, in_size, in_size),
             expected_shape,