diff --git a/monai/networks/nets/resnet.py b/monai/networks/nets/resnet.py index c8be9f0e89..bf5486f06e 100644 --- a/monai/networks/nets/resnet.py +++ b/monai/networks/nets/resnet.py @@ -150,6 +150,9 @@ class ResNet(nn.Module): Args: block: which ResNet block to use, either Basic or Bottleneck. + ResNet block class or str. + for Basic: ResNetBlock or 'basic' + for Bottleneck: ResNetBottleneck or 'bottleneck' layers: how many layers to use. block_inplanes: determine the size of planes at each step. Also tunable with widen_factor. spatial_dims: number of spatial dimensions of the input image. @@ -172,7 +175,7 @@ class ResNet(nn.Module): @deprecated_arg("n_classes", since="0.6") def __init__( self, - block: Type[Union[ResNetBlock, ResNetBottleneck]], + block: Union[Type[Union[ResNetBlock, ResNetBottleneck]], str], layers: List[int], block_inplanes: List[int], spatial_dims: int = 3, @@ -192,6 +195,14 @@ def __init__( if n_classes is not None and num_classes == 400: num_classes = n_classes + if isinstance(block, str): + if block == "basic": + block = ResNetBlock + elif block == "bottleneck": + block = ResNetBottleneck + else: + raise ValueError("Unknown block '%s', use basic or bottleneck" % block) + conv_type: Type[Union[nn.Conv1d, nn.Conv2d, nn.Conv3d]] = Conv[Conv.CONV, spatial_dims] norm_type: Type[Union[nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d]] = Norm[Norm.BATCH, spatial_dims] pool_type: Type[Union[nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d]] = Pool[Pool.MAX, spatial_dims] diff --git a/monai/networks/nets/senet.py b/monai/networks/nets/senet.py index a85d32ba5a..8933cbe7e9 100644 --- a/monai/networks/nets/senet.py +++ b/monai/networks/nets/senet.py @@ -54,10 +54,10 @@ class SENet(nn.Module): Args: spatial_dims: spatial dimension of the input data. in_channels: channel number of the input data. - block: SEBlock class. - for SENet154: SEBottleneck - for SE-ResNet models: SEResNetBottleneck - for SE-ResNeXt models: SEResNeXtBottleneck + block: SEBlock class or str. + for SENet154: SEBottleneck or 'se_bottleneck' + for SE-ResNet models: SEResNetBottleneck or 'se_resnet_bottleneck' + for SE-ResNeXt models: SEResNeXtBottleneck or 'se_resnetxt_bottleneck' layers: number of residual blocks for 4 layers of the network (layer1...layer4). groups: number of groups for the 3x3 convolution in each bottleneck block. for SENet154: 64 @@ -95,7 +95,7 @@ def __init__( self, spatial_dims: int, in_channels: int, - block: Type[Union[SEBottleneck, SEResNetBottleneck, SEResNeXtBottleneck]], + block: Union[Type[Union[SEBottleneck, SEResNetBottleneck, SEResNeXtBottleneck]], str], layers: Sequence[int], groups: int, reduction: int, @@ -109,6 +109,18 @@ def __init__( super().__init__() + if isinstance(block, str): + if block == "se_bottleneck": + block = SEBottleneck + elif block == "se_resnet_bottleneck": + block = SEResNetBottleneck + elif block == "se_resnetxt_bottleneck": + block = SEResNeXtBottleneck + else: + raise ValueError( + "Unknown block '%s', use se_bottleneck, se_resnet_bottleneck or se_resnetxt_bottleneck" % block + ) + relu_type: Type[nn.ReLU] = Act[Act.RELU] conv_type: Type[Union[nn.Conv1d, nn.Conv2d, nn.Conv3d]] = Conv[Conv.CONV, spatial_dims] pool_type: Type[Union[nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d]] = Pool[Pool.MAX, spatial_dims] diff --git a/tests/test_resnet.py b/tests/test_resnet.py index 688f7827b1..88499f78d0 100644 --- a/tests/test_resnet.py +++ b/tests/test_resnet.py @@ -16,7 +16,8 @@ from parameterized import parameterized from monai.networks import eval_mode -from monai.networks.nets import resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200 +from monai.networks.nets import ResNet, resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200 +from monai.networks.nets.resnet import ResNetBlock from monai.utils import optional_import from tests.utils import test_script_save @@ -95,10 +96,57 @@ ((2, 512), (2, 2048)), ] +TEST_CASE_5 = [ # 1D, batch 1, 2 input channels + { + "block": "basic", + "layers": [1, 1, 1, 1], + "block_inplanes": [64, 128, 256, 512], + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), + (1, 3), +] + +TEST_CASE_5_A = [ # 1D, batch 1, 2 input channels + { + "block": ResNetBlock, + "layers": [1, 1, 1, 1], + "block_inplanes": [64, 128, 256, 512], + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), + (1, 3), +] + +TEST_CASE_6 = [ # 1D, batch 1, 2 input channels + { + "block": "bottleneck", + "layers": [3, 4, 6, 3], + "block_inplanes": [64, 128, 256, 512], + "spatial_dims": 1, + "n_input_channels": 2, + "num_classes": 3, + "conv1_t_size": [3], + "conv1_t_stride": 1, + }, + (1, 2, 32), + (1, 3), +] + TEST_CASES = [] for case in [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_2_A, TEST_CASE_3_A]: for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200]: TEST_CASES.append([model, *case]) +for case in [TEST_CASE_5, TEST_CASE_5_A, TEST_CASE_6]: + TEST_CASES.append([ResNet, *case]) TEST_SCRIPT_CASES = [ [model, *TEST_CASE_1] for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200] diff --git a/tests/test_senet.py b/tests/test_senet.py index 80d2b071c3..34f140638e 100644 --- a/tests/test_senet.py +++ b/tests/test_senet.py @@ -19,7 +19,7 @@ import monai.networks.nets.senet as se_mod from monai.networks import eval_mode -from monai.networks.nets import SENet154, SEResNet50, SEResNet101, SEResNet152, SEResNext50, SEResNext101 +from monai.networks.nets import SENet, SENet154, SEResNet50, SEResNet101, SEResNet152, SEResNext50, SEResNext101 from monai.utils import optional_import from tests.utils import test_is_quick, test_pretrained_networks, test_script_save, testing_data_config @@ -41,12 +41,24 @@ TEST_CASE_4 = [SEResNet152, NET_ARGS] TEST_CASE_5 = [SEResNext50, NET_ARGS] TEST_CASE_6 = [SEResNext101, NET_ARGS] +TEST_CASE_7 = [ + SENet, + { + "spatial_dims": 3, + "in_channels": 2, + "num_classes": 2, + "block": "se_bottleneck", + "layers": (3, 8, 36, 3), + "groups": 64, + "reduction": 16, + }, +] TEST_CASE_PRETRAINED_1 = [SEResNet50, {"spatial_dims": 2, "in_channels": 3, "num_classes": 2, "pretrained": True}] class TestSENET(unittest.TestCase): - @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6]) + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7]) def test_senet_shape(self, net, net_args): input_data = torch.randn(2, 2, 64, 64, 64).to(device) expected_shape = (2, 2)