Adjust execution order of activation and masking in MaskedDiceLoss

Signed-off-by: ytl0623 <david89062388@gmail.com>

Author: ytl0623

monai/losses/dice.py

@@ -253,7 +253,34 @@ def forward(self, input: torch.Tensor, target: torch.Tensor, mask: torch.Tensor
             target: the shape should be BNH[WD].
             mask: the shape should B1H[WD] or 11H[WD].
         """
-        return self.spatial_weighted(input=input, target=target, mask=mask)  # type: ignore[no-any-return]
+
+        if self.sigmoid:
+            input = torch.sigmoid(input)
+
+        n_pred_ch = input.shape[1]
+        if self.softmax:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `softmax=True` ignored.")
+            else:
+                input = torch.softmax(input, 1)
+
+        if self.other_act is not None:
+            input = self.other_act(input)
+
+        was_sigmoid = self.sigmoid
+        was_softmax = self.softmax
+        was_other_act = self.other_act
+
+        self.sigmoid = False
+        self.softmax = False
+        self.other_act = None
+
+        try:
+            return self.spatial_weighted(input=input, target=target, mask=mask)  # type: ignore[no-any-return]
+        finally:
+            self.sigmoid = was_sigmoid
+            self.softmax = was_softmax
+            self.other_act = was_other_act
 
 
 class GeneralizedDiceLoss(_Loss):

tests/losses/test_masked_dice_loss.py

@@ -27,7 +27,7 @@
             "target": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]]),
             "mask": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]),
         },
-        0.500,
+        0.333333,
     ],
     [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
         {"include_background": True, "sigmoid": True, "smooth_nr": 1e-4, "smooth_dr": 1e-4},
@@ -36,7 +36,7 @@
             "target": torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[1.0, 0.0], [1.0, 0.0]]]]),
             "mask": torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]], [[[1.0, 1.0], [0.0, 0.0]]]]),
         },
-        0.422969,
+        0.301128,
     ],
     [  # shape: (2, 2, 3), (2, 1, 3)
         {"include_background": False, "to_onehot_y": True, "smooth_nr": 0, "smooth_dr": 0},
@@ -54,7 +54,7 @@
             "target": torch.tensor([[[1.0, 0.0, 0.0]], [[1.0, 1.0, 0.0]]]),
             "mask": torch.tensor([[[1.0, 1.0, 0.0]]]),
         },
-        0.47033,
+        0.579184,
     ],
     [  # shape: (2, 2, 3), (2, 1, 3)
         {