Repository: Project-MONAI/MONAI

Issue:
Unify the input of loss components
So far, the input format of focal loss is different from dice loss, in order to make it more convenient to use, I think it should be modified according to dice loss's input format.


Discussion:

Patch:
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -54,7 +54,7 @@ def __init__(
     ) -> None:
         """
         Args:
-            include_background: if False channel index 0 (background category) is excluded from the calculation.
+            include_background: if False, channel index 0 (background category) is excluded from the calculation.
             to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
             sigmoid: if True, apply a sigmoid function to the prediction.
             softmax: if True, apply a softmax function to the prediction.
@@ -101,10 +101,12 @@ def __init__(
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         Args:
-            input: the shape should be BNH[WD].
-            target: the shape should be BNH[WD].
+            input: the shape should be BNH[WD], where N is the number of classes.
+            target: the shape should be BNH[WD] or B1H[WD], where N is the number of classes.
 
         Raises:
+            AssertionError: When input and target (after one hot transform if setted)
+                have different shapes.
             ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].
 
         """
@@ -136,7 +138,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                 input = input[:, 1:]
 
         if target.shape != input.shape:
-            raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})")
+            raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
 
         # reducing only spatial dimensions (not batch nor channels)
         reduce_axis: List[int] = torch.arange(2, len(input.shape)).tolist()
diff --git a/monai/losses/focal_loss.py b/monai/losses/focal_loss.py
--- a/monai/losses/focal_loss.py
+++ b/monai/losses/focal_loss.py
@@ -9,16 +9,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union
+import warnings
+from typing import Optional, Sequence, Union
 
 import torch
 import torch.nn.functional as F
-from torch.nn.modules.loss import _WeightedLoss
+from torch.nn.modules.loss import _Loss
 
+from monai.networks import one_hot
 from monai.utils import LossReduction
 
 
-class FocalLoss(_WeightedLoss):
+class FocalLoss(_Loss):
     """
     Reimplementation of the Focal Loss described in:
 
@@ -29,15 +31,21 @@ class FocalLoss(_WeightedLoss):
 
     def __init__(
         self,
+        include_background: bool = True,
+        to_onehot_y: bool = False,
         gamma: float = 2.0,
-        weight: Optional[torch.Tensor] = None,
+        weight: Optional[Union[Sequence[float], float, int, torch.Tensor]] = None,
         reduction: Union[LossReduction, str] = LossReduction.MEAN,
     ) -> None:
         """
         Args:
+            include_background: if False, channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
             gamma: value of the exponent gamma in the definition of the Focal loss.
             weight: weights to apply to the voxels of each class. If None no weights are applied.
                 This corresponds to the weights `\alpha` in [1].
+                The input can be a single value (same weight for all classes), a sequence of values (the length
+                of the sequence should be the same as the number of classes).
             reduction: {``"none"``, ``"mean"``, ``"sum"``}
                 Specifies the reduction to apply to the output. Defaults to ``"mean"``.
 
