diff --git a/docs/source/networks.rst b/docs/source/networks.rst index 811084753f..cc0dfb0be5 100644 --- a/docs/source/networks.rst +++ b/docs/source/networks.rst @@ -528,6 +528,13 @@ Nets .. autoclass:: BasicUnet .. autoclass:: Basicunet +`BasicUNetPlusPlus` +~~~~~~~~~~~~~~~~~~~ +.. autoclass:: BasicUNetPlusPlus + :members: +.. autoclass:: BasicUnetPlusPlus +.. autoclass:: BasicunetPlusPlus + `FlexibleUNet` ~~~~~~~~~~~~~~ .. autoclass:: FlexibleUNet diff --git a/monai/networks/nets/__init__.py b/monai/networks/nets/__init__.py index cd9329f61b..b2a21cd88b 100644 --- a/monai/networks/nets/__init__.py +++ b/monai/networks/nets/__init__.py @@ -13,6 +13,7 @@ from .attentionunet import AttentionUnet from .autoencoder import AutoEncoder from .basic_unet import BasicUNet, BasicUnet, Basicunet, basicunet +from .basic_unetplusplus import BasicUNetPlusPlus, BasicUnetPlusPlus, BasicunetPlusPlus, basicunetplusplus from .classifier import Classifier, Critic, Discriminator from .densenet import ( DenseNet, diff --git a/monai/networks/nets/basic_unet.py b/monai/networks/nets/basic_unet.py index 8c87209a86..6fe77038fe 100644 --- a/monai/networks/nets/basic_unet.py +++ b/monai/networks/nets/basic_unet.py @@ -260,13 +260,13 @@ def forward(self, x: torch.Tensor): """ Args: x: input should have spatially N dimensions - ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. + ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N-1])``, N is defined by `spatial_dims`. It is recommended to have ``dim_n % 16 == 0`` to ensure all maxpooling inputs have even edge lengths. Returns: A torch Tensor of "raw" predictions in shape - ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N])``. + ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N-1])``. """ x0 = self.conv_0(x) diff --git a/monai/networks/nets/basic_unetplusplus.py b/monai/networks/nets/basic_unetplusplus.py new file mode 100644 index 0000000000..4f7d319aaa --- /dev/null +++ b/monai/networks/nets/basic_unetplusplus.py @@ -0,0 +1,171 @@ +# 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 typing import Sequence, Union + +import torch +import torch.nn as nn + +from monai.networks.layers.factories import Conv +from monai.networks.nets.basic_unet import Down, TwoConv, UpCat +from monai.utils import ensure_tuple_rep + +__all__ = ["BasicUnetPlusPlus", "BasicunetPlusPlus", "basicunetplusplus", "BasicUNetPlusPlus"] + + +class BasicUNetPlusPlus(nn.Module): + def __init__( + self, + spatial_dims: int = 3, + in_channels: int = 1, + out_channels: int = 2, + features: Sequence[int] = (32, 32, 64, 128, 256, 32), + deep_supervision: bool = False, + act: Union[str, tuple] = ("LeakyReLU", {"negative_slope": 0.1, "inplace": True}), + norm: Union[str, tuple] = ("instance", {"affine": True}), + bias: bool = True, + dropout: Union[float, tuple] = 0.0, + upsample: str = "deconv", + ): + """ + A UNet++ implementation with 1D/2D/3D supports. + + Based on: + + Zhou et al. "UNet++: A Nested U-Net Architecture for Medical Image + Segmentation". 4th Deep Learning in Medical Image Analysis (DLMIA) + Workshop, DOI: https://doi.org/10.48550/arXiv.1807.10165 + + + Args: + spatial_dims: number of spatial dimensions. Defaults to 3 for spatial 3D inputs. + in_channels: number of input channels. Defaults to 1. + out_channels: number of output channels. Defaults to 2. + features: six integers as numbers of features. + Defaults to ``(32, 32, 64, 128, 256, 32)``, + + - the first five values correspond to the five-level encoder feature sizes. + - the last value corresponds to the feature size after the last upsampling. + + deep_supervision: whether to prune the network at inference time. Defaults to False. If true, returns a list, + whose elements correspond to outputs at different nodes. + act: activation type and arguments. Defaults to LeakyReLU. + norm: feature normalization type and arguments. Defaults to instance norm. + bias: whether to have a bias term in convolution blocks. Defaults to True. + According to `Performance Tuning Guide `_, + if a conv layer is directly followed by a batch norm layer, bias should be False. + dropout: dropout ratio. Defaults to no dropout. + upsample: upsampling mode, available options are + ``"deconv"``, ``"pixelshuffle"``, ``"nontrainable"``. + + Examples:: + + # for spatial 2D + >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128)) + + # for spatial 2D, with deep supervision enabled + >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), deep_supervision=True) + + # for spatial 2D, with group norm + >>> net = BasicUNetPlusPlus(spatial_dims=2, features=(64, 128, 256, 512, 1024, 128), norm=("group", {"num_groups": 4})) + + # for spatial 3D + >>> net = BasicUNetPlusPlus(spatial_dims=3, features=(32, 32, 64, 128, 256, 32)) + + See Also + - :py:class:`monai.networks.nets.BasicUNet` + - :py:class:`monai.networks.nets.DynUNet` + - :py:class:`monai.networks.nets.UNet` + + """ + super().__init__() + + self.deep_supervision = deep_supervision + + fea = ensure_tuple_rep(features, 6) + print(f"BasicUNetPlusPlus features: {fea}.") + + self.conv_0_0 = TwoConv(spatial_dims, in_channels, fea[0], act, norm, bias, dropout) + self.conv_1_0 = Down(spatial_dims, fea[0], fea[1], act, norm, bias, dropout) + self.conv_2_0 = Down(spatial_dims, fea[1], fea[2], act, norm, bias, dropout) + self.conv_3_0 = Down(spatial_dims, fea[2], fea[3], act, norm, bias, dropout) + self.conv_4_0 = Down(spatial_dims, fea[3], fea[4], act, norm, bias, dropout) + + self.upcat_0_1 = UpCat(spatial_dims, fea[1], fea[0], fea[0], act, norm, bias, dropout, upsample, halves=False) + self.upcat_1_1 = UpCat(spatial_dims, fea[2], fea[1], fea[1], act, norm, bias, dropout, upsample) + self.upcat_2_1 = UpCat(spatial_dims, fea[3], fea[2], fea[2], act, norm, bias, dropout, upsample) + self.upcat_3_1 = UpCat(spatial_dims, fea[4], fea[3], fea[3], act, norm, bias, dropout, upsample) + + self.upcat_0_2 = UpCat( + spatial_dims, fea[1], fea[0] * 2, fea[0], act, norm, bias, dropout, upsample, halves=False + ) + self.upcat_1_2 = UpCat(spatial_dims, fea[2], fea[1] * 2, fea[1], act, norm, bias, dropout, upsample) + self.upcat_2_2 = UpCat(spatial_dims, fea[3], fea[2] * 2, fea[2], act, norm, bias, dropout, upsample) + + self.upcat_0_3 = UpCat( + spatial_dims, fea[1], fea[0] * 3, fea[0], act, norm, bias, dropout, upsample, halves=False + ) + self.upcat_1_3 = UpCat(spatial_dims, fea[2], fea[1] * 3, fea[1], act, norm, bias, dropout, upsample) + + self.upcat_0_4 = UpCat( + spatial_dims, fea[1], fea[0] * 4, fea[5], act, norm, bias, dropout, upsample, halves=False + ) + + self.final_conv_0_1 = Conv["conv", spatial_dims](fea[0], out_channels, kernel_size=1) + self.final_conv_0_2 = Conv["conv", spatial_dims](fea[0], out_channels, kernel_size=1) + self.final_conv_0_3 = Conv["conv", spatial_dims](fea[0], out_channels, kernel_size=1) + self.final_conv_0_4 = Conv["conv", spatial_dims](fea[5], out_channels, kernel_size=1) + + def forward(self, x: torch.Tensor): + """ + Args: + x: input should have spatially N dimensions + ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N-1])``, N is defined by `dimensions`. + It is recommended to have ``dim_n % 16 == 0`` to ensure all maxpooling inputs have + even edge lengths. + + Returns: + A torch Tensor of "raw" predictions in shape + ``(Batch, out_channels, dim_0[, dim_1, ..., dim_N-1])``. + """ + x_0_0 = self.conv_0_0(x) + x_1_0 = self.conv_1_0(x_0_0) + x_0_1 = self.upcat_0_1(x_1_0, x_0_0) + + x_2_0 = self.conv_2_0(x_1_0) + x_1_1 = self.upcat_1_1(x_2_0, x_1_0) + x_0_2 = self.upcat_0_2(x_1_1, torch.cat([x_0_0, x_0_1], dim=1)) + + x_3_0 = self.conv_3_0(x_2_0) + x_2_1 = self.upcat_2_1(x_3_0, x_2_0) + x_1_2 = self.upcat_1_2(x_2_1, torch.cat([x_1_0, x_1_1], dim=1)) + x_0_3 = self.upcat_0_3(x_1_2, torch.cat([x_0_0, x_0_1, x_0_2], dim=1)) + + x_4_0 = self.conv_4_0(x_3_0) + x_3_1 = self.upcat_3_1(x_4_0, x_3_0) + x_2_2 = self.upcat_2_2(x_3_1, torch.cat([x_2_0, x_2_1], dim=1)) + x_1_3 = self.upcat_1_3(x_2_2, torch.cat([x_1_0, x_1_1, x_1_2], dim=1)) + x_0_4 = self.upcat_0_4(x_1_3, torch.cat([x_0_0, x_0_1, x_0_2, x_0_3], dim=1)) + + output_0_1 = self.final_conv_0_1(x_0_1) + output_0_2 = self.final_conv_0_2(x_0_2) + output_0_3 = self.final_conv_0_3(x_0_3) + output_0_4 = self.final_conv_0_4(x_0_4) + + if self.deep_supervision: + output = [output_0_1, output_0_2, output_0_3, output_0_4] + else: + output = [output_0_4] + + return output + + +BasicUnetPlusPlus = BasicunetPlusPlus = basicunetplusplus = BasicUNetPlusPlus diff --git a/tests/test_basic_unetplusplus.py b/tests/test_basic_unetplusplus.py new file mode 100644 index 0000000000..3bca65676a --- /dev/null +++ b/tests/test_basic_unetplusplus.py @@ -0,0 +1,107 @@ +# 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. + +import unittest + +import torch +from parameterized import parameterized + +from monai.networks import eval_mode +from monai.networks.nets import BasicUNetPlusPlus +from tests.utils import test_script_save + +CASES_1D = [] +for mode in ["pixelshuffle", "nontrainable", "deconv", None]: + kwargs = {"spatial_dims": 1, "in_channels": 5, "out_channels": 8} + if mode is not None: + kwargs["upsample"] = mode # type: ignore + CASES_1D.append([kwargs, (10, 5, 33), (10, 8, 33)]) + +CASES_2D = [] +for mode in ["pixelshuffle", "nontrainable", "deconv"]: + for d1 in range(33, 64, 14): + for d2 in range(63, 33, -21): + in_channels, out_channels = 2, 3 + CASES_2D.append( + [ + { + "spatial_dims": 2, + "in_channels": in_channels, + "out_channels": out_channels, + "features": (12, 12, 13, 14, 15, 16), + "upsample": mode, + }, + (2, in_channels, d1, d2), + (2, out_channels, d1, d2), + ] + ) +CASES_3D = [ + [ # single channel 3D, batch 2 + { + "spatial_dims": 3, + "in_channels": 1, + "out_channels": 2, + "features": (16, 20, 21, 22, 23, 11), + "upsample": "pixelshuffle", + }, + (2, 1, 33, 34, 35), + (2, 2, 33, 34, 35), + ], + [ # 2-channel 3D, batch 3 + { + "spatial_dims": 3, + "in_channels": 2, + "out_channels": 7, + "features": (14, 15, 16, 17, 18, 11), + "upsample": "deconv", + }, + (3, 2, 33, 37, 34), + (3, 7, 33, 37, 34), + ], + [ # 4-channel 3D, batch 5 + { + "spatial_dims": 3, + "in_channels": 4, + "out_channels": 2, + "features": (14, 15, 16, 17, 18, 10), + "upsample": "nontrainable", + }, + (5, 4, 34, 35, 37), + (5, 2, 34, 35, 37), + ], +] + + +class TestBasicUNETPlusPlus(unittest.TestCase): + @parameterized.expand(CASES_1D + CASES_2D + CASES_3D) + def test_shape(self, input_param, input_shape, expected_shape): + device = "cuda" if torch.cuda.is_available() else "cpu" + print(input_param) + net = BasicUNetPlusPlus(**input_param).to(device) + with eval_mode(net): + result = net(torch.randn(input_shape).to(device)) + self.assertEqual(result[0].shape, expected_shape) + + def test_deep_supervision_shape(self): + net = BasicUNetPlusPlus(spatial_dims=2, deep_supervision=True, in_channels=3, out_channels=3) + test_data = torch.randn(16, 3, 32, 32) + with eval_mode(net): + result = net(test_data) + self.assertEqual(result[0].shape, test_data.shape) + + def test_script(self): + net = BasicUNetPlusPlus(spatial_dims=2, deep_supervision=True, in_channels=1, out_channels=3) + test_data = torch.randn(16, 1, 32, 32) + test_script_save(net, test_data) + + +if __name__ == "__main__": + unittest.main()