fix: reduction=NONE

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

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -71,6 +71,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_pred = torch.sigmoid(y_pred)
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
             is_already_prob = True
+            # Expand y_true to match if it's single channel
             if y_true.shape[1] == 1:
                 y_true = one_hot(y_true, num_classes=2)
         else:
@@ -213,17 +214,15 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         # Concatenate losses
         all_ce = torch.cat([back_ce.unsqueeze(1), fore_ce], dim=1)
 
-        # Sum over classes (dim=1) to get total loss per pixel
-        total_loss = torch.sum(all_ce, dim=1)
-
         # Apply reduction
         if self.reduction == LossReduction.MEAN.value:
-            return torch.mean(total_loss)
+            return torch.mean(torch.sum(all_ce, dim=1))
         if self.reduction == LossReduction.SUM.value:
-            return torch.sum(total_loss)
+            return torch.sum(all_ce)
         if self.reduction == LossReduction.NONE.value:
-            return total_loss
-        return torch.mean(total_loss)
+            return all_ce
+
+        return torch.mean(torch.sum(all_ce, dim=1))
 
 
 class AsymmetricUnifiedFocalLoss(_Loss):
@@ -268,14 +267,14 @@ def __init__(
             delta=self.delta,
             use_softmax=self.use_softmax,
             to_onehot_y=to_onehot_y,
-            reduction=reduction,
+            reduction=LossReduction.NONE,
         )
         self.asy_focal_tversky_loss = AsymmetricFocalTverskyLoss(
             gamma=self.gamma,
             delta=self.delta,
             use_softmax=self.use_softmax,
             to_onehot_y=to_onehot_y,
-            reduction=reduction,
+            reduction=LossReduction.NONE,
         )
 
     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
@@ -293,6 +292,21 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         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
+        # Align Focal Loss to (B, C) by averaging over spatial dimensions
+        spatial_dims = list(range(2, len(asy_focal_loss.shape)))
+        focal_aligned = torch.mean(asy_focal_loss, dim=spatial_dims)
+
+        # Calculate weighted sum. Result shape: (B, C)
+        combined_loss = self.weight * focal_aligned + (1 - self.weight) * asy_focal_tversky_loss
+
+        loss: torch.Tensor
+        if self.reduction == LossReduction.MEAN.value:
+            loss = torch.mean(combined_loss)
+        elif self.reduction == LossReduction.SUM.value:
+            loss = torch.sum(combined_loss)
+        elif self.reduction == LossReduction.NONE.value:
+            loss = combined_loss
+        else:
+            loss = torch.mean(combined_loss)
 
         return loss