Fix a bug in dice loss when there exist ignored pixels

mmseg/models/losses/dice_loss.py

@@ -40,10 +40,12 @@ def binary_dice_loss(pred, target, valid_mask, smooth=1, exponent=2, **kwargs):
     pred = pred.reshape(pred.shape[0], -1)
     target = target.reshape(target.shape[0], -1)
     valid_mask = valid_mask.reshape(valid_mask.shape[0], -1)
-
+    # A quick check. When valid_mask is all 0, dice loss should also be 0,
+    #  regardless of pred and target.
     num = torch.sum(torch.mul(pred, target) * valid_mask, dim=1) * 2 + smooth
-    den = torch.sum(pred.pow(exponent) + target.pow(exponent), dim=1) + smooth
-
+    den = torch.sum(
+        (pred.pow(exponent) + target.pow(exponent)) * valid_mask,
+        dim=1) + smooth
     return 1 - num / den
 
 
@@ -105,9 +107,16 @@ def forward(self,
 
         pred = F.softmax(pred, dim=1)
         num_classes = pred.shape[1]
+
+        # Convert target from B x H x W to B x H x W x (C + 1)
+        # For example, given one batch of labels [0, 1, 255] with C = 2
+        # The one_hot label should be [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
+        # Note that the last dim is not used. Therefore, the final one-hot
+        # label is [[1, 0, 0], [0, 1, 0]]
         one_hot_target = F.one_hot(
-            torch.clamp(target.long(), 0, num_classes - 1),
-            num_classes=num_classes)
+            torch.clamp(target.long(), 0, num_classes),
+            num_classes=num_classes + 1)
+        one_hot_target = one_hot_target[..., :num_classes]
         valid_mask = (target != self.ignore_index).long()
 
         loss = self.loss_weight * dice_loss(

tests/test_models/test_losses/test_dice_loss.py

@@ -18,6 +18,20 @@ def test_dice_lose():
     labels = (torch.rand(8, 4, 4) * 3).long()
     dice_loss(logits, labels)
 
+    # test dice loss with ignore_index
+    loss_cfg = dict(
+        type='DiceLoss',
+        reduction='mean',
+        class_weight=[1.0, 2.0, 3.0],
+        loss_weight=1.0,
+        ignore_index=255,
+        loss_name='loss_dice')
+    dice_loss = build_loss(loss_cfg)
+    logits = torch.rand(8, 3, 4, 4)
+    # all labels are ignored
+    labels = (torch.ones(8, 4, 4) * 255).long()
+    assert dice_loss(logits, labels).item() == 0
+
     # test loss with class weights from file
     import os
     import tempfile