From 6ffa1ba256fdba32300eb77bc3f0ed11d0b34648 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Tue, 22 Sep 2026 12:11:45 +0100 Subject: [PATCH] feat(networks): opt-in native-resolution deep supervision output for DynUNet Deep supervision heads were always upsampled (nearest) and stacked on dim 1, while the reference nnU-Net keeps each head at its native resolution and downsamples the target instead. Add deep_supr_output="stack"|"list" (default "stack", behaviour unchanged) so the native-resolution form can go straight to DeepSupervisionLoss, and document how the network diverges from nnU-Net: architecture-only pointer to monai.apps.nnunet, weighting left to the loss, tutorial link fixed to the tutorials repo main branch. Signed-off-by: Soumya Snigdha Kundu Assisted-by: OpenAI Codex --- monai/networks/nets/dynunet.py | 27 +++++++- tests/networks/nets/test_dynunet.py | 102 ++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 3 deletions(-) diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py index d130b886a78..b31c99415e7 100644 --- a/monai/networks/nets/dynunet.py +++ b/monai/networks/nets/dynunet.py @@ -60,6 +60,17 @@ class DynUNet(nn.Module): `nnU-Net: Self-adapting Framework for U-Net-Based Medical Image Segmentation `_. `Optimized U-Net for Brain Tumor Segmentation `_. + This is the network architecture only; for the full nnU-Net pipeline (planning, preprocessing, training, + ensembling) see :py:mod:`monai.apps.nnunet`. + + Differences from the reference nnU-Net network: deep supervision heads are optional and limited by + ``deep_supr_num`` (nnU-Net supervises every decoder stage); by default their low-resolution logits are + upsampled with nearest-neighbour interpolation and stacked, whereas nnU-Net keeps each output at its native + resolution and downsamples the target instead (``deep_supr_output="list"`` gives that behaviour, to be used + with :py:class:`monai.losses.DeepSupervisionLoss`); no per-level loss weighting is built in (nnU-Net uses a + normalized ``1/2**level`` schedule with the deepest level dropped; here weighting is left to the loss); + convolutions have no bias, which is immaterial before affine instance norm. + This model is more flexible compared with ``monai.networks.nets.UNet`` in three places: @@ -87,7 +98,7 @@ class DynUNet(nn.Module): For backwards compatibility with old weights, please set `strict=False` when calling `load_state_dict`. Usage example with medical segmentation decathlon dataset is available at: - https://github.com/Project-MONAI/tutorials/tree/master/modules/dynunet_pipeline. + https://github.com/Project-MONAI/tutorials/tree/main/modules/dynunet_pipeline. Args: spatial_dims: number of spatial dimensions. @@ -125,6 +136,10 @@ class DynUNet(nn.Module): res_block: whether to use residual connection based convolution blocks during the network. Defaults to ``False``. trans_bias: whether to set the bias parameter in transposed convolution layers. Defaults to ``False``. + deep_supr_output: format of the training-mode output when ``deep_supervision=True``. ``"stack"`` (default): + heads are upsampled with nearest interpolation to the final output size and stacked along dim 1, as + described above. ``"list"``: returns ``[final, head_1, ..., head_k]`` with each head at its native + resolution, highest first, as expected by :py:class:`monai.losses.DeepSupervisionLoss`. """ def __init__( @@ -143,6 +158,7 @@ def __init__( deep_supr_num: int = 1, res_block: bool = False, trans_bias: bool = False, + deep_supr_output: str = "stack", ): super().__init__() self.spatial_dims = spatial_dims @@ -168,6 +184,9 @@ def __init__( self.output_block = self.get_output_block(0) self.deep_supervision = deep_supervision self.deep_supr_num = deep_supr_num + if deep_supr_output not in ("stack", "list"): + raise ValueError(f"deep_supr_output should be 'stack' or 'list', got {deep_supr_output!r}.") + self.deep_supr_output = deep_supr_output # initialize the typed list of supervision head outputs so that Torchscript can recognize what's going on self.heads: list[torch.Tensor] = [torch.rand(1)] * self.deep_supr_num if self.deep_supervision: @@ -265,10 +284,12 @@ def check_filters(self): else: self.filters = filters[: len(self.strides)] - def forward(self, x): - out = self.skip_layers(x) + def forward(self, x: torch.Tensor) -> torch.Tensor | list[torch.Tensor]: + out: torch.Tensor = self.skip_layers(x) out = self.output_block(out) if self.training and self.deep_supervision: + if self.deep_supr_output == "list": + return [out] + list(self.heads) out_all = [out] for feature_map in self.heads: out_all.append(interpolate(feature_map, out.shape[2:])) diff --git a/tests/networks/nets/test_dynunet.py b/tests/networks/nets/test_dynunet.py index c2c9369923a..a573956296e 100644 --- a/tests/networks/nets/test_dynunet.py +++ b/tests/networks/nets/test_dynunet.py @@ -17,6 +17,7 @@ import torch from parameterized import parameterized +from monai.losses import DeepSupervisionLoss, DiceCELoss from monai.networks import eval_mode from monai.networks.nets import DynUNet from monai.utils import optional_import @@ -123,6 +124,40 @@ ] TEST_CASE_DEEP_SUPERVISION.append(test_case) +TEST_CASE_DEEP_SUPERVISION_LIST = [] +for params in dict_product( + spatial_dims=[2, 3], + res_block=[True, False], + deep_supr_num=[1, 2], + strides=[(1, 2, 1, 2, 1), (2, 2, 2, 1), (2, 1, 1, 2, 2)], +): + spatial_dims = params["spatial_dims"] + deep_supr_num = params["deep_supr_num"] + strides = params["strides"] + res_block = params["res_block"] + # each head sits at a skip resolution: cumulative stride products of the input/downsample blocks + cum_strides = [strides[0]] + for stride in strides[1:-1]: + cum_strides.append(cum_strides[-1] * stride) + test_case = [ + { + "spatial_dims": spatial_dims, + "in_channels": 1, + "out_channels": 2, + "kernel_size": [3] * len(strides), + "strides": strides, + "upsample_kernel_size": strides[1:], + "norm_name": ("group", {"num_groups": 16}), + "deep_supervision": True, + "deep_supr_num": deep_supr_num, + "deep_supr_output": "list", + "res_block": res_block, + }, + (1, 1, *[in_size_ds] * spatial_dims), + [(1, 2, *[in_size_ds // cum_strides[level]] * spatial_dims) for level in range(deep_supr_num + 1)], + ] + TEST_CASE_DEEP_SUPERVISION_LIST.append(test_case) + class TestDynUNet(unittest.TestCase): @parameterized.expand(TEST_CASE_DYNUNET_3D) @@ -185,5 +220,72 @@ def test_shape(self, input_param, input_shape, expected_shape): self.assertEqual(results.shape, expected_shape) +TEST_CASE_DEEP_SUPERVISION_LIST_2D = { + "spatial_dims": 2, + "in_channels": 1, + "out_channels": 2, + "kernel_size": [3] * 5, + "strides": [1, 2, 1, 2, 1], + "upsample_kernel_size": [2, 1, 2, 1], + "deep_supervision": True, + "deep_supr_num": 2, + "deep_supr_output": "list", +} + + +class TestDynUNetDeepSupervisionList(unittest.TestCase): + @parameterized.expand(TEST_CASE_DEEP_SUPERVISION_LIST) + def test_shape(self, input_param, input_shape, expected_shapes): + net = DynUNet(**input_param).to(device) + with torch.no_grad(): + results = net(torch.randn(input_shape).to(device)) + self.assertIsInstance(results, list) + self.assertEqual(len(results), len(expected_shapes)) + for result, expected_shape in zip(results, expected_shapes): + self.assertEqual(tuple(result.shape), expected_shape) + + def test_with_deep_supervision_loss(self): + net = DynUNet(**TEST_CASE_DEEP_SUPERVISION_LIST_2D).to(device) + net.train() + inputs = torch.randn(1, 1, 32, 32).to(device) + target = (torch.rand(1, 2, 32, 32).to(device) > 0.5).float() + outputs = net(inputs) + self.assertIsInstance(outputs, list) + loss = DeepSupervisionLoss(DiceCELoss(sigmoid=True))(outputs, target) + self.assertTrue(torch.isfinite(loss)) + loss.backward() + for head in net.deep_supervision_heads: + for param in head.parameters(): + self.assertIsNotNone(param.grad) + + def test_eval_returns_single_tensor(self): + net = DynUNet(**TEST_CASE_DEEP_SUPERVISION_LIST_2D).to(device) + inputs = torch.randn(1, 1, 32, 32).to(device) + with eval_mode(net): + result = net(inputs) + self.assertIsInstance(result, torch.Tensor) + self.assertEqual(tuple(result.shape), (1, 2, 32, 32)) + + def test_list_output_ignored_without_deep_supervision(self): + net = DynUNet(**{**TEST_CASE_DEEP_SUPERVISION_LIST_2D, "deep_supervision": False}).to(device) + net.train() + inputs = torch.randn(1, 1, 32, 32).to(device) + result = net(inputs) + self.assertIsInstance(result, torch.Tensor) + self.assertEqual(tuple(result.shape), (1, 2, 32, 32)) + + def test_invalid_deep_supr_output(self): + with self.assertRaises(ValueError): + DynUNet( + spatial_dims=2, + in_channels=1, + out_channels=2, + kernel_size=[3] * 4, + strides=[1, 2, 2, 1], + upsample_kernel_size=[2, 2, 1], + deep_supr_output="invalid", + ) + + if __name__ == "__main__": unittest.main()