fix(networks): honor explicit padding=0 in ResidualUnit

monai/networks/blocks/convolutions.py

@@ -269,7 +269,8 @@ def __init__(
         self.out_channels = out_channels
         self.conv = nn.Sequential()
         self.residual = nn.Identity()
-        if not padding:
+        same_pad = padding is None
+        if padding is None:
             padding = same_padding(kernel_size, dilation)
         schannels = in_channels
         sstrides = strides
@@ -301,16 +302,33 @@ def __init__(
             sstrides = 1
 
         # apply convolution to input to change number of output channels and size to match that coming from self.conv
-        if np.prod(strides) != 1 or in_channels != out_channels:
-            rkernel_size = kernel_size
-            rpadding = padding
-
-            if np.prod(strides) == 1:  # if only adapting number of channels a 1x1 kernel is used with no padding
-                rkernel_size = 1
-                rpadding = 0
-
+        if np.prod(strides) != 1 or in_channels != out_channels or not same_pad:
             conv_type = Conv[Conv.CONV, self.spatial_dims]
-            self.residual = conv_type(in_channels, out_channels, rkernel_size, strides, rpadding, bias=bias)
+            if not same_pad and subunits > 1:
+                # mirror the subunit chain so the residual shrinks exactly like the main path
+                residual = nn.Sequential()
+                schannels = in_channels
+                sstrides = strides
+                for su in range(subunits):
+                    residual.add_module(
+                        f"unit{su:d}",
+                        conv_type(
+                            schannels, out_channels, kernel_size, sstrides, padding, bias=bias, dilation=dilation
+                        ),
+                    )
+                    schannels = out_channels
+                    sstrides = 1
+                self.residual = residual
+            else:
+                rkernel_size = kernel_size
+                rpadding = padding
+
+                # if only adapting number of channels a 1x1 kernel is used with no padding
+                if np.prod(strides) == 1 and same_pad:
+                    rkernel_size = 1
+                    rpadding = 0
+
+                self.residual = conv_type(in_channels, out_channels, rkernel_size, strides, rpadding, bias=bias)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         res: torch.Tensor = self.residual(x)  # create the additive residual from x

tests/networks/blocks/test_convolutions.py

@@ -13,6 +13,8 @@
 
 import unittest
 
+import torch
+
 from monai.networks.blocks import Convolution, ResidualUnit
 from tests.test_utils import TorchImageTestCase2D, TorchImageTestCase3D
 
@@ -67,6 +69,12 @@ def test_transpose2(self):
         expected_shape = (1, self.output_channels, self.im_shape[0] * 2, self.im_shape[1] * 2)
         self.assertEqual(out.shape, expected_shape)
 
+    def test_padding0(self):
+        conv = Convolution(2, self.input_channels, self.output_channels, kernel_size=3, padding=0)
+        out = conv(self.imt)
+        expected_shape = (1, self.output_channels, self.im_shape[0] - 2, self.im_shape[1] - 2)
+        self.assertEqual(out.shape, expected_shape)
+
 
 class TestConvolution3D(TorchImageTestCase3D):
     def test_conv1(self):
@@ -151,6 +159,24 @@ def test_dropout1(self):
         expected_shape = (1, self.output_channels, self.im_shape[0], self.im_shape[1])
         self.assertEqual(out.shape, expected_shape)
 
+    def test_padding0(self):
+        conv = ResidualUnit(2, 1, self.output_channels, kernel_size=3, subunits=1, padding=0)
+        out = conv(self.imt)
+        expected_shape = (1, self.output_channels, self.im_shape[0] - 2, self.im_shape[1] - 2)
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_padding0_identity_residual(self):
+        conv = ResidualUnit(2, 2, 2, kernel_size=3, subunits=1, padding=0)
+        out = conv(torch.rand(1, 2, 8, 8))
+        expected_shape = (1, 2, 6, 6)
+        self.assertEqual(out.shape, expected_shape)
+
+    def test_padding0_default_subunits(self):
+        conv = ResidualUnit(2, 2, 2, kernel_size=3, padding=0)
+        out = conv(torch.rand(1, 2, 8, 8))
+        expected_shape = (1, 2, 4, 4)
+        self.assertEqual(out.shape, expected_shape)
+
 
 if __name__ == "__main__":
     unittest.main()