Repository: kornia/kornia
Pull request: Add ignore_index to dice, focal and tversky loss
#### Changes
* adds `ignore_index` argument to dice and tversky losses. 
* a second commit also adds code to the focal loss, but without tests. I couldn't easily extend an existing test case to properly test the new parameter. IMHO, I would drop this second commit for now.

fixes #839

#### Type of change
- [x] 🧪 Tests Cases
- [x] 🔬 New feature (non-breaking change which adds functionality)


#### Checklist

- [x] My code follows the style guidelines of this project
- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings


Problem statement:
"ignore_index" argment for focal loss and dice loss
When I try to train for segmentation on VOC 2012 with Dice loss, I always run into cuda illegal error and I finally found that this may be caused by out-of-range indexes (255 out of 20 classes). Since pytorch CrossEntropyLoss has an argment "ignore_index" to deal with this issue, I wonder if there is a proper way to deal with uncertain indexes?

Hints from issue discussion:

Commit message: fix copy&paste bug in binary_focal_loss_with_logits

Patch to review:
diff --git a/kornia/losses/dice.py b/kornia/losses/dice.py
index 2b78a6838a..bbf032468f 100644
--- a/kornia/losses/dice.py
+++ b/kornia/losses/dice.py
@@ -15,7 +15,12 @@
 
 
 def dice_loss(
-    pred: Tensor, target: Tensor, average: str = "micro", eps: float = 1e-8, weight: Optional[Tensor] = None
+    pred: Tensor,
+    target: Tensor,
+    average: str = "micro",
+    eps: float = 1e-8,
+    weight: Optional[Tensor] = None,
+    ignore_index: int = -100,
 ) -> Tensor:
     r"""Criterion that computes Sørensen-Dice Coefficient loss.
 
@@ -48,6 +53,7 @@ def dice_loss(
             - ``'macro'``: Calculate the loss for each class separately and average the metrics across classes.
         eps: Scalar to enforce numerical stabiliy.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Return:
         One-element tensor of the computed loss.
@@ -76,9 +82,23 @@ def dice_loss(
     # compute softmax over the classes axis
     pred_soft: Tensor = pred.softmax(dim=1)
 
+    # check if all pixels are valid
+    valid_idx = target != ignore_index
+    all_valid = valid_idx.all()
+    if not all_valid:
+        # map invalid pixels to a valid mask for one hot encoding
+        # will be mapped to zero after
+        target = target.where(valid_idx, target.new_zeros(1))
+
     # create the labels one hot tensor
     target_one_hot: Tensor = one_hot(target, num_classes=pred.shape[1], device=pred.device, dtype=pred.dtype)
 
+    # mask ignore pixels
+    if not all_valid:
+        valid_idx.unsqueeze_(1)
+        target_one_hot = target_one_hot * valid_idx
+        pred_soft = pred_soft * valid_idx
+
     # set dimensions for the appropriate averaging
     dims: tuple[int, ...] = (2, 3)
     if average == "micro":
@@ -138,6 +158,7 @@ class DiceLoss(nn.Module):
             - ``'macro'``: Calculate the loss for each class separately and average the metrics across classes.
         eps: Scalar to enforce numerical stabiliy.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Shape:
         - Pred: :math:`(N, C, H, W)` where C = number of classes.
@@ -153,11 +174,14 @@ class DiceLoss(nn.Module):
         >>> output.backward()
     """
 
-    def __init__(self, average: str = "micro", eps: float = 1e-8, weight: Optional[Tensor] = None) -> None:
+    def __init__(
+        self, average: str = "micro", eps: float = 1e-8, weight: Optional[Tensor] = None, ignore_index: int = -100
+    ) -> None:
         super().__init__()
         self.average = average
         self.eps = eps
         self.weight = weight
+        self.ignore_index = ignore_index
 
     def forward(self, pred: Tensor, target: Tensor) -> Tensor:
