Describe the bug
Even though the DynUNet does have an option for Dropout, the Dropout layers are never created.
To Reproduce
model = DynUNet(
spatial_dims=2,
in_channels=1,
out_channels=1,
kernel_size=(3, 3,3),
strides=(1,2,1),
upsample_kernel_size=(1,2,1),
dropout=0.2
).to(device)
print(model)Expected behavior
The Dropout layers should be added after each norm layer. See UnetBasicBlock
E.g. something like this:
class UnetBasicBlock(nn.Module):
"""
A CNN module module that can be used for DynUNet, based on:
`Automated Design of Deep Learning Methods for Biomedical Image Segmentation <https://arxiv.org/abs/1904.08128>`_.
`nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation <https://arxiv.org/abs/1809.10486>`_.
Args:
spatial_dims: number of spatial dimensions.
in_channels: number of input channels.
out_channels: number of output channels.
kernel_size: convolution kernel size.
stride: convolution stride.
norm_name: feature normalization type and arguments.
act_name: activation layer type and arguments.
dropout: dropout probability.
"""
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: Union[Sequence[int], int],
stride: Union[Sequence[int], int],
norm_name: Union[Tuple, str],
act_name: Union[Tuple, str] = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}),
dropout: Optional[Union[Tuple, str, float]] = None,
):
super().__init__()
self.conv1 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
dropout=dropout,
conv_only=True,
)
self.conv2 = get_conv_layer(
spatial_dims, out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, conv_only=True
)
self.lrelu = get_act_layer(name=act_name)
self.norm1 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)
self.norm2 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)
# CREATE DROPOUT LAYERS
self.dropout1 = nn.Dropout2d(dropout if dropout is not None else 0)
self.dropout2 = nn.Dropout2d(dropout if dropout is not None else 0)
def forward(self, inp):
out = self.conv1(inp)
out = self.norm1(out)
# USE HERE
out = self.dropout1(out)
out = self.lrelu(out)
out = self.conv2(out)
out = self.norm2(out)
# AND HERE
out = self.dropout2(out)
out = self.lrelu(out)
return out
Describe the bug
Even though the DynUNet does have an option for Dropout, the Dropout layers are never created.
To Reproduce
Expected behavior
The Dropout layers should be added after each norm layer. See UnetBasicBlock
E.g. something like this: