fix(networks): mirror subunit chain in ResidualUnit residual for explicit padding

With subunits>1 the single projection conv only reproduced one subunit's shrinkage, so cx + res mismatched. Build a Sequential mirroring the subunit geometry when padding is explicit.

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/networks/blocks/convolutions.py

@@ -303,16 +303,32 @@ def __init__(
 
         # 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 or not same_pad:
-            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
-
             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

@@ -171,6 +171,12 @@ def test_padding0_identity_residual(self):
         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()