Source code for monai.networks.blocks.dynunet_block

# Copyright (c) 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 __future__ import annotations

from collections.abc import Sequence

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
from monai.networks.layers.utils import get_act_layer, get_norm_layer


[docs] class UnetResBlock(nn.Module): """ A skip-connection based 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: Sequence[int] | int, stride: Sequence[int] | int, norm_name: tuple | str, act_name: tuple | str = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}), dropout: tuple | str | float | None = None, ): super().__init__() self.conv1 = get_conv_layer( spatial_dims, in_channels, out_channels, kernel_size=kernel_size, stride=stride, dropout=dropout, act=None, norm=None, conv_only=False, ) self.conv2 = get_conv_layer( spatial_dims, out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, act=None, norm=None, conv_only=False, ) 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) self.downsample = in_channels != out_channels stride_np = np.atleast_1d(stride) if not np.all(stride_np == 1): self.downsample = True if self.downsample: self.conv3 = get_conv_layer( spatial_dims, in_channels, out_channels, kernel_size=1, stride=stride, dropout=dropout, act=None, norm=None, conv_only=False, ) self.norm3 = get_norm_layer(name=norm_name, spatial_dims=spatial_dims, channels=out_channels)
[docs] def forward(self, inp): residual = inp out = self.conv1(inp) out = self.norm1(out) out = self.lrelu(out) out = self.conv2(out) out = self.norm2(out) if hasattr(self, "conv3"): residual = self.conv3(residual) if hasattr(self, "norm3"): residual = self.norm3(residual) out += residual out = self.lrelu(out) return out
[docs] class UnetBasicBlock(nn.Module): """ A CNN 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: Sequence[int] | int, stride: Sequence[int] | int, norm_name: tuple | str, act_name: tuple | str = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}), dropout: tuple | str | float | None = None, ): super().__init__() self.conv1 = get_conv_layer( spatial_dims, in_channels, out_channels, kernel_size=kernel_size, stride=stride, dropout=dropout, act=None, norm=None, conv_only=False, ) self.conv2 = get_conv_layer( spatial_dims, out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, act=None, norm=None, conv_only=False, ) 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)
[docs] def forward(self, inp): out = self.conv1(inp) out = self.norm1(out) out = self.lrelu(out) out = self.conv2(out) out = self.norm2(out) out = self.lrelu(out) return out
[docs] class UnetUpBlock(nn.Module): """ An upsampling 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. upsample_kernel_size: convolution kernel size for transposed convolution layers. norm_name: feature normalization type and arguments. act_name: activation layer type and arguments. dropout: dropout probability. trans_bias: transposed convolution bias. """ def __init__( self, spatial_dims: int, in_channels: int, out_channels: int, kernel_size: Sequence[int] | int, stride: Sequence[int] | int, upsample_kernel_size: Sequence[int] | int, norm_name: tuple | str, act_name: tuple | str = ("leakyrelu", {"inplace": True, "negative_slope": 0.01}), dropout: tuple | str | float | None = None, trans_bias: bool = False, ): super().__init__() upsample_stride = upsample_kernel_size self.transp_conv = get_conv_layer( spatial_dims, in_channels, out_channels, kernel_size=upsample_kernel_size, stride=upsample_stride, dropout=dropout, bias=trans_bias, act=None, norm=None, conv_only=False, is_transposed=True, ) self.conv_block = UnetBasicBlock( spatial_dims, out_channels + out_channels, out_channels, kernel_size=kernel_size, stride=1, dropout=dropout, norm_name=norm_name, act_name=act_name, )
[docs] def forward(self, inp, skip): # number of channels for skip should equals to out_channels out = self.transp_conv(inp) out = torch.cat((out, skip), dim=1) out = self.conv_block(out) return out
[docs] class UnetOutBlock(nn.Module): def __init__( self, spatial_dims: int, in_channels: int, out_channels: int, dropout: tuple | str | float | None = None ): super().__init__() self.conv = get_conv_layer( spatial_dims, in_channels, out_channels, kernel_size=1, stride=1, dropout=dropout, bias=True, act=None, norm=None, conv_only=False, )
[docs] def forward(self, inp): return self.conv(inp)
def get_conv_layer( spatial_dims: int, in_channels: int, out_channels: int, kernel_size: Sequence[int] | int = 3, stride: Sequence[int] | int = 1, act: tuple | str | None = Act.PRELU, norm: tuple | str | None = Norm.INSTANCE, dropout: tuple | str | float | None = None, bias: bool = False, conv_only: bool = True, is_transposed: bool = False, ): padding = get_padding(kernel_size, stride) output_padding = None if is_transposed: output_padding = get_output_padding(kernel_size, stride, padding) return Convolution( spatial_dims, in_channels, out_channels, strides=stride, kernel_size=kernel_size, act=act, norm=norm, dropout=dropout, bias=bias, conv_only=conv_only, is_transposed=is_transposed, padding=padding, output_padding=output_padding, ) def get_padding(kernel_size: Sequence[int] | int, stride: Sequence[int] | int) -> tuple[int, ...] | int: kernel_size_np = np.atleast_1d(kernel_size) stride_np = np.atleast_1d(stride) padding_np = (kernel_size_np - stride_np + 1) / 2 if np.min(padding_np) < 0: raise AssertionError("padding value should not be negative, please change the kernel size and/or stride.") padding = tuple(int(p) for p in padding_np) return padding if len(padding) > 1 else padding[0] def get_output_padding( kernel_size: Sequence[int] | int, stride: Sequence[int] | int, padding: Sequence[int] | int ) -> tuple[int, ...] | int: kernel_size_np = np.atleast_1d(kernel_size) stride_np = np.atleast_1d(stride) padding_np = np.atleast_1d(padding) out_padding_np = 2 * padding_np + stride_np - kernel_size_np if np.min(out_padding_np) < 0: raise AssertionError("out_padding value should not be negative, please change the kernel size and/or stride.") out_padding = tuple(int(p) for p in out_padding_np) return out_padding if len(out_padding) > 1 else out_padding[0]