diff --git a/monai/networks/nets/resnet.py b/monai/networks/nets/resnet.py index c8be9f0e89..16d1d26360 100644 --- a/monai/networks/nets/resnet.py +++ b/monai/networks/nets/resnet.py @@ -172,7 +172,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 +192,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/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]