Repository: Project-MONAI/MONAI

Issue:
Incorrect results for MaskedDiceLoss when softmax=True
When the softmax parameter `softmax=True` is used in `MaskedDiceLoss`, the masked pixels are not excluded from the computation since the softmax is still applied to the zeroed-out pixels. Example:

```python
import math
import torch
from monai.losses import MaskedDiceLoss, DiceLoss

# shape: (batch, classes, height, width)
input_prob = torch.tensor([[[[0.1, 0.8]], [[0.9, 0.2]]]])  # [1, 2, 1, 2]
input_logits = torch.tensor([[[[math.log(0.1), math.log(0.8)]], [[math.log(0.9), math.log(0.2)]]]]) # [1, 2, 1, 2]
target = torch.tensor([[[[1, 0]]]])  # [1, 1, 1, 2]
mask = torch.tensor([[[[1, 0]]]])  # [1, 1, 1, 2]

dice_loss_prob = MaskedDiceLoss(to_onehot_y=True, softmax=False, reduction='none')
dice_loss_logits = MaskedDiceLoss(to_onehot_y=True, softmax=True, reduction='none')

# This yields the same result
dice_loss_prob(input_prob, target)  # tensor([[0.1579, 0.1429]])
dice_loss_logits(input_logits, target)  # tensor([[0.1579, 0.1429]])

# This does not (the first one is the correct result)
dice_loss_prob(input_prob, target, mask=mask)  # tensor([[1.0000, 0.0526]])
dice_loss_logits(input_logits, target, mask=mask)  # tensor([[0.3750, 0.2500]])
```

The problem is that the [zeroing happens before the softmax](https://github.com/Project-MONAI/MONAI/blob/master/monai/losses/dice.py#L173).

My solution would be to apply the masking before [this line](https://github.com/Project-MONAI/MONAI/blob/master/monai/losses/dice.py#L118). Actually, I don't know why there is a `MaskedDiceLoss` class in the first place. I think an additional (optional) argument to the forward function of `DiceLoss` would be sufficient. So, I propose:
- Remove the `MaskedDiceLoss` class.
- Add an additional (optional) argument to the `DiceLoss` class for the mask.
- Add the example above to the tests.

If this is fine with you, I could make a PR :-)

Bug is present in commit 11fc31dc1b6b37cf3ed6a3a89e39739629cd3147


Discussion:
I feel the docstring of `MaskedDiceLoss` is kind of misleading https://github.com/Project-MONAI/MONAI/blob/master/monai/losses/dice.py#L149 IMO the mask is more like a differentiable attention map and this is different from `DiceLoss`

Resolving patch:
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -159,7 +159,12 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, smooth: float = 1e-
 
 class MaskedDiceLoss(DiceLoss):
     """
-    Same as DiceLoss, but accepts a binary mask ([0,1]) indicating a region over which to compute the dice.
+    Add an additional `masking` process before `DiceLoss`, accept a binary mask ([0, 1]) indicating a region,
+    `input` and `target` will be masked by the region: region with mask `1` will keep the original value,
+    region with `0` mask will be converted to `0`. Then feed `input` and `target` to normal `DiceLoss` computation.
+    This has the effect of ensuring only the masked region contributes to the loss computation and
+    hence gradient calculation.
+
     """
 
     def forward(
@@ -187,6 +192,8 @@ def forward(
 
             input = input * mask
             target = target * mask
+        else:
+            warnings.warn("no mask value specified for the MaskedDiceLoss.")
 
         return super().forward(input=input, target=target, smooth=smooth)
 


Test patch: