Repository: Project-MONAI/MONAI

Issue:
Unified Focal Loss
**Is your feature request related to a problem? Please describe.**
Having the best loss available.

**Describe the solution you'd like**
Implement the Unified Focal Loss

**Additional context**
The Unified Focal Loss has been shown to be constantly better than quite a few other alternatives that are implemented din MONAI.

Paper: https://www.sciencedirect.com/science/article/pii/S0895611121001750
Original Tensorflow Code: https://github.com/mlyg/unified-focal-loss/blob/main/loss_functions.py



Discussion:
hi @Nic-Ma ,

I think it's an exciting idea. 

I am trying to go code with this proposal. but i'm a first-time contributor, I had to spend some time reading CONTRIBUTING.md.

Mason
Hi @JohnMasoner,

Internally a coworker of mine coded the Symmetrical Unified Focal Loss like this:
```
import torch
from monai.losses import TverskyLoss, FocalLoss

class Unified_FL(torch.nn.Module):
    def __init__(self, n_classes, delta = 0.6, gamma = 0.5, lambda_ = 0.5):
        super(Unified_FL, self).__init__()
        self.N = n_classes
        self.delta = delta
        self.gamma = gamma
        self.lambda_ = lambda_
        self.tversky = TverskyLoss(include_background=True, alpha = self.delta, beta = (1-self.delta),
                                   reduction="mean")
        self.focal = FocalLoss(include_background=True, gamma = 1-self.gamma, reduction="mean")
    def forward(self, x, y):
        return self.lambda_ * self.delta * self.focal(x, y) + (1 - self.lambda_) * self.tversky(x, y)
```

She said that for the Asymmetrical one the TverskyLoss needs to be modified.

Hope it helps ;) 

And thanks for picking it up <3 

Resolving patch:
diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py
--- a/monai/losses/__init__.py
+++ b/monai/losses/__init__.py
@@ -32,3 +32,4 @@
 from .multi_scale import MultiScaleLoss
 from .spatial_mask import MaskedLoss
 from .tversky import TverskyLoss
