diff --git a/monai/networks/nets/daf3d.py b/monai/networks/nets/daf3d.py index 4e47d79a2d..9d4397319b 100644 --- a/monai/networks/nets/daf3d.py +++ b/monai/networks/nets/daf3d.py @@ -172,13 +172,21 @@ class Daf3dResNetBottleneck(ResNetBottleneck): spatial_dims: number of spatial dimensions of the input image. stride: stride to use for second conv layer. downsample: which downsample layer to use. + act: activation type and arguments. Defaults to relu. norm: which normalization layer to use. Defaults to group. """ expansion = 2 def __init__( - self, in_planes, planes, spatial_dims=3, stride=1, downsample=None, norm=("group", {"num_groups": 32}) + self, + in_planes, + planes, + spatial_dims=3, + stride=1, + downsample=None, + act=("relu", {"inplace": True}), + norm=("group", {"num_groups": 32}), ): conv_type: Callable = Conv[Conv.CONV, spatial_dims] @@ -191,7 +199,7 @@ def __init__( norm_layer(channels=planes * self.expansion), ) - super().__init__(in_planes, planes, spatial_dims, stride, downsample) + super().__init__(in_planes, planes, spatial_dims, stride, downsample, act) # change norm from batch to group norm self.bn1 = norm_layer(channels=planes) @@ -216,12 +224,21 @@ class Daf3dResNetDilatedBottleneck(Daf3dResNetBottleneck): spatial_dims: number of spatial dimensions of the input image. stride: stride to use for second conv layer. downsample: which downsample layer to use. + act: activation type and arguments. Defaults to relu. + norm: which normalization layer to use. Defaults to group. """ def __init__( - self, in_planes, planes, spatial_dims=3, stride=1, downsample=None, norm=("group", {"num_groups": 32}) + self, + in_planes, + planes, + spatial_dims=3, + stride=1, + downsample=None, + act=("relu", {"inplace": True}), + norm=("group", {"num_groups": 32}), ): - super().__init__(in_planes, planes, spatial_dims, stride, downsample, norm) + super().__init__(in_planes, planes, spatial_dims, stride, downsample, act, norm) # add dilation in second convolution conv_type: Callable = Conv[Conv.CONV, spatial_dims] diff --git a/monai/networks/nets/resnet.py b/monai/networks/nets/resnet.py index 9142116ae4..b99d890834 100644 --- a/monai/networks/nets/resnet.py +++ b/monai/networks/nets/resnet.py @@ -271,10 +271,18 @@ def __init__( self.bn1 = norm_layer self.act = get_act_layer(name=act) self.maxpool = pool_type(kernel_size=3, stride=2, padding=1) - self.layer1 = self._make_layer(block, block_inplanes[0], layers[0], spatial_dims, shortcut_type) - self.layer2 = self._make_layer(block, block_inplanes[1], layers[1], spatial_dims, shortcut_type, stride=2) - self.layer3 = self._make_layer(block, block_inplanes[2], layers[2], spatial_dims, shortcut_type, stride=2) - self.layer4 = self._make_layer(block, block_inplanes[3], layers[3], spatial_dims, shortcut_type, stride=2) + self.layer1 = self._make_layer( + block, block_inplanes[0], layers[0], spatial_dims, shortcut_type, act=act, norm=norm + ) + self.layer2 = self._make_layer( + block, block_inplanes[1], layers[1], spatial_dims, shortcut_type, stride=2, act=act, norm=norm + ) + self.layer3 = self._make_layer( + block, block_inplanes[2], layers[2], spatial_dims, shortcut_type, stride=2, act=act, norm=norm + ) + self.layer4 = self._make_layer( + block, block_inplanes[3], layers[3], spatial_dims, shortcut_type, stride=2, act=act, norm=norm + ) self.avgpool = avgp_type(block_avgpool[spatial_dims]) self.fc = nn.Linear(block_inplanes[3] * block.expansion, num_classes) if feed_forward else None @@ -282,8 +290,11 @@ def __init__( if isinstance(m, conv_type): nn.init.kaiming_normal_(torch.as_tensor(m.weight), mode="fan_out", nonlinearity="relu") elif isinstance(m, type(norm_layer)): - nn.init.constant_(torch.as_tensor(m.weight), 1) - nn.init.constant_(torch.as_tensor(m.bias), 0) + # non-affine norm layers (e.g. instance/layer norm defaults) have no weight/bias + if m.weight is not None: + nn.init.constant_(torch.as_tensor(m.weight), 1) + if m.bias is not None: + nn.init.constant_(torch.as_tensor(m.bias), 0) elif isinstance(m, nn.Linear): nn.init.constant_(torch.as_tensor(m.bias), 0) @@ -301,6 +312,7 @@ def _make_layer( spatial_dims: int, shortcut_type: str, stride: int = 1, + act: str | tuple = ("relu", {"inplace": True}), norm: str | tuple = "batch", ) -> nn.Sequential: conv_type: Callable = Conv[Conv.CONV, spatial_dims] @@ -333,13 +345,14 @@ def _make_layer( spatial_dims=spatial_dims, stride=stride, downsample=downsample, + act=act, norm=norm, ) ] self.in_planes = planes * block.expansion for _i in range(1, blocks): - layers.append(block(self.in_planes, planes, spatial_dims=spatial_dims, norm=norm)) + layers.append(block(self.in_planes, planes, spatial_dims=spatial_dims, act=act, norm=norm)) return nn.Sequential(*layers) diff --git a/tests/networks/nets/test_resnet.py b/tests/networks/nets/test_resnet.py index 241f57c78d..e792c9d81e 100644 --- a/tests/networks/nets/test_resnet.py +++ b/tests/networks/nets/test_resnet.py @@ -202,7 +202,7 @@ (1, 3), ] -TEST_CASE_9 = [ # Layer norm +TEST_CASE_9 = [ # Group norm { "block": ResNetBlock, "layers": [3, 4, 6, 3], @@ -213,7 +213,7 @@ "conv1_t_size": [3], "conv1_t_stride": 1, "act": ("relu", {"inplace": False}), - "norm": ("layer", {"normalized_shape": (64, 32)}), + "norm": ("group", {"num_groups": 8}), }, (1, 2, 32), (1, 3), @@ -232,6 +232,16 @@ [model, *TEST_CASE_1] for model in [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200] ] +# small 2D net used by the norm/act threading tests +TEST_CASE_NORM_ACT = { + "block": "basic", + "layers": [1, 1, 1, 1], + "block_inplanes": [8, 16, 32, 64], + "spatial_dims": 2, + "n_input_channels": 1, + "num_classes": 2, +} + CASE_EXTRACT_FEATURES = [ ( {"model_name": "resnet10", "pretrained": True, "spatial_dims": 3, "in_channels": 1}, @@ -316,6 +326,37 @@ def test_script(self, model, input_param, input_shape, expected_shape): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_norm_act_reach_blocks(self): + """`norm`/`act` given to the constructor must be used by the residual blocks and downsamples.""" + net = ResNet(**{**TEST_CASE_NORM_ACT, "norm": ("instance", {"affine": True}), "act": ("leakyrelu", {})}) + for layer in (net.layer1, net.layer2, net.layer3, net.layer4): + for block in layer: + self.assertIsInstance(block.bn1, torch.nn.InstanceNorm2d) + self.assertIsInstance(block.bn2, torch.nn.InstanceNorm2d) + self.assertIsInstance(block.act, torch.nn.LeakyReLU) + if block.downsample is not None: + self.assertIsInstance(block.downsample[1], torch.nn.InstanceNorm2d) + + def test_non_affine_norm_init(self): + """Non-affine norms have no weight/bias, the init loop must not choke on them.""" + net = ResNet(**{**TEST_CASE_NORM_ACT, "norm": "instance"}) + self.assertIsInstance(net.layer1[0].bn1, torch.nn.InstanceNorm2d) + with eval_mode(net): + net.forward(torch.randn(1, 1, 32, 32)) + + def test_default_norm_act_unchanged(self): + """Default construction must keep the same module tree as before the norm/act threading.""" + net = ResNet(**TEST_CASE_NORM_ACT) + for layer in (net.layer1, net.layer2, net.layer3, net.layer4): + for block in layer: + self.assertIsInstance(block.bn1, torch.nn.BatchNorm2d) + self.assertIsInstance(block.bn2, torch.nn.BatchNorm2d) + self.assertIsInstance(block.act, torch.nn.ReLU) + self.assertTrue(block.act.inplace) + # BatchNorm weights/biases are still initialised to 1/0 by the guarded init loop + self.assertTrue(torch.equal(net.layer1[0].bn1.weight, torch.ones_like(net.layer1[0].bn1.weight))) + self.assertTrue(torch.equal(net.layer1[0].bn1.bias, torch.zeros_like(net.layer1[0].bn1.bias))) + @SkipIfNoModule("hf_hub_download") class TestExtractFeatures(unittest.TestCase):