3918 Enhance the 1st conv layer of ResNet

monai/networks/nets/resnet.py

@@ -10,13 +10,14 @@
 # limitations under the License.
 
 from functools import partial
-from typing import Any, Callable, List, Optional, Type, Union
+from typing import Any, Callable, List, Optional, Tuple, Type, Union
 
 import torch
 import torch.nn as nn
 
 from monai.networks.layers.factories import Conv, Norm, Pool
 from monai.networks.layers.utils import get_pool_layer
+from monai.utils import ensure_tuple_rep
 from monai.utils.module import look_up_option
 
 __all__ = ["ResNet", "resnet10", "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnet200"]
@@ -32,14 +33,6 @@ def get_avgpool():
     return [0, 1, (1, 1), (1, 1, 1)]
 
 
-def get_conv1(conv1_t_size: int, conv1_t_stride: int):
-    return (
-        [0, conv1_t_size, (conv1_t_size, 7), (conv1_t_size, 7, 7)],
-        [0, conv1_t_stride, (conv1_t_stride, 2), (conv1_t_stride, 2, 2)],
-        [0, (conv1_t_size // 2), (conv1_t_size // 2, 3), (conv1_t_size // 2, 3, 3)],
-    )
-
-
 class ResNetBlock(nn.Module):
     expansion = 1
 
@@ -184,8 +177,8 @@ def __init__(
         block_inplanes: List[int],
         spatial_dims: int = 3,
         n_input_channels: int = 3,
-        conv1_t_size: int = 7,
-        conv1_t_stride: int = 1,
+        conv1_t_size: Union[Tuple[int], int] = 7,
+        conv1_t_stride: Union[Tuple[int], int] = 1,
         no_max_pool: bool = False,
         shortcut_type: str = "B",
         widen_factor: float = 1.0,
@@ -207,18 +200,20 @@ def __init__(
         ]
 
         block_avgpool = get_avgpool()
-        conv1_kernel, conv1_stride, conv1_padding = get_conv1(conv1_t_size, conv1_t_stride)
         block_inplanes = [int(x * widen_factor) for x in block_inplanes]
 
         self.in_planes = block_inplanes[0]
         self.no_max_pool = no_max_pool
 
+        conv1_kernel_size = ensure_tuple_rep(conv1_t_size, spatial_dims)
+        conv1_stride = ensure_tuple_rep(conv1_t_stride, spatial_dims)
+
         self.conv1 = conv_type(
             n_input_channels,
             self.in_planes,
-            kernel_size=conv1_kernel[spatial_dims],
-            stride=conv1_stride[spatial_dims],
-            padding=conv1_padding[spatial_dims],
+            kernel_size=conv1_kernel_size,  # type: ignore
+            stride=conv1_stride,  # type: ignore
+            padding=tuple(k // 2 for k in conv1_kernel_size),  # type: ignore
             bias=False,
         )
         self.bn1 = norm_type(self.in_planes)

tests/test_net_adapter.py

@@ -41,7 +41,9 @@
 class TestNetAdapter(unittest.TestCase):
     @parameterized.expand([TEST_CASE_0, TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
     def test_shape(self, input_param, input_shape, expected_shape):
-        model = resnet18(spatial_dims=input_param["dim"])
+        spatial_dims = input_param["dim"]
+        stride = (1, 2, 2)[:spatial_dims]
+        model = resnet18(spatial_dims=spatial_dims, conv1_t_stride=stride)
         input_param["model"] = model
         net = NetAdapter(**input_param).to(device)
         with eval_mode(net):

tests/test_resnet.py

@@ -31,25 +31,54 @@
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 TEST_CASE_1 = [  # 3D, batch 3, 2 input channel
-    {"pretrained": False, "spatial_dims": 3, "n_input_channels": 2, "num_classes": 3},
+    {
+        "pretrained": False,
+        "spatial_dims": 3,
+        "n_input_channels": 2,
+        "num_classes": 3,
+        "conv1_t_size": 7,
+        "conv1_t_stride": (2, 2, 2),
+    },
     (3, 2, 32, 64, 48),
     (3, 3),
 ]
 
 TEST_CASE_2 = [  # 2D, batch 2, 1 input channel
-    {"pretrained": False, "spatial_dims": 2, "n_input_channels": 1, "num_classes": 3},
+    {
+        "pretrained": False,
+        "spatial_dims": 2,
+        "n_input_channels": 1,
+        "num_classes": 3,
+        "conv1_t_size": [7, 7],
+        "conv1_t_stride": [2, 2],
+    },
     (2, 1, 32, 64),
     (2, 3),
 ]
 
 TEST_CASE_2_A = [  # 2D, batch 2, 1 input channel, shortcut type A
-    {"pretrained": False, "spatial_dims": 2, "n_input_channels": 1, "num_classes": 3, "shortcut_type": "A"},
+    {
+        "pretrained": False,
+        "spatial_dims": 2,
+        "n_input_channels": 1,
+        "num_classes": 3,
+        "shortcut_type": "A",
+        "conv1_t_size": (7, 7),
+        "conv1_t_stride": 2,
+    },
     (2, 1, 32, 64),
     (2, 3),
 ]
 
 TEST_CASE_3 = [  # 1D, batch 1, 2 input channels
-    {"pretrained": False, "spatial_dims": 1, "n_input_channels": 2, "num_classes": 3},
+    {
+        "pretrained": False,
+        "spatial_dims": 1,
+        "n_input_channels": 2,
+        "num_classes": 3,
+        "conv1_t_size": [3],
+        "conv1_t_stride": 1,
+    },
     (1, 2, 32),
     (1, 3),
 ]