@@ -57,80 +65,75 @@ def __init__(
                 fl(pred, grnd)
 
         """
-        super(FocalLoss, self).__init__(weight=weight, reduction=LossReduction(reduction).value)
+        super(FocalLoss, self).__init__(reduction=LossReduction(reduction).value)
+        self.include_background = include_background
+        self.to_onehot_y = to_onehot_y
         self.gamma = gamma
-        self.weight: Optional[torch.Tensor] = None
+        self.weight: Optional[Union[Sequence[float], float, int, torch.Tensor]] = weight
 
-    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         Args:
-            logits: the shape should be BCH[WD].
-                where C (greater than 1) is the number of classes.
-                Softmax over the logits is integrated in this module for improved numerical stability.
-            target: the shape should be B1H[WD] or BCH[WD].
-                If the target's shape is B1H[WD], the target that this loss expects should be a class index
-                in the range [0, C-1] where C is the number of classes.
+            input: the shape should be BNH[WD], where N is the number of classes.
+                The input should be the original logits since it will be transferred by
+                `F.log_softmax` in the forward function.
+            target: the shape should be BNH[WD] or B1H[WD], where N is the number of classes.
 
         Raises:
-            ValueError: When ``target`` ndim differs from ``logits``.
-            ValueError: When ``target`` channel is not 1 and ``target`` shape differs from ``logits``.
+            AssertionError: When input and target (after one hot transform if setted)
+                have different shapes.
             ValueError: When ``self.reduction`` is not one of ["mean", "sum", "none"].
 
         """
-        i = logits
+        n_pred_ch = input.shape[1]
+
+        if self.to_onehot_y:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
+            else:
+                target = one_hot(target, num_classes=n_pred_ch)
+
+        if not self.include_background:
+            if n_pred_ch == 1:
+                warnings.warn("single channel prediction, `include_background=False` ignored.")
+            else:
+                # if skipping background, removing first channel
+                target = target[:, 1:]
+                input = input[:, 1:]
+
+        if target.shape != input.shape:
+            raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
+
+        i = input
         t = target
 
-        if i.ndim != t.ndim:
-            raise ValueError(f"logits and target ndim must match, got logits={i.ndim} target={t.ndim}.")
-
-        if t.shape[1] != 1 and t.shape[1] != i.shape[1]:
-            raise ValueError(
-                "target must have one channel or have the same shape as the logits. "
-                "If it has one channel, it should be a class index in the range [0, C-1] "
-                f"where C is the number of classes inferred from 'logits': C={i.shape[1]}. "
-            )
-        if i.shape[1] == 1:
-            raise NotImplementedError("Single-channel predictions not supported.")
-
-        # Change the shape of logits and target to
-        # num_batch x num_class x num_voxels.
-        if i.dim() > 2:
-            i = i.view(i.size(0), i.size(1), -1)  # N,C,H,W => N,C,H*W
-            t = t.view(t.size(0), t.size(1), -1)  # N,1,H,W => N,1,H*W or N,C,H*W
-        else:  # Compatibility with classification.
-            i = i.unsqueeze(2)  # N,C => N,C,1
-            t = t.unsqueeze(2)  # N,1 => N,1,1 or N,C,1
-
-        # Compute the log proba (more stable numerically than softmax).
-        logpt = F.log_softmax(i, dim=1)  # N,C,H*W
-        # Keep only log proba values of the ground truth class for each voxel.
-        if target.shape[1] == 1:
-            logpt = logpt.gather(1, t.long())  # N,C,H*W => N,1,H*W
-            logpt = torch.squeeze(logpt, dim=1)  # N,1,H*W => N,H*W
+        # Change the shape of input and target to B x N x num_voxels.
+        i = i.view(i.size(0), i.size(1), -1)
+        t = t.view(t.size(0), t.size(1), -1)
 
+        # Compute the log proba.
+        logpt = F.log_softmax(i, dim=1)
         # Get the proba
-        pt = torch.exp(logpt)  # N,H*W or N,C,H*W
+        pt = torch.exp(logpt)  # B,H*W or B,N,H*W
 
         if self.weight is not None:
-            self.weight = self.weight.to(i)
+            class_weight: Optional[torch.Tensor] = None
+            if isinstance(self.weight, (float, int)):
+                class_weight = torch.as_tensor([self.weight] * i.size(1))
+            else:
+                class_weight = torch.as_tensor(self.weight)
+            class_weight = class_weight.to(i)
             # Convert the weight to a map in which each voxel
             # has the weight associated with the ground-truth label
             # associated with this voxel in target.
-            at = self.weight[None, :, None]  # C => 1,C,1
-            at = at.expand((t.size(0), -1, t.size(2)))  # 1,C,1 => N,C,H*W
-            if target.shape[1] == 1:
-                at = at.gather(1, t.long())  # selection of the weights  => N,1,H*W
-                at = torch.squeeze(at, dim=1)  # N,1,H*W => N,H*W
+            at = class_weight[None, :, None]  # N => 1,N,1
+            at = at.expand((t.size(0), -1, t.size(2)))  # 1,N,1 => B,N,H*W
             # Multiply the log proba by their weights.
             logpt = logpt * at
 
         # Compute the loss mini-batch.
         weight = torch.pow(-pt + 1.0, self.gamma)
-        if target.shape[1] == 1:
-            loss = torch.mean(-weight * logpt, dim=1)  # N
-        else:
-            loss = torch.mean(-weight * t * logpt, dim=-1)  # N,C
-
+        loss = torch.mean(-weight * t * logpt, dim=-1)
         if self.reduction == LossReduction.SUM.value:
             return loss.sum()
         if self.reduction == LossReduction.NONE.value:


Test patch:
diff --git a/tests/test_focal_loss.py b/tests/test_focal_loss.py
--- a/tests/test_focal_loss.py
+++ b/tests/test_focal_loss.py
@@ -16,13 +16,14 @@
 import torch.nn.functional as F
 
 from monai.losses import FocalLoss
+from monai.networks import one_hot
 from tests.utils import SkipIfBeforePyTorchVersion, test_script_save
 
 
 class TestFocalLoss(unittest.TestCase):
     def test_consistency_with_cross_entropy_2d(self):
         # For gamma=0 the focal loss reduces to the cross entropy loss
-        focal_loss = FocalLoss(gamma=0.0, reduction="mean")
+        focal_loss = FocalLoss(to_onehot_y=True, gamma=0.0, reduction="mean", weight=1.0)
         ce = nn.CrossEntropyLoss(reduction="mean")
         max_error = 0
         class_num = 10
@@ -36,7 +37,30 @@ def test_consistency_with_cross_entropy_2d(self):
                 x = x.cuda()
                 l = l.cuda()
             output0 = focal_loss(x, l)
-            output1 = ce(x, l[:, 0])
+            output1 = ce(x, l[:, 0]) / class_num
+            a = float(output0.cpu().detach())
+            b = float(output1.cpu().detach())
+            if abs(a - b) > max_error:
+                max_error = abs(a - b)
+        self.assertAlmostEqual(max_error, 0.0, places=3)
+
+    def test_consistency_with_cross_entropy_2d_onehot_label(self):
+        # For gamma=0 the focal loss reduces to the cross entropy loss
+        focal_loss = FocalLoss(to_onehot_y=False, gamma=0.0, reduction="mean")
+        ce = nn.CrossEntropyLoss(reduction="mean")
+        max_error = 0
+        class_num = 10
+        batch_size = 128
+        for _ in range(100):
+            # Create a random tensor of shape (batch_size, class_num, 8, 4)
+            x = torch.rand(batch_size, class_num, 8, 4, requires_grad=True)
+            # Create a random batch of classes
+            l = torch.randint(low=0, high=class_num, size=(batch_size, 1, 8, 4))
+            if torch.cuda.is_available():
+                x = x.cuda()
+                l = l.cuda()
+            output0 = focal_loss(x, one_hot(l, num_classes=class_num))
+            output1 = ce(x, l[:, 0]) / class_num
             a = float(output0.cpu().detach())
             b = float(output1.cpu().detach())
             if abs(a - b) > max_error:
@@ -45,7 +69,7 @@ def test_consistency_with_cross_entropy_2d(self):
 
     def test_consistency_with_cross_entropy_classification(self):
         # for gamma=0 the focal loss reduces to the cross entropy loss
-        focal_loss = FocalLoss(gamma=0.0, reduction="mean")
+        focal_loss = FocalLoss(to_onehot_y=True, gamma=0.0, reduction="mean")
         ce = nn.CrossEntropyLoss(reduction="mean")
         max_error = 0
         class_num = 10
@@ -60,7 +84,7 @@ def test_consistency_with_cross_entropy_classification(self):
                 x = x.cuda()
                 l = l.cuda()
             output0 = focal_loss(x, l)
-            output1 = ce(x, l[:, 0])
+            output1 = ce(x, l[:, 0]) / class_num
             a = float(output0.cpu().detach())
             b = float(output1.cpu().detach())
             if abs(a - b) > max_error:
@@ -75,7 +99,7 @@ def test_bin_seg_2d(self):
         pred_very_good = 1000 * F.one_hot(target, num_classes=2).permute(0, 3, 1, 2).float()
 
         # initialize the mean dice loss
-        loss = FocalLoss()
+        loss = FocalLoss(to_onehot_y=True)
 
         # focal loss for pred_very_good should be close to 0
         target = target.unsqueeze(1)  # shape (1, 1, H, W)
@@ -91,7 +115,7 @@ def test_empty_class_2d(self):
         pred_very_good = 1000 * F.one_hot(target, num_classes=num_classes).permute(0, 3, 1, 2).float()
 
         # initialize the mean dice loss
-        loss = FocalLoss()
+        loss = FocalLoss(to_onehot_y=True)
 
         # focal loss for pred_very_good should be close to 0
         target = target.unsqueeze(1)  # shape (1, 1, H, W)
@@ -106,7 +130,8 @@ def test_multi_class_seg_2d(self):
         target = target.unsqueeze(0)  # shape (1, H, W)
         pred_very_good = 1000 * F.one_hot(target, num_classes=num_classes).permute(0, 3, 1, 2).float()
         # initialize the mean dice loss
-        loss = FocalLoss()
+        loss = FocalLoss(to_onehot_y=True)
+        loss_onehot = FocalLoss(to_onehot_y=False)
 
         # focal loss for pred_very_good should be close to 0
         target_one_hot = F.one_hot(target, num_classes=num_classes).permute(0, 3, 1, 2)  # test one hot
@@ -115,7 +140,7 @@ def test_multi_class_seg_2d(self):
         focal_loss_good = float(loss(pred_very_good, target).cpu())
         self.assertAlmostEqual(focal_loss_good, 0.0, places=3)
 
-        focal_loss_good = float(loss(pred_very_good, target_one_hot).cpu())
+        focal_loss_good = float(loss_onehot(pred_very_good, target_one_hot).cpu())
         self.assertAlmostEqual(focal_loss_good, 0.0, places=3)
 
     def test_bin_seg_3d(self):
@@ -137,33 +162,30 @@ def test_bin_seg_3d(self):
         pred_very_good = 1000 * F.one_hot(target, num_classes=num_classes).permute(0, 4, 1, 2, 3).float()
 
         # initialize the mean dice loss
-        loss = FocalLoss()
+        loss = FocalLoss(to_onehot_y=True)
+        loss_onehot = FocalLoss(to_onehot_y=False)
 
         # focal loss for pred_very_good should be close to 0
         target = target.unsqueeze(1)  # shape (1, 1, H, W)
         focal_loss_good = float(loss(pred_very_good, target).cpu())
         self.assertAlmostEqual(focal_loss_good, 0.0, places=3)
 
-        focal_loss_good = float(loss(pred_very_good, target_one_hot).cpu())
+        focal_loss_good = float(loss_onehot(pred_very_good, target_one_hot).cpu())
         self.assertAlmostEqual(focal_loss_good, 0.0, places=3)
 
     def test_ill_opts(self):
         chn_input = torch.ones((1, 2, 3))
-        chn_target = torch.ones((1, 1, 3))
+        chn_target = torch.ones((1, 2, 3))
         with self.assertRaisesRegex(ValueError, ""):
             FocalLoss(reduction="unknown")(chn_input, chn_target)
-        with self.assertRaisesRegex(ValueError, ""):
-            FocalLoss(reduction=None)(chn_input, chn_target)
+        with self.assertRaisesRegex(TypeError, ""):
+            FocalLoss(other_act="tanh")(chn_input, chn_target)
 
     def test_ill_shape(self):
         chn_input = torch.ones((1, 2, 3))
         chn_target = torch.ones((1, 3))
-        with self.assertRaisesRegex(ValueError, ""):
+        with self.assertRaisesRegex(AssertionError, ""):
             FocalLoss(reduction="mean")(chn_input, chn_target)
-        chn_input = torch.ones((1, 1, 30))
-        chn_target = torch.ones((1, 1, 30))
-        with self.assertRaisesRegex(NotImplementedError, ""):
-            FocalLoss()(chn_input, chn_target)
 
     @SkipIfBeforePyTorchVersion((1, 7, 0))
     def test_script(self):
diff --git a/tests/test_seg_loss_integration.py b/tests/test_seg_loss_integration.py
--- a/tests/test_seg_loss_integration.py
+++ b/tests/test_seg_loss_integration.py
@@ -18,23 +18,29 @@
 from parameterized import parameterized
 
 from monai.losses import DiceLoss, FocalLoss, GeneralizedDiceLoss, TverskyLoss
+from monai.networks import one_hot
 
 TEST_CASES = [
     [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "smooth_nr": 1e-4, "smooth_dr": 1e-4}, {}],
     [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "smooth_nr": 0, "smooth_dr": 1e-3}, {}],
