4691 support block type strings in resnet and senet

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]

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]

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]

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)