added num_classes parameter

Signed-off-by: ytl0623 <[email protected]>

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -34,6 +34,7 @@ class AsymmetricFocalTverskyLoss(_Loss):
     def __init__(
         self,
         to_onehot_y: bool = False,
+        num_classes: int = 2,
         delta: float = 0.7,
         gamma: float = 0.75,
         epsilon: float = 1e-7,
@@ -43,6 +44,7 @@ def __init__(
         """
         Args:
             to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
+            num_classes: number of classes. Defaults to 2.
             delta: weight of the background class (used in the Tversky index denominator). Defaults to 0.7.
             gamma: focal exponent value to down-weight easy foreground examples. Defaults to 0.75.
             epsilon: a small value to prevent division by zero. Defaults to 1e-7.
@@ -54,6 +56,7 @@ def __init__(
         """
         super().__init__(reduction=LossReduction(reduction).value)
         self.to_onehot_y = to_onehot_y
+        self.num_classes = num_classes
         self.delta = delta
         self.gamma = gamma
         self.epsilon = epsilon
@@ -72,7 +75,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             is_already_prob = True
 
             if y_true.shape[1] == 1:
-                y_true = one_hot(y_true, num_classes=2)
+                y_true = one_hot(y_true, num_classes=self.num_classes)
         else:
             is_already_prob = False
 
@@ -120,13 +123,13 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         # Apply reduction
         if self.reduction == LossReduction.MEAN.value:
-            return torch.mean(all_losses)
+            return torch.mean(torch.sum(all_losses, dim=1))
         elif self.reduction == LossReduction.SUM.value:
             return torch.sum(all_losses)
         elif self.reduction == LossReduction.NONE.value:
             return all_losses
         else:
-            return torch.mean(all_losses)
+            return torch.mean(torch.sum(all_losses, dim=1))
 
 
 class AsymmetricFocalLoss(_Loss):
@@ -143,6 +146,7 @@ class AsymmetricFocalLoss(_Loss):
     def __init__(
         self,
         to_onehot_y: bool = False,
+        num_classes: int = 2,
         delta: float = 0.7,
         gamma: float = 2.0,
         epsilon: float = 1e-7,
@@ -152,6 +156,7 @@ def __init__(
         """
         Args:
             to_onehot_y: whether to convert `y` into the one-hot format. Defaults to False.
+            num_classes: number of classes. Defaults to 2.
             delta: weight for the foreground classes. Defaults to 0.7.
             gamma: focusing parameter for the background class (to down-weight easy background examples). Defaults to 2.0.
             epsilon: a small value to prevent calculation errors. Defaults to 1e-7.
@@ -160,6 +165,7 @@ def __init__(
         """
         super().__init__(reduction=LossReduction(reduction).value)
         self.to_onehot_y = to_onehot_y
+        self.num_classes = num_classes
         self.delta = delta
         self.gamma = gamma
         self.epsilon = epsilon
@@ -172,12 +178,12 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_true: ground truth labels.
         """
 
-        if y_pred.shape[1] == 1 and not self.use_softmax:
+        if y_pred.shape[1] == 1:
             y_pred = torch.sigmoid(y_pred)
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
             is_already_prob = True
             if y_true.shape[1] == 1:
-                y_true = one_hot(y_true, num_classes=2)
+                y_true = one_hot(y_true, num_classes=self.num_classes)
         else:
             is_already_prob = False
 
@@ -270,13 +276,15 @@ def __init__(
         self.use_softmax = use_softmax
 
         self.asy_focal_loss = AsymmetricFocalLoss(
+            num_classes=self.num_classes,
             gamma=self.gamma,
             delta=self.delta,
             use_softmax=self.use_softmax,
             to_onehot_y=to_onehot_y,
             reduction=LossReduction.NONE,
         )
         self.asy_focal_tversky_loss = AsymmetricFocalTverskyLoss(
+            num_classes=self.num_classes,
             gamma=self.gamma,
             delta=self.delta,
             use_softmax=self.use_softmax,