+    [DiceLoss, {"to_onehot_y": False, "squared_pred": True, "smooth_nr": 0, "smooth_dr": 1e-3}, {}],
     [DiceLoss, {"to_onehot_y": True, "squared_pred": True, "batch": True}, {}],
     [DiceLoss, {"to_onehot_y": True, "sigmoid": True}, {}],
     [DiceLoss, {"to_onehot_y": True, "softmax": True}, {}],
-    [FocalLoss, {"gamma": 1.5, "weight": torch.tensor([1, 2])}, {}],
-    [FocalLoss, {"gamma": 1.5}, {}],
+    [FocalLoss, {"to_onehot_y": True, "gamma": 1.5, "weight": torch.tensor([1, 2])}, {}],
+    [FocalLoss, {"to_onehot_y": False, "gamma": 1.5, "weight": [1, 2]}, {}],
+    [FocalLoss, {"to_onehot_y": False, "gamma": 1.5, "weight": 1.0}, {}],
+    [FocalLoss, {"to_onehot_y": True, "gamma": 1.5}, {}],
     [GeneralizedDiceLoss, {"to_onehot_y": True, "softmax": True}, {}],
     [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True}, {}],
     [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "simple"}, {}],
     [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "uniform"}, {}],
     [GeneralizedDiceLoss, {"to_onehot_y": True, "sigmoid": True, "w_type": "uniform", "batch": True}, {}],
