Repository: Project-MONAI/MONAI

Issue:
DynUnet get_feature_maps() issue when using distributed training
**Describe the bug**
If we use `DistributedDataParallel` to wrap the network, calling the 'get_feature_maps' function will raise the following error:

![Screen Shot 2021-02-08 at 6 51 01 PM](https://user-images.githubusercontent.com/68361391/107210065-9aa5b600-6a3e-11eb-92f6-bbe4823d90ce.png)

It is common to return multiple variables in a network's forward function, not only for this case, but also for other situations such as for calculating triplet loss in metric learning based tasks. Therefore, here I think we should re-consider about what to return in the forward function of DynUnet.

Let me think about it and submit a PR then we can discuss it later.
@Nic-Ma @wyli @rijobro 



Discussion:
I recommend to do below change to fix this issue:
Remove the `get_feature_maps()` and change the logic of `forward()`:
1. if in `training=true` and `deep_supervision=true`, return a list of `Tensor`.
2. else, return the output `Tensor`.

It's slightly different from the original code which always returns a list of `Tensor`.

@wyli  @rijobro @ericspod , what do you guys think?

Thanks.

This will revert the behaviour of https://github.com/Project-MONAI/MONAI/pull/1393. I suppose this is fine as long as we make the documentation clear so that future users aren't confused.
Hi @rijobro ,

I think most of the previous problems are due to the `list` output during validation or inference.
So I suggest to return a list of data during training, return only the output instead of [out] during validation or inference.

Thanks.

> Hi @rijobro ,
> 
> I think most of the previous problems are due to the `list` output during validation or inference.
> So I suggest to return a list of data during training, return only the output instead of [out] during validation or inference.
> 
> Thanks.

Hi @Nic-Ma @rijobro , actually at the beginning, in val or infer modes, it does only return the output tensor. However, return different types are not permitted for torchscript, thus start from [this version](https://github.com/Project-MONAI/MONAI/commit/a2922bb7c4e582303cebf1ca2e3122c40b4bae29#diff-0095d268268e7ee415e079c1f0ccabed72f2c722f92cb0780d2ba8133a088334), list based results will be returned.
If we all return a list, then our default inferrer is not compatible since it relies on tensor based input.
@ericspod @wyli @rijobro ,

Do you guys have any better idea to solve this issue? Seems it's not easy to get a perfect solution..
Thanks in advance.

So it seems that we need to be consistent with what we return, regardless of train mode, etc., is that correct?

If so, what if we return:
```python
if self.deep_supervision:
     return dict{"data": self.output_block(out), "feature_map": self.heads[1 : self.deep_supr_num + 1]}
else:
    return self.output_block(out) 
```
In this fashion the output is always consistent (since `deep_supervision` is set at construction time and will not change over the lifetime of the object).


By returning a dictionary, hopefully it will be clearer to users what we are returning (since users were confused by the list previously returned).
> So it seems that we need to be consistent with what we return, regardless of train mode, etc., is that correct?
> 
> If so, what if we return:
> 
> ```python
> if self.deep_supervision:
>      return dict{"data": self.output_block(out), "feature_map": self.heads[1 : self.deep_supr_num + 1]}
> else:
>     return self.output_block(out) 
> ```
> 
> In this fashion the output is always consistent (since `deep_supervision` is set at construction time and will not change over the lifetime of the object).
> 
> By returning a dictionary, hopefully it will be clearer to users what we are returning (since users were confused by the list previously returned).

thanks this looks good to me as well, we could have the default value `self.deep_supervision=False`, so that the user wouldn't get this error by default
```
seg_prob = predictor(window_data, *args, **kwargs).to(device) # batched patch segmentation
AttributeError: 'list' object has no attribute 'to'
```
@yiheng-wang-nv ,

Please have a try with dict return first, I am afraid it can't fix the TorchScript issue...
Thanks.
If the default inferrer requires a tensor, then presumably you could wrap that function:

```python
def default_dict_inferrer(input, key, *args, **kwargs):
    return default_inferrer(input[key], *args, **kwargs)
```
Hi @rijobro , thanks for the advice, but return this kind of dict doesn't fix the torchscript issue as @Nic-Ma mentioned. I suggest here we still return a list, and add the corresponding docstrings. In order to help users use this network, I will also add/update tutorials. Let me submit a PR first for you to review.

Resolving patch:
diff --git a/monai/networks/nets/dynunet.py b/monai/networks/nets/dynunet.py
--- a/monai/networks/nets/dynunet.py
+++ b/monai/networks/nets/dynunet.py
@@ -14,6 +14,7 @@
 
 import torch
 import torch.nn as nn
+from torch.nn.functional import interpolate
 
 from monai.networks.blocks.dynunet_block import UnetBasicBlock, UnetOutBlock, UnetResBlock, UnetUpBlock
 
@@ -80,8 +81,22 @@ class DynUNet(nn.Module):
         upsample_kernel_size: convolution kernel size for transposed convolution layers.
         norm_name: [``"batch"``, ``"instance"``, ``"group"``]
             feature normalization type and arguments.
+        deep_supervision: whether to add deep supervision head before output. Defaults to ``False``.
+            If ``True``, in training mode, the forward function will output not only the last feature
+            map, but also the previous feature maps that come from the intermediate up sample layers.
+            In order to unify the return type (the restriction of TorchScript), all intermediate
+            feature maps are interpolated into the same size as the last feature map and stacked together
+            (with a new dimension in the first axis)into one single tensor.
+            For instance, if there are three feature maps with shapes: (1, 2, 32, 24), (1, 2, 16, 12) and
+            (1, 2, 8, 6). The last two will be interpolated into (1, 2, 32, 24), and the stacked tensor
+            will has the shape (1, 3, 2, 8, 6).
+            When calculating the loss, you can use torch.unbind to get all feature maps can compute the loss
+            one by one with the groud truth, then do a weighted average for all losses to achieve the final loss.
+            (To be added: a corresponding tutorial link)
+
         deep_supr_num: number of feature maps that will output during deep supervision head. The
-            value should be less than the number of up sample layers. Defaults to 1.
+            value should be larger than 0 and less than the number of up sample layers.
+            Defaults to 1.
         res_block: whether to use residual connection based convolution blocks during the network.
             Defaults to ``True``.
     """
@@ -95,6 +110,7 @@ def __init__(
         strides: Sequence[Union[Sequence[int], int]],
         upsample_kernel_size: Sequence[Union[Sequence[int], int]],
         norm_name: str = "instance",
+        deep_supervision: bool = False,
         deep_supr_num: int = 1,
         res_block: bool = False,
     ):
@@ -113,6 +129,7 @@ def __init__(
         self.bottleneck = self.get_bottleneck()
         self.upsamples = self.get_upsamples()
         self.output_block = self.get_output_block(0)
+        self.deep_supervision = deep_supervision
         self.deep_supervision_heads = self.get_deep_supervision_heads()
         self.deep_supr_num = deep_supr_num
         self.apply(self.initialize_weights)
@@ -140,6 +157,8 @@ def create_skips(index, downsamples, upsamples, superheads, bottleneck):
                 return bottleneck
             if index == 0:  # don't associate a supervision head with self.input_block
                 current_head, rest_heads = nn.Identity(), superheads
+            elif not self.deep_supervision:  # bypass supervision heads by passing nn.Identity in place of a real one
+                current_head, rest_heads = nn.Identity(), superheads[1:]
             else:
                 current_head, rest_heads = superheads[0], superheads[1:]
 
@@ -176,19 +195,21 @@ def check_kernel_stride(self):
     def check_deep_supr_num(self):
         deep_supr_num, strides = self.deep_supr_num, self.strides
         num_up_layers = len(strides) - 1
-        if deep_supr_num < 1 or deep_supr_num >= num_up_layers:
+        if deep_supr_num >= num_up_layers:
             raise AssertionError("deep_supr_num should be less than the number of up sample layers.")
+        if deep_supr_num < 1:
+            raise AssertionError("deep_supr_num should be larger than 0.")
 
     def forward(self, x):
         out = self.skip_layers(x)
-        return self.output_block(out)
-
-    def get_feature_maps(self):
-        """
-        Return the feature maps.
-
-        """
-        return self.heads[1 : self.deep_supr_num + 1]
+        out = self.output_block(out)
+        if self.training and self.deep_supervision:
+            out_all = [out]
+            feature_maps = self.heads[1 : self.deep_supr_num + 1]
+            for feature_map in feature_maps:
+                out_all.append(interpolate(feature_map, out.shape[2:]))
+            return torch.stack(out_all, dim=1)
+        return out
 
     def get_input_block(self):
         return self.conv_block(


Test patch:
diff --git a/tests/test_dynunet.py b/tests/test_dynunet.py
--- a/tests/test_dynunet.py
+++ b/tests/test_dynunet.py
@@ -43,6 +43,7 @@
                         "strides": strides,
                         "upsample_kernel_size": strides[1:],
                         "norm_name": "batch",
+                        "deep_supervision": False,
                         "res_block": res_block,
                     },
                     (1, in_channels, in_size, in_size),
@@ -65,6 +66,7 @@
                 "strides": ((1, 2, 1), 2, 2, 1),
                 "upsample_kernel_size": (2, 2, 1),
                 "norm_name": "instance",
+                "deep_supervision": False,
                 "res_block": res_block,
             },
             (1, in_channels, in_size, in_size, in_size),
@@ -77,6 +79,7 @@
     for res_block in [True, False]:
         for deep_supr_num in [1, 2]:
             for strides in [(1, 2, 1, 2, 1), (2, 2, 2, 1), (2, 1, 1, 2, 2)]:
+                scale = strides[0]
                 test_case = [
                     {
                         "spatial_dims": spatial_dims,
@@ -86,18 +89,13 @@
                         "strides": strides,
                         "upsample_kernel_size": strides[1:],
                         "norm_name": "group",
+                        "deep_supervision": True,
                         "deep_supr_num": deep_supr_num,
                         "res_block": res_block,
                     },
                     (1, 1, *[in_size] * spatial_dims),
+                    (1, 1 + deep_supr_num, 2, *[in_size // scale] * spatial_dims),
                 ]
-                scale = 1
-                all_expected_shapes = []
-                for stride in strides[: 1 + deep_supr_num]:
-                    scale *= stride
-                    deep_out_shape = (1, 2, *[in_size // scale] * spatial_dims)
-                    all_expected_shapes.append(deep_out_shape)
-                test_case.append(all_expected_shapes)
                 TEST_CASE_DEEP_SUPERVISION.append(test_case)
 
 
@@ -121,11 +119,8 @@ class TestDynUNetDeepSupervision(unittest.TestCase):
     def test_shape(self, input_param, input_shape, expected_shape):
         net = DynUNet(**input_param).to(device)
         with torch.no_grad():
-            results = [net(torch.randn(input_shape).to(device))] + net.get_feature_maps()
-            self.assertEqual(len(results), len(expected_shape))
-            for idx in range(len(results)):
-                result, sub_expected_shape = results[idx], expected_shape[idx]
-                self.assertEqual(result.shape, sub_expected_shape)
+            results = net(torch.randn(input_shape).to(device))
+            self.assertEqual(results.shape, expected_shape)
 
 
 if __name__ == "__main__":