previous dynunet (#2534)

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

monai/networks/blocks/dynunet_block_v1.py

@@ -0,0 +1,152 @@
+# 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 Sequence, Union
+
+import numpy as np
+import torch.nn as nn
+
+from monai.networks.blocks.dynunet_block import UnetBasicBlock, UnetResBlock, UnetUpBlock, get_conv_layer
+from monai.networks.layers.factories import Norm
+from monai.networks.layers.utils import get_act_layer
+
+
+class _UnetResBlockV1(UnetResBlock):
+    """
+    UnetResBlock for backward compatibility purpose.
+    """
+
+    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: str,
+    ):
+        nn.Module.__init__(self)
+        self.conv1 = get_conv_layer(
+            spatial_dims,
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            conv_only=True,
+        )
+        self.conv2 = get_conv_layer(
+            spatial_dims,
+            out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            conv_only=True,
+        )
+        self.conv3 = get_conv_layer(
+            spatial_dims,
+            in_channels,
+            out_channels,
+            kernel_size=1,
+            stride=stride,
+            conv_only=True,
+        )
+        self.lrelu = get_act_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.downsample = in_channels != out_channels
+        stride_np = np.atleast_1d(stride)
+        if not np.all(stride_np == 1):
+            self.downsample = True
+
+
+class _UnetBasicBlockV1(UnetBasicBlock):
+    """
+    UnetBasicBlock for backward compatibility purpose.
+    """
+
+    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: str,
+    ):
+        nn.Module.__init__(self)
+        self.conv1 = get_conv_layer(
+            spatial_dims,
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            conv_only=True,
+        )
+        self.conv2 = get_conv_layer(
+            spatial_dims,
+            out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            conv_only=True,
+        )
+        self.lrelu = get_act_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)
+
+
+class _UnetUpBlockV1(UnetUpBlock):
+    """
+    UnetUpBlock for backward compatibility purpose.
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[Sequence[int], int],
+        stride: Union[Sequence[int], int],
+        upsample_kernel_size: Union[Sequence[int], int],
+        norm_name: str,
+    ):
+        nn.Module.__init__(self)
+        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,
+            conv_only=True,
+            is_transposed=True,
+        )
+        self.conv_block = _UnetBasicBlockV1(
+            spatial_dims,
+            out_channels + out_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            norm_name=norm_name,
+        )
+
+
+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, spatial_dims](num_groups=num_groups, num_channels=out_channels, affine=True)
+    else:
+        norm = Norm[norm_name, spatial_dims](out_channels, affine=True)
+    return norm

monai/networks/nets/dynunet_v1.py

@@ -0,0 +1,140 @@
+# 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 List, Sequence, Union
+
+import torch
+import torch.nn as nn
+
+from monai.networks.blocks.dynunet_block_v1 import _UnetBasicBlockV1, _UnetResBlockV1, _UnetUpBlockV1
+from monai.networks.nets.dynunet import DynUNet, DynUNetSkipLayer
+from monai.utils import deprecated
+
+__all__ = ["DynUNetV1", "DynUnetV1", "DynunetV1"]
+
+
+@deprecated(
+    since="0.6.0",
+    removed="0.7.0",
+    msg_suffix="This module is for backward compatibility purpose only. Please use `DynUNet` instead.",
+)
+class DynUNetV1(DynUNet):
+    """
+    This a deprecated reimplementation of a dynamic UNet (DynUNet), please use `monai.networks.nets.DynUNet` instead.
+
+    Args:
+        spatial_dims: number of spatial dimensions.
+        in_channels: number of input channels.
+        out_channels: number of output channels.
+        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"``]. 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.
+        res_block: whether to use residual connection based convolution blocks during the network.
+            Defaults to ``False``.
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        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",
+        deep_supervision: bool = False,
+        deep_supr_num: int = 1,
+        res_block: bool = False,
+    ):
+        nn.Module.__init__(self)
+        self.spatial_dims = spatial_dims
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.strides = strides
+        self.upsample_kernel_size = upsample_kernel_size
+        self.norm_name = norm_name
+        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()
+        self.downsamples = self.get_downsamples()
+        self.bottleneck = self.get_bottleneck()
+        self.upsamples = self.get_upsamples()
+        self.output_block = self.get_output_block(0)
+        self.deep_supervision = deep_supervision
+        self.deep_supervision_heads = self.get_deep_supervision_heads()
+        self.deep_supr_num = deep_supr_num
+        self.apply(self.initialize_weights)
+        self.check_kernel_stride()
+        self.check_deep_supr_num()
+
+        # initialize the typed list of supervision head outputs so that Torchscript can recognize what's going on
+        self.heads: List[torch.Tensor] = [torch.rand(1)] * (len(self.deep_supervision_heads) + 1)
+
+        def create_skips(index, downsamples, upsamples, superheads, bottleneck):
+            """
+            Construct the UNet topology as a sequence of skip layers terminating with the bottleneck layer. This is
+            done recursively from the top down since a recursive nn.Module subclass is being used to be compatible
+            with Torchscript. Initially the length of `downsamples` will be one more than that of `superheads`
+            since the `input_block` is passed to this function as the first item in `downsamples`, however this
+            shouldn't be associated with a supervision head.
+            """
+
+            if len(downsamples) != len(upsamples):
+                raise AssertionError(f"{len(downsamples)} != {len(upsamples)}")
+            if (len(downsamples) - len(superheads)) not in (1, 0):
+                raise AssertionError(f"{len(downsamples)}-(0,1) != {len(superheads)}")
+
+            if len(downsamples) == 0:  # bottom of the network, pass the bottleneck block
+                return bottleneck
+            if index == 0:  # don't associate a supervision head with self.input_block
+                current_head, rest_heads = nn.Identity(), superheads
+            elif not self.deep_supervision:  # bypass supervision heads by passing nn.Identity in place of a real one
+                current_head, rest_heads = nn.Identity(), superheads[1:]
+            else:
+                current_head, rest_heads = superheads[0], superheads[1:]
+
+            # create the next layer down, this will stop at the bottleneck layer
+            next_layer = create_skips(1 + index, downsamples[1:], upsamples[1:], rest_heads, bottleneck)
+
+            return DynUNetSkipLayer(index, self.heads, downsamples[0], upsamples[0], current_head, next_layer)
+
+        self.skip_layers = create_skips(
+            0,
+            [self.input_block] + list(self.downsamples),
+            self.upsamples[::-1],
+            self.deep_supervision_heads,
+            self.bottleneck,
+        )
+
+    def get_upsamples(self):
+        inp, out = self.filters[1:][::-1], self.filters[:-1][::-1]
+        strides, kernel_size = self.strides[1:][::-1], self.kernel_size[1:][::-1]
+        upsample_kernel_size = self.upsample_kernel_size[::-1]
+        return self.get_module_list(inp, out, kernel_size, strides, _UnetUpBlockV1, upsample_kernel_size)
+
+    @staticmethod
+    def initialize_weights(module):
+        name = module.__class__.__name__.lower()
+        if "conv3d" in name or "conv2d" in name:
+            nn.init.kaiming_normal_(module.weight, a=0.01)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif "norm" in name:
+            nn.init.normal_(module.weight, 1.0, 0.02)
+            nn.init.zeros_(module.bias)
+
+
+DynUnetV1 = DynunetV1 = DynUNetV1