+from .unified_focal_loss import AsymmetricUnifiedFocalLoss
diff --git a/monai/losses/unified_focal_loss.py b/monai/losses/unified_focal_loss.py
new file mode 100644
--- /dev/null
+++ b/monai/losses/unified_focal_loss.py
@@ -0,0 +1,239 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Union
+
+import torch
+from torch.nn.modules.loss import _Loss
+
+from monai.networks import one_hot
+from monai.utils import LossReduction
+
+
+class AsymmetricFocalTverskyLoss(_Loss):
+    """
+    AsymmetricFocalTverskyLoss is a variant of FocalTverskyLoss, which attentions to the foreground class.
+
+    Actually, it's only supported for binary image segmentation now.
+
+    Reimplementation of the Asymmetric Focal Tversky Loss described in:
+
+    - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation",
+    Michael Yeung, Computerized Medical Imaging and Graphics
+    """
+
+    def __init__(
+        self,
+        to_onehot_y: bool = False,
+        delta: float = 0.7,
+        gamma: float = 0.75,
+        epsilon: float = 1e-7,
+        reduction: Union[LossReduction, str] = LossReduction.MEAN,
+    ) -> None:
+        """
+        Args:
+            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
+            delta : weight of the background. Defaults to 0.7.
+            gamma : value of the exponent gamma in the definition of the Focal loss  . Defaults to 0.75.
+            epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7.
+        """
+        super().__init__(reduction=LossReduction(reduction).value)
+        self.to_onehot_y = to_onehot_y
+        self.delta = delta
+        self.gamma = gamma
+        self.epsilon = epsilon
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        n_pred_ch = y_pred.shape[1]
+
+        if self.to_onehot_y:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
+            else:
+                y_true = one_hot(y_true, num_classes=n_pred_ch)
+
+        if y_true.shape != y_pred.shape:
+            raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
+
+        # clip the prediction to avoid NaN
+        y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
+        axis = list(range(2, len(y_pred.shape)))
+
+        # Calculate true positives (tp), false negatives (fn) and false positives (fp)
+        tp = torch.sum(y_true * y_pred, dim=axis)
+        fn = torch.sum(y_true * (1 - y_pred), dim=axis)
+        fp = torch.sum((1 - y_true) * y_pred, dim=axis)
+        dice_class = (tp + self.epsilon) / (tp + self.delta * fn + (1 - self.delta) * fp + self.epsilon)
+
+        # Calculate losses separately for each class, enhancing both classes
+        back_dice = 1 - dice_class[:, 0]
+        fore_dice = (1 - dice_class[:, 1]) * torch.pow(1 - dice_class[:, 1], -self.gamma)
+
+        # Average class scores
+        loss = torch.mean(torch.stack([back_dice, fore_dice], dim=-1))
+        return loss
+
+
+class AsymmetricFocalLoss(_Loss):
+    """
+    AsymmetricFocalLoss is a variant of FocalTverskyLoss, which attentions to the foreground class.
+
+    Actually, it's only supported for binary image segmentation now.
+
+    Reimplementation of the Asymmetric Focal Loss described in:
+
+    - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation",
+    Michael Yeung, Computerized Medical Imaging and Graphics
+    """
+
+    def __init__(
+        self,
+        to_onehot_y: bool = False,
+        delta: float = 0.7,
+        gamma: float = 2,
+        epsilon: float = 1e-7,
+        reduction: Union[LossReduction, str] = LossReduction.MEAN,
+    ):
+        """
+        Args:
+            to_onehot_y : whether to convert `y` into the one-hot format. Defaults to False.
+            delta : weight of the background. Defaults to 0.7.
+            gamma : value of the exponent gamma in the definition of the Focal loss  . Defaults to 0.75.
+            epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7.
+        """
+        super().__init__(reduction=LossReduction(reduction).value)
+        self.to_onehot_y = to_onehot_y
+        self.delta = delta
+        self.gamma = gamma
+        self.epsilon = epsilon
+
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        n_pred_ch = y_pred.shape[1]
+
+        if self.to_onehot_y:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
+            else:
+                y_true = one_hot(y_true, num_classes=n_pred_ch)
+
+        if y_true.shape != y_pred.shape:
+            raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
+
+        y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
+        cross_entropy = -y_true * torch.log(y_pred)
+
+        back_ce = torch.pow(1 - y_pred[:, 0], self.gamma) * cross_entropy[:, 0]
+        back_ce = (1 - self.delta) * back_ce
+
+        fore_ce = cross_entropy[:, 1]
+        fore_ce = self.delta * fore_ce
+
+        loss = torch.mean(torch.sum(torch.stack([back_ce, fore_ce], dim=1), dim=1))
+        return loss
+
+
+class AsymmetricUnifiedFocalLoss(_Loss):
+    """
+    AsymmetricUnifiedFocalLoss is a variant of Focal Loss.
+
+    Actually, it's only supported for binary image segmentation now
+
+    Reimplementation of the Asymmetric Unified Focal Tversky Loss described in:
+
+    - "Unified Focal Loss: Generalising Dice and Cross Entropy-based Losses to Handle Class Imbalanced Medical Image Segmentation",
+    Michael Yeung, Computerized Medical Imaging and Graphics
+    """
+
+    def __init__(
+        self,
+        to_onehot_y: bool = False,
+        num_classes: int = 2,
+        weight: float = 0.5,
+        gamma: float = 0.5,
+        delta: float = 0.7,
+        reduction: Union[LossReduction, str] = LossReduction.MEAN,
+    ):
+        """
+        Args:
+            to_onehot_y : whether to convert `y` into the one-hot format. Defaults to False.
+            num_classes : number of classes, it only supports 2 now. Defaults to 2.
+            delta : weight of the background. Defaults to 0.7.
+            gamma : value of the exponent gamma in the definition of the Focal loss. Defaults to 0.75.
+            epsilon : it defines a very small number each time. simmily smooth value. Defaults to 1e-7.
+            weight : weight for each loss function, if it's none it's 0.5. Defaults to None.
+
+        Example:
+            >>> import torch
+            >>> from monai.losses import AsymmetricUnifiedFocalLoss
+            >>> pred = torch.ones((1,1,32,32), dtype=torch.float32)
+            >>> grnd = torch.ones((1,1,32,32), dtype=torch.int64)
+            >>> fl = AsymmetricUnifiedFocalLoss(to_onehot_y=True)
+            >>> fl(pred, grnd)
+        """
+        super().__init__(reduction=LossReduction(reduction).value)
+        self.to_onehot_y = to_onehot_y
+        self.num_classes = num_classes
+        self.gamma = gamma
+        self.delta = delta
+        self.weight: float = weight
+        self.asy_focal_loss = AsymmetricFocalLoss(gamma=self.gamma, delta=self.delta)
+        self.asy_focal_tversky_loss = AsymmetricFocalTverskyLoss(gamma=self.gamma, delta=self.delta)
+
+    # TODO: Implement this  function to support multiple classes segmentation
+    def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            y_pred : the shape should be BNH[WD], where N is the number of classes.
+                It only supports binary segmentation.
+                The input should be the original logits since it will be transformed by
+                    a sigmoid in the forward function.
+            y_true : the shape should be BNH[WD], where N is the number of classes.
+                It only supports binary segmentation.
+
+        Raises:
+            ValueError: When input and target are different shape
+            ValueError: When len(y_pred.shape) != 4 and len(y_pred.shape) != 5
+            ValueError: When num_classes
+            ValueError: When the number of classes entered does not match the expected number
+        """
+        if y_pred.shape != y_true.shape:
+            raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
+
+        if len(y_pred.shape) != 4 and len(y_pred.shape) != 5:
+            raise ValueError(f"input shape must be 4 or 5, but got {y_pred.shape}")
+
+        if y_pred.shape[1] == 1:
+            y_pred = one_hot(y_pred, num_classes=self.num_classes)
+            y_true = one_hot(y_true, num_classes=self.num_classes)
+
+        if torch.max(y_true) != self.num_classes - 1:
+            raise ValueError(f"Pelase make sure the number of classes is {self.num_classes-1}")
+
+        n_pred_ch = y_pred.shape[1]
+        if self.to_onehot_y:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
+            else:
+                y_true = one_hot(y_true, num_classes=n_pred_ch)
+
+        asy_focal_loss = self.asy_focal_loss(y_pred, y_true)
+        asy_focal_tversky_loss = self.asy_focal_tversky_loss(y_pred, y_true)
+
+        loss: torch.Tensor = self.weight * asy_focal_loss + (1 - self.weight) * asy_focal_tversky_loss
+
+        if self.reduction == LossReduction.SUM.value:
+            return torch.sum(loss)  # sum over the batch and channel dims
+        if self.reduction == LossReduction.NONE.value:
+            return loss  # returns [N, num_classes] losses
+        if self.reduction == LossReduction.MEAN.value:
+            return torch.mean(loss)
+        raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].')


Test patch:
diff --git a/tests/test_unified_focal_loss.py b/tests/test_unified_focal_loss.py
new file mode 100644
--- /dev/null
+++ b/tests/test_unified_focal_loss.py
@@ -0,0 +1,63 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.losses import AsymmetricUnifiedFocalLoss
+
+TEST_CASES = [
+    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+        {
+            "y_pred": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+            "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+        },
+        0.0,
+    ],
+    [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
+        {
+            "y_pred": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+            "y_true": torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]),
+        },
+        0.0,
+    ],
+]
+
+
+class TestAsymmetricUnifiedFocalLoss(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_result(self, input_data, expected_val):
+        loss = AsymmetricUnifiedFocalLoss()
+        result = loss(**input_data)
+        np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+
+    def test_ill_shape(self):
+        loss = AsymmetricUnifiedFocalLoss()
+        with self.assertRaisesRegex(ValueError, ""):
+            loss(torch.ones((2, 2, 2)), torch.ones((2, 2, 2, 2)))
+
+    def test_with_cuda(self):
+        loss = AsymmetricUnifiedFocalLoss()
+        i = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]])
+        j = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]])
+        if torch.cuda.is_available():
+            i = i.cuda()
+            j = j.cuda()
+        output = loss(i, j)
+        print(output)
+        np.testing.assert_allclose(output.detach().cpu().numpy(), 0.0, atol=1e-4, rtol=1e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()