add implementation of TverskyLoss and TverskyLossFocall (#405)

* add implementation of TverskyLoss and TverskyLossFocall
* add tests for TverskyLoss

Author: loopdigga96

segmentation_models_pytorch/losses/__init__.py

@@ -6,3 +6,4 @@
 from .lovasz import LovaszLoss
 from .soft_bce import SoftBCEWithLogitsLoss
 from .soft_ce import SoftCrossEntropyLoss
+from .tversky import TverskyLoss, TverskyLossFocal

segmentation_models_pytorch/losses/_functional.py

@@ -6,7 +6,6 @@
 import torch
 import torch.nn.functional as F
 
-
 __all__ = [
     "focal_loss_with_logits",
     "softmax_focal_loss_with_logits",
@@ -35,14 +34,14 @@ def to_tensor(x, dtype=None) -> torch.Tensor:
 
 
 def focal_loss_with_logits(
-    output: torch.Tensor,
-    target: torch.Tensor,
-    gamma: float = 2.0,
-    alpha: Optional[float] = 0.25,
-    reduction: str = "mean",
-    normalized: bool = False,
-    reduced_threshold: Optional[float] = None,
-    eps: float = 1e-6,
+        output: torch.Tensor,
+        target: torch.Tensor,
+        gamma: float = 2.0,
+        alpha: Optional[float] = 0.25,
+        reduction: str = "mean",
+        normalized: bool = False,
+        reduced_threshold: Optional[float] = None,
+        eps: float = 1e-6,
 ) -> torch.Tensor:
     """Compute binary focal loss between target and output logits.
     See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.
@@ -98,13 +97,13 @@ def focal_loss_with_logits(
 
 
 def softmax_focal_loss_with_logits(
-    output: torch.Tensor,
-    target: torch.Tensor,
-    gamma: float = 2.0,
-    reduction="mean",
-    normalized=False,
-    reduced_threshold: Optional[float] = None,
-    eps: float = 1e-6,
+        output: torch.Tensor,
+        target: torch.Tensor,
+        gamma: float = 2.0,
+        reduction="mean",
+        normalized=False,
+        reduced_threshold: Optional[float] = None,
+        eps: float = 1e-6,
 ) -> torch.Tensor:
     """Softmax version of focal loss between target and output logits.
     See :class:`~pytorch_toolbelt.losses.FocalLoss` for details.
@@ -151,7 +150,7 @@ def softmax_focal_loss_with_logits(
 
 
 def soft_jaccard_score(
-    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
+        output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
 ) -> torch.Tensor:
     assert output.size() == target.size()
     if dims is not None:
@@ -167,7 +166,7 @@ def soft_jaccard_score(
 
 
 def soft_dice_score(
-    output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
+        output: torch.Tensor, target: torch.Tensor, smooth: float = 0.0, eps: float = 1e-7, dims=None
 ) -> torch.Tensor:
     assert output.size() == target.size()
     if dims is not None:
@@ -180,6 +179,22 @@ def soft_dice_score(
     return dice_score
 
 
+def soft_tversky_score(output: torch.Tensor, target: torch.Tensor, alpha: float, beta: float,
+                       smooth: float = 0.0, eps: float = 1e-7, dims=None) -> torch.Tensor:
+    assert output.size() == target.size()
+    if dims is not None:
+        intersection = torch.sum(output * target, dim=dims)  # TP
+        fp = torch.sum(output * (1. - target), dim=dims)
+        fn = torch.sum((1 - output) * target, dim=dims)
+    else:
+        intersection = torch.sum(output * target)  # TP
+        fp = torch.sum(output * (1. - target))
+        fn = torch.sum((1 - output) * target)
+
+    tversky_score = (intersection + smooth) / (intersection + alpha * fp + beta * fn + smooth).clamp_min(eps)
+    return tversky_score
+
+
 def wing_loss(output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"):
     """
     https://arxiv.org/pdf/1711.06753.pdf
@@ -211,7 +226,7 @@ def wing_loss(output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5
 
 
 def label_smoothed_nll_loss(
-    lprobs: torch.Tensor, target: torch.Tensor, epsilon: float, ignore_index=None, reduction="mean", dim=-1
+        lprobs: torch.Tensor, target: torch.Tensor, epsilon: float, ignore_index=None, reduction="mean", dim=-1
 ) -> torch.Tensor:
     """
     Source: https://github.com/pytorch/fairseq/blob/master/fairseq/criterions/label_smoothed_cross_entropy.py

segmentation_models_pytorch/losses/dice.py

@@ -12,14 +12,14 @@
 class DiceLoss(_Loss):
 
     def __init__(
-        self,
-        mode: str,
-        classes: Optional[List[int]] = None,
-        log_loss: bool = False,
-        from_logits: bool = True,
-        smooth: float = 0.0,
-        ignore_index: Optional[int] = None,
-        eps: float = 1e-7,
+            self,
+            mode: str,
+            classes: Optional[List[int]] = None,
+            log_loss: bool = False,
+            from_logits: bool = True,
+            smooth: float = 0.0,
+            ignore_index: Optional[int] = None,
+            eps: float = 1e-7,
     ):
         """Implementation of Dice loss for image segmentation task.
         It supports binary, multiclass and multilabel cases
@@ -104,7 +104,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                 y_pred = y_pred * mask
                 y_true = y_true * mask
 
-        scores = soft_dice_score(y_pred, y_true.type_as(y_pred), smooth=self.smooth, eps=self.eps, dims=dims)
+        scores = self.compute_score(y_pred, y_true.type_as(y_pred), smooth=self.smooth, eps=self.eps, dims=dims)
 
         if self.log_loss:
             loss = -torch.log(scores.clamp_min(self.eps))
@@ -122,4 +122,10 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         if self.classes is not None:
             loss = loss[self.classes]
 
+        return self.aggregate_loss(loss)
+
+    def aggregate_loss(self, loss):
         return loss.mean()
+
+    def compute_score(self, output, target, smooth=0.0, eps=1e-7, dims=None) -> torch.Tensor:
+        return soft_dice_score(output, target, smooth, eps, dims)

segmentation_models_pytorch/losses/tversky.py

@@ -0,0 +1,88 @@
+from typing import List
+
+import torch
+from ._functional import soft_tversky_score
+from .constants import BINARY_MODE, MULTICLASS_MODE, MULTILABEL_MODE
+from .dice import DiceLoss
+
+__all__ = ["TverskyLoss", "TverskyLossFocal"]
+
+
+class TverskyLoss(DiceLoss):
+    """
+    Implementation of Tversky loss for image segmentation task. Where TP and FP is weighted by alpha and beta params.
+    With alpha == beta == 0.5, this loss becomes equal DiceLoss.
+    It supports binary, multiclass and multilabel cases
+    """
+
+    def __init__(
+            self,
+            mode: str,
+            classes: List[int] = None,
+            log_loss=False,
+            from_logits=True,
+            smooth: float = 0.0,
+            ignore_index=None,
+            eps=1e-7,
+            alpha=0.5,
+            beta=0.5
+    ):
+        """
+        :param mode: Metric mode {'binary', 'multiclass', 'multilabel'}
+        :param classes: Optional list of classes that contribute in loss computation;
+        By default, all channels are included.
+        :param log_loss: If True, loss computed as `-log(jaccard)`; otherwise `1 - jaccard`
+        :param from_logits: If True assumes input is raw logits
+        :param smooth:
+        :param ignore_index: Label that indicates ignored pixels (does not contribute to loss)
+        :param eps: Small epsilon for numerical stability
+        :param alpha: Weight constant that penalize model for FPs (False Positives)
+        :param beta: Weight constant that penalize model for FNs (False Positives)
+        """
+        assert mode in {BINARY_MODE, MULTILABEL_MODE, MULTICLASS_MODE}
+        super().__init__(mode, classes, log_loss, from_logits, smooth, ignore_index, eps)
+        self.alpha = alpha
+        self.beta = beta
+
+    def compute_score(self, output, target, smooth=0.0, eps=1e-7, dims=None) -> torch.Tensor:
+        return soft_tversky_score(output, target, self.alpha, self.beta, smooth, eps, dims)
+
+
+class TverskyLossFocal(TverskyLoss):
+    """
+    A variant on the Tversky loss that also includes the gamma modifier from Focal Loss https://arxiv.org/abs/1708.02002
+    It supports binary, multiclass and multilabel cases
+    """
+
+    def __init__(
+            self,
+            mode: str,
+            classes: List[int] = None,
+            log_loss=False,
+            from_logits=True,
+            smooth: float = 0.0,
+            ignore_index=None,
+            eps=1e-7,
+            alpha=0.5,
+            beta=0.5,
+            gamma=1
+    ):
+        """
+        :param mode: Metric mode {'binary', 'multiclass', 'multilabel'}
+        :param classes: Optional list of classes that contribute in loss computation;
+        By default, all channels are included.
+        :param log_loss: If True, loss computed as `-log(jaccard)`; otherwise `1 - jaccard`
+        :param from_logits: If True assumes input is raw logits
+        :param smooth:
+        :param ignore_index: Label that indicates ignored pixels (does not contribute to loss)
+        :param eps: Small epsilon for numerical stability
+        :param alpha: Weight constant that penalize model for FPs (False Positives)
+        :param beta: Weight constant that penalize model for FNs (False Positives)
+        :param gamma: Constant that squares the error function
+        """
+        assert mode in {BINARY_MODE, MULTILABEL_MODE, MULTICLASS_MODE}
+        super().__init__(mode, classes, log_loss, from_logits, smooth, ignore_index, eps, alpha, beta)
+        self.gamma = gamma
+
+    def aggregate_loss(self, loss):
+        return loss.mean() ** self.gamma

tests/test_losses.py

@@ -2,7 +2,8 @@
 import torch
 import segmentation_models_pytorch as smp
 import segmentation_models_pytorch.losses._functional as F
-from segmentation_models_pytorch.losses import DiceLoss, JaccardLoss, SoftBCEWithLogitsLoss, SoftCrossEntropyLoss
+from segmentation_models_pytorch.losses import DiceLoss, JaccardLoss, SoftBCEWithLogitsLoss, SoftCrossEntropyLoss, \
+    TverskyLoss, TverskyLossFocal
 
 
 def test_focal_loss_with_logits():
@@ -71,6 +72,21 @@ def test_soft_dice_score(y_true, y_pred, expected, eps):
     assert float(actual) == pytest.approx(expected, eps)
 
 
+@pytest.mark.parametrize(
+    ["y_true", "y_pred", "expected", "eps", "alpha", "beta"],
+    [
+        [[1, 1, 1, 1], [1, 1, 1, 1], 1.0, 1e-5, 0.5, 0.5],
+        [[0, 1, 1, 0], [0, 1, 1, 0], 1.0, 1e-5, 0.5, 0.5],
+        [[1, 1, 1, 1], [1, 1, 0, 0], 2.0 / 3.0, 1e-5, 0.5, 0.5],
+    ],
+)
+def test_soft_tversky_score(y_true, y_pred, expected, eps, alpha, beta):
+    y_true = torch.tensor(y_true, dtype=torch.float32)
+    y_pred = torch.tensor(y_pred, dtype=torch.float32)
+    actual = F.soft_tversky_score(y_pred, y_true, eps=eps, alpha=alpha, beta=beta)
+    assert float(actual) == pytest.approx(expected, eps)
+
+
 @torch.no_grad()
 def test_dice_loss_binary():
     eps = 1e-5
@@ -109,6 +125,45 @@ def test_dice_loss_binary():
     assert float(loss) == pytest.approx(1.0, abs=eps)
 
 
+@torch.no_grad()
+def test_tversky_loss_binary():
+    eps = 1e-5
+    # with alpha=0.5; beta=0.5 it is equal to DiceLoss
+    criterion = TverskyLoss(mode=smp.losses.BINARY_MODE, from_logits=False, alpha=0.5, beta=0.5)
+
+    # Ideal case
+    y_pred = torch.tensor([1.0, 1.0, 1.0]).view(1, 1, 1, -1)
+    y_true = torch.tensor(([1, 1, 1])).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(0.0, abs=eps)
+
+    y_pred = torch.tensor([1.0, 0.0, 1.0]).view(1, 1, 1, -1)
+    y_true = torch.tensor(([1, 0, 1])).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(0.0, abs=eps)
+
+    y_pred = torch.tensor([0.0, 0.0, 0.0]).view(1, 1, 1, -1)
+    y_true = torch.tensor(([0, 0, 0])).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(0.0, abs=eps)
+
+    # Worst case
+    y_pred = torch.tensor([1.0, 1.0, 1.0]).view(1, 1, -1)
+    y_true = torch.tensor([0, 0, 0]).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(0.0, abs=eps)
+
+    y_pred = torch.tensor([1.0, 0.0, 1.0]).view(1, 1, -1)
+    y_true = torch.tensor([0, 1, 0]).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(1.0, abs=eps)
+
+    y_pred = torch.tensor([0.0, 0.0, 0.0]).view(1, 1, -1)
+    y_true = torch.tensor([1, 1, 1]).view(1, 1, 1, -1)
+    loss = criterion(y_pred, y_true)
+    assert float(loss) == pytest.approx(1.0, abs=eps)
+
+
 @torch.no_grad()
 def test_binary_jaccard_loss():
     eps = 1e-5