+    [GeneralizedDiceLoss, {"to_onehot_y": False, "sigmoid": True, "w_type": "uniform", "batch": True}, {}],
     [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 0.8, "beta": 0.2}, {}],
     [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 0.8, "beta": 0.2, "batch": True}, {}],
     [TverskyLoss, {"to_onehot_y": True, "softmax": True, "alpha": 1.0, "beta": 0.0}, {}],
+    [TverskyLoss, {"to_onehot_y": False, "softmax": True, "alpha": 1.0, "beta": 0.0}, {}],
 ]
 
 
@@ -80,6 +86,8 @@ def test_convergence(self, loss_type, loss_args, forward_args):
         num_classes = 2
         num_voxels = 3 * 4 * 4
 
+        target_onehot = one_hot(target_seg, num_classes=num_classes)
+
         # define a one layer model
         class OnelayerNet(nn.Module):
             def __init__(self):
@@ -118,7 +126,10 @@ def forward(self, x):
             if init_output is None:
                 init_output = torch.argmax(output, 1).detach().cpu().numpy()
 
-            loss_val = loss(output, target_seg, **forward_args)
+            if loss_args["to_onehot_y"] is False:
+                loss_val = loss(output, target_onehot, **forward_args)
+            else:
+                loss_val = loss(output, target_seg, **forward_args)
 
             if iter_i % 10 == 0:
                 pred = torch.argmax(output, 1).detach().cpu().numpy()