diff --git a/monai/networks/blocks/dynunet_block.py b/monai/networks/blocks/dynunet_block.py index bb654d841c..fc37fc8999 100644 --- a/monai/networks/blocks/dynunet_block.py +++ b/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, diff --git a/monai/networks/blocks/dynunet_block_v1.py b/monai/networks/blocks/dynunet_block_v1.py index d5d9bbf3dc..b5b88dd0df 100644 --- a/monai/networks/blocks/dynunet_block_v1.py +++ b/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, ) diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py index 4af70b22c7..d65cd9f5f4 100644 --- a/monai/networks/nets/dynunet.py +++ b/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) diff --git a/monai/networks/nets/dynunet_v1.py b/monai/networks/nets/dynunet_v1.py index feb05d1762..c6a54807e4 100644 --- a/monai/networks/nets/dynunet_v1.py +++ b/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() diff --git a/tests/test_dynunet.py b/tests/test_dynunet.py index 81ed239461..18fe146a40 100644 --- a/tests/test_dynunet.py +++ b/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,6 +45,7 @@ "norm_name": "batch", "deep_supervision": False, "res_block": res_block, + "dropout": None, }, (1, in_channels, in_size, in_size), expected_shape, @@ -52,11 +53,11 @@ 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,