-        return dice_loss(pred, target, self.average, self.eps, self.weight)
+        return dice_loss(pred, target, self.average, self.eps, self.weight, self.ignore_index)
diff --git a/kornia/losses/focal.py b/kornia/losses/focal.py
index 1a5813dd6d..35e1a836d6 100644
--- a/kornia/losses/focal.py
+++ b/kornia/losses/focal.py
@@ -20,6 +20,7 @@ def focal_loss(
     gamma: float = 2.0,
     reduction: str = "none",
     weight: Optional[Tensor] = None,
+    ignore_index: int = -100,
 ) -> Tensor:
     r"""Criterion that computes Focal loss.
 
@@ -44,6 +45,7 @@ def focal_loss(
           the number of elements in the output, ``'sum'``: the output will be
           summed.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Return:
         the computed loss.
@@ -68,9 +70,22 @@ def focal_loss(
         f"pred and target must be in the same device. Got: {pred.device} and {target.device}",
     )
 
+    # check if all pixels are valid
+    valid_idx = target != ignore_index
+    all_valid = valid_idx.all()
+    if not all_valid:
+        # map invalid pixels to a valid mask for one hot encoding
+        # will be mapped to zero after
+        target = target.where(valid_idx, target.new_zeros(1))
+
     # create the labels one hot tensor
     target_one_hot: Tensor = one_hot(target, num_classes=pred.shape[1], device=pred.device, dtype=pred.dtype)
 
+    # mask ignore pixels
+    if not all_valid:
+        valid_idx.unsqueeze_(1)
+        target_one_hot = target_one_hot * valid_idx
+
     # compute softmax over the classes axis
     log_pred_soft: Tensor = pred.log_softmax(1)
 
@@ -130,6 +145,7 @@ class FocalLoss(nn.Module):
           the number of elements in the output, ``'sum'``: the output will be
           summed.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Shape:
         - Pred: :math:`(N, C, *)` where C = number of classes.
@@ -147,16 +163,22 @@ class FocalLoss(nn.Module):
     """
 
     def __init__(
-        self, alpha: Optional[float], gamma: float = 2.0, reduction: str = "none", weight: Optional[Tensor] = None
+        self,
+        alpha: Optional[float],
+        gamma: float = 2.0,
+        reduction: str = "none",
+        weight: Optional[Tensor] = None,
+        ignore_index: int = -100,
     ) -> None:
         super().__init__()
         self.alpha: Optional[float] = alpha
         self.gamma: float = gamma
         self.reduction: str = reduction
         self.weight: Optional[Tensor] = weight
+        self.ignore_index = ignore_index
 
     def forward(self, pred: Tensor, target: Tensor) -> Tensor:
-        return focal_loss(pred, target, self.alpha, self.gamma, self.reduction, self.weight)
+        return focal_loss(pred, target, self.alpha, self.gamma, self.reduction, self.weight, self.ignore_index)
 
 
 def binary_focal_loss_with_logits(
@@ -167,6 +189,7 @@ def binary_focal_loss_with_logits(
     reduction: str = "none",
     pos_weight: Optional[Tensor] = None,
     weight: Optional[Tensor] = None,
+    ignore_index: int = -100,
 ) -> Tensor:
     r"""Criterion that computes Binary Focal loss.
 
@@ -193,6 +216,7 @@ def binary_focal_loss_with_logits(
         pos_weight: a weight of positive examples with shape :math:`(num\_of\_classes,)`.
           It is possible to trade off recall and precision by adding weights to positive examples.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Returns:
         the computed loss.
@@ -216,6 +240,18 @@ def binary_focal_loss_with_logits(
     log_probs_pos: Tensor = nn.functional.logsigmoid(pred)
     log_probs_neg: Tensor = nn.functional.logsigmoid(-pred)
 
+    # check if all pixels are valid
+    valid_idx = target != ignore_index
+    all_valid = valid_idx.all()
+    if not all_valid:
+        # map invalid pixels to a valid class
+        # will be mapped to zero after
+        target = target.where(valid_idx, target.new_zeros(1))
+
+        #  mask ignore pixels
+        log_probs_neg = log_probs_neg * valid_idx
+        log_probs_pos = log_probs_pos * valid_idx
+
     pos_term: Tensor = -log_probs_neg.exp().pow(gamma) * target * log_probs_pos
     neg_term: Tensor = -log_probs_pos.exp().pow(gamma) * (1.0 - target) * log_probs_neg
     if alpha is not None:
@@ -287,6 +323,7 @@ class BinaryFocalLossWithLogits(nn.Module):
         pos_weight: a weight of positive examples with shape :math:`(num\_of\_classes,)`.
           It is possible to trade off recall and precision by adding weights to positive examples.
         weight: weights for classes with shape :math:`(num\_of\_classes,)`.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Shape:
         - Pred: :math:`(N, C, *)` where C = number of classes.
@@ -310,6 +347,7 @@ def __init__(
         reduction: str = "none",
         pos_weight: Optional[Tensor] = None,
         weight: Optional[Tensor] = None,
+        ignore_index: int = -100,
     ) -> None:
         super().__init__()
         self.alpha: Optional[float] = alpha
@@ -317,8 +355,9 @@ def __init__(
         self.reduction: str = reduction
         self.pos_weight: Optional[Tensor] = pos_weight
         self.weight: Optional[Tensor] = weight
+        self.ignore_index = ignore_index
 
     def forward(self, pred: Tensor, target: Tensor) -> Tensor:
         return binary_focal_loss_with_logits(
-            pred, target, self.alpha, self.gamma, self.reduction, self.pos_weight, self.weight
+            pred, target, self.alpha, self.gamma, self.reduction, self.pos_weight, self.weight, self.ignore_index
         )
diff --git a/kornia/losses/tversky.py b/kornia/losses/tversky.py
index cfbbf56a8a..805f015d64 100644
--- a/kornia/losses/tversky.py
+++ b/kornia/losses/tversky.py
@@ -11,7 +11,12 @@
 
 
 def tversky_loss(
-    pred: torch.Tensor, target: torch.Tensor, alpha: float, beta: float, eps: float = 1e-8
+    pred: torch.Tensor,
+    target: torch.Tensor,
+    alpha: float,
+    beta: float,
+    eps: float = 1e-8,
+    ignore_index: int = -100,
 ) -> torch.Tensor:
     r"""Criterion that computes Tversky Coefficient loss.
 
@@ -40,6 +45,7 @@ def tversky_loss(
         alpha: the first coefficient in the denominator.
         beta: the second coefficient in the denominator.
         eps: scalar for numerical stability.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Return:
         the computed loss.
@@ -66,9 +72,23 @@ def tversky_loss(
     # compute softmax over the classes axis
     pred_soft: torch.Tensor = F.softmax(pred, dim=1)
 
+    # check if all pixels are valid
+    valid_idx = target != ignore_index
+    all_valid = valid_idx.all()
+    if not all_valid:
+        # map invalid pixels to a valid mask for one hot encoding
+        # will be mapped to zero after
+        target = target.where(valid_idx, target.new_zeros(1))
+
     # create the labels one hot tensor
     target_one_hot: torch.Tensor = one_hot(target, num_classes=pred.shape[1], device=pred.device, dtype=pred.dtype)
 
+    # mask ignore pixels
+    if not all_valid:
+        valid_idx.unsqueeze_(1)
+        target_one_hot = target_one_hot * valid_idx
+        pred_soft = pred_soft * valid_idx
+
     # compute the actual dice score
     dims = (1, 2, 3)
     intersection = torch.sum(pred_soft * target_one_hot, dims)
@@ -107,6 +127,7 @@ class TverskyLoss(nn.Module):
         alpha: the first coefficient in the denominator.
         beta: the second coefficient in the denominator.
         eps: scalar for numerical stability.
+        ignore_index: labels with this value are ignored in the loss computation.
 
     Shape:
         - Pred: :math:`(N, C, H, W)` where C = number of classes.
@@ -122,11 +143,12 @@ class TverskyLoss(nn.Module):
         >>> output.backward()
     """
 
-    def __init__(self, alpha: float, beta: float, eps: float = 1e-8) -> None:
+    def __init__(self, alpha: float, beta: float, eps: float = 1e-8, ignore_index: int = -100) -> None:
         super().__init__()
         self.alpha: float = alpha
         self.beta: float = beta
         self.eps: float = eps
+        self.ignore_index = ignore_index
 
     def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return tversky_loss(pred, target, self.alpha, self.beta, self.eps)
+        return tversky_loss(pred, target, self.alpha, self.beta, self.eps, self.ignore_index)
diff --git a/tests/losses/test_dice.py b/tests/losses/test_dice.py
index 96d61c7626..599522104a 100644
--- a/tests/losses/test_dice.py
+++ b/tests/losses/test_dice.py
@@ -105,10 +105,28 @@ def test_averaging_macro(self, device, dtype):
         loss = criterion(logits, labels)
         self.assert_close(loss, expected_loss, rtol=1e-3, atol=1e-3)
 
+    @pytest.mark.parametrize("ignore_index", [-100, 255])
+    def test_ignore_index(self, device, dtype, ignore_index):
+        num_classes = 2
+        eps = 1e-8
+
+        logits = torch.zeros(2, num_classes, 1, 4, device=device, dtype=dtype)
+        logits[:, 0, :, 0] = 100.0
+        logits[:, 1, :, 1:] = 100.0
+        labels = torch.zeros(2, 1, 4, device=device, dtype=torch.int64)
+
+        labels[..., 2:] = ignore_index
+        expected_loss = torch.tensor([1.0 / 2.0], device=device, dtype=dtype).squeeze()
+        criterion = kornia.losses.DiceLoss(average="micro", eps=eps, ignore_index=ignore_index)
+        loss = criterion(logits, labels)
+        self.assert_close(loss, expected_loss, rtol=1e-3, atol=1e-3)
+
     def test_gradcheck(self, device):
         num_classes = 3
         logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=torch.float64)
         labels = torch.randint(0, num_classes, (2, 3, 2), device=device)
+        ignore = torch.rand(2, 3, 2, device=device) > 0.8
+        labels[ignore] = -100
         self.gradcheck(kornia.losses.dice_loss, (logits, labels), dtypes=[torch.float64, torch.int64])
 
     def test_dynamo(self, device, dtype, torch_optimizer):
diff --git a/tests/losses/test_focal_loss.py b/tests/losses/test_focal_loss.py
index fe737e7dd3..9d15b4fd74 100644
--- a/tests/losses/test_focal_loss.py
+++ b/tests/losses/test_focal_loss.py
@@ -63,6 +63,20 @@ def test_shape_pos_weight_weight(self, device, dtype, reduction, expected_shape,
         ).shape
         assert actual_shape == expected_shape
 
+    @pytest.mark.parametrize("reduction,expected_shape", [("none", (2, 3, 2)), ("mean", ()), ("sum", ())])
+    @pytest.mark.parametrize("ignore_index", [-100, 255])
+    def test_shape_ignore_index(self, device, dtype, reduction, expected_shape, ignore_index):
+        logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
+        labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
+
+        ignore = torch.rand(2, 3, 2, device=device) > 0.6
+        labels[ignore] = ignore_index
+
+        actual_shape = kornia.losses.binary_focal_loss_with_logits(
+            logits, labels, alpha=0.8, gamma=0.5, reduction=reduction, ignore_index=ignore_index
+        ).shape
+        assert actual_shape == expected_shape
+
     def test_dynamo(self, device, dtype, torch_optimizer):
         logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
@@ -83,6 +97,16 @@ def test_gradcheck(self, device):
         op = kornia.losses.binary_focal_loss_with_logits
         self.gradcheck(op, (logits, labels, *args))
 
+    def test_gradcheck_ignore_index(self, device):
+        logits = torch.rand(2, 3, 2, device=device, dtype=torch.float64)
+        labels = torch.rand(2, 3, 2, device=device, dtype=torch.float64)
+        ignore = torch.rand(2, 3, 2, device=device) > 0.8
+        labels[ignore] = -100
+
+        args = (0.25, 2.0)
+        op = kornia.losses.binary_focal_loss_with_logits
+        self.gradcheck(op, (logits, labels, *args), requires_grad=[True, False, False, False])
+
     def test_module(self, device, dtype):
         logits = torch.rand(2, 3, 2, dtype=dtype, device=device)
         labels = torch.rand(2, 3, 2, dtype=dtype, device=device)
@@ -137,6 +161,44 @@ def test_shape_weight(self, device, dtype, reduction, expected_shape, weight):
         ).shape
         assert actual_shape == expected_shape
 
+    @pytest.mark.parametrize("reduction,expected_shape", [("none", (2, 3, 3, 2)), ("mean", ()), ("sum", ())])
+    @pytest.mark.parametrize("ignore_index", [-100, 255])
+    def test_shape_ignore_index(self, device, dtype, reduction, expected_shape, ignore_index):
+        num_classes = 3
+        logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
+        labels = torch.randint(num_classes, (2, 3, 2), device=device)
+
+        ignore = torch.rand(2, 3, 2, device=device) > 0.6
+        labels[ignore] = ignore_index
+
+        actual_shape = kornia.losses.focal_loss(
+            logits, labels, alpha=0.8, gamma=0.5, reduction=reduction, ignore_index=ignore_index
+        ).shape
+        assert actual_shape == expected_shape
+
+    @pytest.mark.parametrize("ignore_index", [-100, 255])
+    def test_value_ignore_index(self, device, dtype, ignore_index):
+        num_classes = 3
+        logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=dtype)
+        labels = torch.randint(num_classes, (2, 3, 2), device=device)
+
+        ignore = torch.rand(2, 3, 2, device=device) > 0.6
+        labels[ignore] = ignore_index
+
+        labels_extra_class = labels.clone()
+        labels_extra_class[ignore] = num_classes
+        logits_extra_class = torch.cat([logits, logits.new_full((2, 1, 3, 2), float("-inf"))], dim=1)
+
+        expected_values = kornia.losses.focal_loss(
+            logits_extra_class, labels_extra_class, alpha=0.8, gamma=0.5, reduction="none"
+        )[:, :-1, ...]
+
+        actual_values = kornia.losses.focal_loss(
+            logits, labels, alpha=0.8, gamma=0.5, reduction="none", ignore_index=ignore_index
+        )
+
+        self.assert_close(actual_values, expected_values)
+
     def test_dynamo(self, device, dtype, torch_optimizer):
         num_classes = 3
         logits = torch.rand(2, num_classes, device=device, dtype=dtype)
@@ -154,6 +216,8 @@ def test_gradcheck(self, device):
         num_classes = 3
         logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=torch.float64)
         labels = torch.randint(num_classes, (2, 3, 2), device=device).long()
+        ignore = torch.rand(2, 3, 2, device=device) > 0.8
+        labels[ignore] = -100
 
         self.gradcheck(
             kornia.losses.focal_loss, (logits, labels, 0.25, 2.0), dtypes=[torch.float64, torch.int64, None, None]
diff --git a/tests/losses/test_tversky.py b/tests/losses/test_tversky.py
index 5cdc8b75a0..d8e45bb4d0 100644
--- a/tests/losses/test_tversky.py
+++ b/tests/losses/test_tversky.py
@@ -47,11 +47,28 @@ def test_all_zeros(self, device, dtype):
         loss = criterion(logits, labels)
         self.assert_close(loss, torch.zeros_like(loss), atol=1e-3, rtol=1e-3)
 
+    @pytest.mark.parametrize("ignore_index", [-100, 255])
+    def test_ignore_index(self, device, dtype, ignore_index):
+        num_classes = 2
+
+        logits = torch.zeros(2, num_classes, 1, 4, device=device, dtype=dtype)
+        logits[:, 0, :, 0] = 100.0
+        logits[:, 1, :, 1:] = 100.0
+        labels = torch.zeros(2, 1, 4, device=device, dtype=torch.int64)
+
+        labels[..., 2:] = ignore_index
+        expected_loss = torch.tensor([1.0 / 2.0], device=device, dtype=dtype).squeeze()
+        criterion = kornia.losses.TverskyLoss(alpha=0.5, beta=0.5, ignore_index=ignore_index)
+        loss = criterion(logits, labels)
+        self.assert_close(loss, expected_loss, rtol=1e-3, atol=1e-3)
+
     def test_gradcheck(self, device):
         num_classes = 3
         alpha, beta = 0.5, 0.5  # for tversky loss
         logits = torch.rand(2, num_classes, 3, 2, device=device, dtype=torch.float64)
         labels = torch.randint(0, num_classes, (2, 3, 2), device=device)
+        ignore = torch.rand(2, 3, 2, device=device) > 0.8
+        labels[ignore] = -100
 
         self.gradcheck(
             kornia.losses.tversky_loss, (logits, labels, alpha, beta), dtypes=[torch.float64, torch.int64, None, None]


Code review comments:
diff_hunk: @@ -20,6 +20,7 @@ def focal_loss(
     gamma: float = 2.0,
     reduction: str = "none",
     weight: Optional[Tensor] = None,
+    ignore_index: int = -100,
path: kornia/losses/focal.py
text: @user1:
can this be set as None too ?

@author:
I initially followed value and typing from pytorch, e.g.
https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html

considering that the kornia losses are pure python, I added a short circuit for None that saves some computation. I left the default value the same as in pytorch though. It is not a breaking change as -100 would crash if fed into one_hot. As such no existing code should have it as a valid label.

diff_hunk: @@ -66,9 +72,23 @@ def tversky_loss(
     # compute softmax over the classes axis
     pred_soft: torch.Tensor = F.softmax(pred, dim=1)
 
+    # check if all pixels are valid
+    valid_idx = target != ignore_index
path: kornia/losses/tversky.py
text: @user1:
maybe this can be wrapped into a reusable function across all the losses 

@author:
I pulled it out into a helper function, not sure where to place it though. Any suggestions appreciated.

@user1:
This sounds good