UpSample optional kernel_size for deconv mode

monai/networks/blocks/upsample.py

@@ -40,6 +40,7 @@ def __init__(
         in_channels: Optional[int] = None,
         out_channels: Optional[int] = None,
         scale_factor: Union[Sequence[float], float] = 2,
+        kernel_size: Optional[Union[Sequence[float], float]] = None,
         size: Optional[Union[Tuple[int], int]] = None,
         mode: Union[UpsampleMode, str] = UpsampleMode.DECONV,
         pre_conv: Optional[Union[nn.Module, str]] = "default",
@@ -54,6 +55,7 @@ def __init__(
             in_channels: number of channels of the input image.
             out_channels: number of channels of the output image. Defaults to `in_channels`.
             scale_factor: multiplier for spatial size. Has to match input size if it is a tuple. Defaults to 2.
+            kernel_size: kernel size used during UpsampleMode.DECONV. Defaults to `scale_factor`.
             size: spatial size of the output image.
                 Only used when ``mode`` is ``UpsampleMode.NONTRAINABLE``.
                 In torch.nn.functional.interpolate, only one of `size` or `scale_factor` should be defined,
@@ -83,13 +85,24 @@ def __init__(
         if up_mode == UpsampleMode.DECONV:
             if not in_channels:
                 raise ValueError(f"in_channels needs to be specified in the '{mode}' mode.")
+
+            if not kernel_size:
+                kernel_size_ = scale_factor_
+                output_padding = padding = 0
+            else:
+                kernel_size_ = ensure_tuple_rep(kernel_size, spatial_dims)
+                padding = tuple((k - 1) // 2 for k in kernel_size_)  # type: ignore
+                output_padding = tuple(s - 1 - (k - 1) % 2 for k, s in zip(kernel_size_, scale_factor_))  # type: ignore
+
             self.add_module(
                 "deconv",
                 Conv[Conv.CONVTRANS, spatial_dims](
                     in_channels=in_channels,
                     out_channels=out_channels or in_channels,
-                    kernel_size=scale_factor_,
+                    kernel_size=kernel_size_,
                     stride=scale_factor_,
+                    padding=padding,
+                    output_padding=output_padding,
                     bias=bias,
                 ),
             )

tests/test_milmodel.py

@@ -17,7 +17,7 @@
 from monai.networks import eval_mode
 from monai.networks.nets import MILModel
 from monai.utils.module import optional_import
-from tests.utils import test_script_save
+from tests.utils import skip_if_downloading_fails, test_script_save
 
 models, _ = optional_import("torchvision.models")
 
@@ -65,7 +65,8 @@
 class TestMilModel(unittest.TestCase):
     @parameterized.expand(TEST_CASE_MILMODEL)
     def test_shape(self, input_param, input_shape, expected_shape):
-        net = MILModel(**input_param).to(device)
+        with skip_if_downloading_fails():
+            net = MILModel(**input_param).to(device)
         with eval_mode(net):
             result = net(torch.randn(input_shape, dtype=torch.float).to(device))
             self.assertEqual(result.shape, expected_shape)

tests/test_upsample_block.py

@@ -88,13 +88,34 @@
         TEST_CASES_EQ.append(test_case)
 
 
+TEST_CASES_EQ2 = []  # type: ignore
+for s in range(2, 5):
+    for k in range(1, 7):
+        expected_shape = (16, 5, 4 * s, 5 * s, 6 * s)
+        for t in UpsampleMode:
+            test_case = [
+                {
+                    "spatial_dims": 3,
+                    "in_channels": 3,
+                    "out_channels": 5,
+                    "mode": t,
+                    "scale_factor": s,
+                    "kernel_size": k,
+                    "align_corners": False,
+                },
+                (16, 3, 4, 5, 6),
+                expected_shape,
+            ]
+            TEST_CASES_EQ.append(test_case)
+
+
 class TestUpsample(unittest.TestCase):
-    @parameterized.expand(TEST_CASES + TEST_CASES_EQ)
+    @parameterized.expand(TEST_CASES + TEST_CASES_EQ + TEST_CASES_EQ2)
     def test_shape(self, input_param, input_shape, expected_shape):
         net = UpSample(**input_param)
         with eval_mode(net):
             result = net(torch.randn(input_shape))
-            self.assertEqual(result.shape, expected_shape)
+            self.assertEqual(result.shape, expected_shape, msg=str(input_param))
 
 
 if __name__ == "__main__":