DynUNet gives different shapes in training step and validation step

[YerePhy]

The network monai.networks.nets.DynUNet is giving different shapes in training_step and validation_step in Pytorch Lightning. You can reproduce the problem with this notebook.

[JohnMasoner]

Hi @YerePhy,

Thank you for providing more information. I tried to reproduce your question.

import monai
import torch

net = monai.networks.nets.DynUNet(2, in_channels=3, out_channels=3,
                                kernel_size=[3, 3, 3, 3], strides=[1, 2, 2, 2], 
                                upsample_kernel_size=[2, 2, 2, 2], filters=[64, 64, 64, 64],
                                deep_supervision=True, deep_supr_num=2)

rand = torch.randn((1, 3, 72, 72))
print('the net outputs shape:', net(rand).shape)
# the net outputs shape: torch.Size([1, 3, 3, 72, 72])

your problems seem is the dimensions are different.

As the DynUNet docs said,

When calculating the loss, you can use torch.unbind to get all feature maps can compute the loss one by one with the ground truth, then do a weighted average for all losses to achieve the final loss.

Maybe you can take a look at the following code(It is very unregulated code).

rand = torch.randn((1, 3, 72, 72))
net_rand = net(rand)
print('the net outputs shape:', net_rand.shape)
rand_unbinds = torch.unbind(net_rand, 1)
print('the net unbinds shape:', [rand_unbind.shape for rand_unbind in rand_unbinds])
criterion = monai.losses.DiceLoss()
loss = torch.mean(torch.FloatTensor([criterion(rand_unbind, torch.zeros_like(rand_unbind)) for rand_unbind in rand_unbinds]))
print('loss value:', loss)

Hope this is useful to you.

[YerePhy]

Hi John,

Thanks for your comment!

I'm aware of the expected shape of the output, the problem is that where you get # the net outputs shape: torch.Size([1, 3, 3, 72, 72]) I get # the net outputs shape: torch.Size([1, 3, 72, 72]) when I should get torch.Size([1, 3, 3, 72, 72]) just like you. And this only happens inside the training_step of my LightningModule. If I stop the execution in the __init__ of the LightningModule with the debugger and call the equivalent net(rand).shape I get the right shape torch.Size([1, 3, 3, 72, 72]).

I'm working on a jupyter notebook to reproduce this, let me finish that so we can understand each other better. But, generally, this is the problem:

import pytorch_lightning
import torch

class MyLightningModule(pytorch_lightning.LightningModule):

    def __init__(self, model: torch.nn.Module, other_kwargs: Any) -> None:
        self.model = model    # DynUnet
        self.other_kwargs = other_kwargs
        self.save_hyperparameters()
        # >>>> stop the excution here and do `self.model(torch.randn((1, 3, 72, 72))).shape`
        # >>>> Output: torch.Size([1, 3, 3, 72, 72]) as expected, everything looks great
    
    @property
    def otehr_kwargs(self) -> Any:
        return self._other_kwargs       
    
    @other_kwargs.setter
    def otehr_kwargs(self, other_kwargs: Any) -> Any:
        """
        Yes, I'm sure that `do_some_operation_that_not_changes_the_model`
        does not changes `self.model`.
        """
        self._other_kwargs = do_some_operation_that_not_changes_the_model(self.model, other_kwargs)
    
    def training_step(self, batch: Dict[str, Any], batch_idx: int) -> Dict[str, torch.Tensor]:
        input = batch["INPUTS"]    # has shape `N=1 C=3 H=72 W=72`
        
        outputs = self.model.forward(inputs)
        # >>>> stop the execution here and do `self.model(torch.randn((1, 3, 72, 72))).shape`
        # >>>> Ouptut: torch.Size([1, 3, 72, 72]) !!!!
        # Expected Output: torch.Size([1, 3, 3, 72, 72]) (one or more of DynUNet, Pytorch Lightning or myself are trolling)  !!!!
        # do also `print(self.model.deep_supervison) >>>> True` !!!!
        # do also `print(self.model.deep_supr_num)` >>>> 2` !!!!
        # do also `print(self.model.spatial_dims) >>>> 2`

    def configure_optimizers(self):
        ...

[YerePhy]

@JohnMasoner take a look to this notebook

[YerePhy]

Solution: monai.networks.nets.DynUNet behaves different in evaluation mode (e.g. .eval()) in that mode won't ouput the outputs of the supervision heads. So be careful because you will get different shapes in training_step and validation_step (where .eval() is applied under the hood).

[YerePhy]

Thanks beforehand,

Current docstring for deep_supervision param:

whether to add deep supervision head before output. Defaults to False. If True, in training mode, the forward function will output not only the final feature map (from output_block), but also the 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 final feature map and stacked together (with a new dimension in the first axis)into one single tensor. For instance, if there are two intermediate feature maps with shapes: (1, 2, 16, 12) and (1, 2, 8, 6), and the final feature map has the shape (1, 2, 32, 24), then all intermediate feature maps will be interpolated into (1, 2, 32, 24), and the stacked tensor will has the shape (1, 3, 2, 32, 24). When calculating the loss, you can use torch.unbind to get all feature maps can compute the loss one by one with the ground truth, then do a weighted average for all losses to achieve the final loss.

Suggested change:

whether to add deep supervision head before output. Defaults to False. If True, in training mode, the forward function will output not only the final feature map (from output_block), but also the feature maps that come from the intermediate up sample layers. In evaluation mode, the forward function will output only the final feature map. In order to unify the return type (the restriction of TorchScript), all intermediate feature maps are interpolated into the same size as the final feature map and stacked together (with a new dimension in the first axis)into one single tensor. For instance, if there are two intermediate feature maps with shapes: (1, 2, 16, 12) and (1, 2, 8, 6), and the final feature map has the shape (1, 2, 32, 24), then all intermediate feature maps will be interpolated into (1, 2, 32, 24), and the stacked tensor will has the shape (1, 3, 2, 32, 24). When calculating the loss, you can use torch.unbind to get all feature maps can compute the loss one by one with the ground truth, then do a weighted average for all losses to achieve the final loss.

I know that it is a small change, but with this sentence I'd avoid several days of debugging.

[Nic-Ma]

Hi @YerePhy ,

Thanks for your feedback! Would you like to submit a pull request for it directly?

[YerePhy]

Okay, let me try it and if I have any problem I'll delegate this to you or @yiheng-wang-nv

Thanks for this opportunity!

[Nic-Ma]

Cool, you can refer to the contribution guide doc:
https://github.com/Project-MONAI/MONAI/blob/dev/CONTRIBUTING.md

Thanks in advance!

[yiheng-wang-nv]

Thanks @YerePhy for the advice!