5542 amend extract levels in localnet (#5543)

Fixes #5542 

### Description

A few sentences describing the changes proposed in this pull request.

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: kate-sann5100 <[email protected]>
Signed-off-by: monai-bot <[email protected]>
Signed-off-by: Wenqi Li <[email protected]>
Co-authored-by: monai-bot <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wenqi Li <[email protected]>

Author: 4 people

monai/networks/nets/regunet.py

@@ -342,7 +342,7 @@ def __init__(self, spatial_dims: int, in_channels: int, out_channels: int):
         self.deconv = get_deconv_block(spatial_dims=spatial_dims, in_channels=in_channels, out_channels=out_channels)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        output_size = (size * 2 for size in x.shape[2:])
+        output_size = [size * 2 for size in x.shape[2:]]
         deconved = self.deconv(x)
         resized = F.interpolate(x, output_size)
         resized = torch.sum(torch.stack(resized.split(split_size=resized.shape[1] // 2, dim=1), dim=-1), dim=-1)
@@ -372,6 +372,7 @@ def __init__(
         out_activation: Optional[str] = None,
         out_channels: int = 3,
         pooling: bool = True,
+        use_addictive_sampling: bool = True,
         concat_skip: bool = False,
     ):
         """
@@ -384,12 +385,15 @@ def __init__(
             out_channels: number of channels for the output
             extract_levels: list, which levels from net to extract. The maximum level must equal to ``depth``
             pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv3d
+            use_addictive_sampling: whether use additive up-sampling layer for decoding.
             concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition
         """
+        self.use_additive_upsampling = use_addictive_sampling
         super().__init__(
             spatial_dims=spatial_dims,
             in_channels=in_channels,
             num_channel_initial=num_channel_initial,
+            extract_levels=extract_levels,
             depth=max(extract_levels),
             out_kernel_initializer=out_kernel_initializer,
             out_activation=out_activation,
@@ -406,7 +410,7 @@ def build_bottom_block(self, in_channels: int, out_channels: int):
         )
 
     def build_up_sampling_block(self, in_channels: int, out_channels: int) -> nn.Module:
-        if self._use_additive_upsampling:
+        if self.use_additive_upsampling:
             return AdditiveUpSampleBlock(
                 spatial_dims=self.spatial_dims, in_channels=in_channels, out_channels=out_channels
             )

tests/test_localnet.py

@@ -65,6 +65,12 @@ def test_shape(self, input_param, input_shape, expected_shape):
             result = net(torch.randn(input_shape).to(device))
             self.assertEqual(result.shape, expected_shape)
 
+    @parameterized.expand(TEST_CASE_LOCALNET_2D + TEST_CASE_LOCALNET_3D)
+    def test_extract_levels(self, input_param, input_shape, expected_shape):
+        net = LocalNet(**input_param).to(device)
+        self.assertEqual(len(net.decode_deconvs), len(input_param["extract_levels"]) - 1)
+        self.assertEqual(len(net.decode_convs), len(input_param["extract_levels"]) - 1)
+
     def test_script(self):
         input_param, input_shape, _ = TEST_CASE_LOCALNET_2D[0]
         net = LocalNet(**input_param)