feat(networks): opt-in native-resolution deep supervision output for DynUNet

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 <https://arxiv.org/abs/1809.10486>`_.
     `Optimized U-Net for Brain Tumor Segmentation <https://arxiv.org/pdf/2110.03352.pdf>`_.
 
+    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:]))

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()