dice_similarity_coefficient_with_logits(): RuntimeError: Boolean value of Tensor with more than one value is ambiguous (#210)

* Fix ambiguity. (#209)

* Add `dice_similarity_coefficient_with_logits_clip` and update usages.

* Removed `min_percentage_per_class` from loss and metric functions. (#209)

Author: ATATC

mipcandy/__init__.py

@@ -8,9 +8,9 @@
 from mipcandy.layer import batch_int_multiply, batch_int_divide, LayerT, HasDevice, auto_device, WithPaddingModule, \
     WithNetwork
 from mipcandy.metrics import do_reduction, dice_similarity_coefficient_binary, \
-    dice_similarity_coefficient_multiclass, dice_similarity_coefficient_with_logits, soft_dice_coefficient, \
-    accuracy_binary, accuracy_multiclass, precision_binary, precision_multiclass, recall_binary, recall_multiclass, \
-    iou_binary, iou_multiclass
+    dice_similarity_coefficient_multiclass, dice_similarity_coefficient_with_logits, \
+    dice_similarity_coefficient_with_logits_clip, soft_dice_coefficient, accuracy_binary, accuracy_multiclass, \
+    precision_binary, precision_multiclass, recall_binary, recall_multiclass, iou_binary, iou_multiclass
 from mipcandy.presets import *
 from mipcandy.profiler import ProfilerFrame, Profiler
 from mipcandy.run import config

mipcandy/common/optim/loss.py

@@ -5,7 +5,7 @@
 
 from mipcandy.data import convert_ids_to_logits, convert_logits_to_ids
 from mipcandy.metrics import do_reduction, soft_dice_coefficient, dice_similarity_coefficient_binary, \
-    dice_similarity_coefficient_with_logits
+    dice_similarity_coefficient_with_logits_clip
 
 
 class FocalBCEWithLogits(nn.Module):
@@ -56,21 +56,19 @@ def forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Te
                 class_dice = dice_similarity_coefficient_binary(outputs == i, labels == i).item()
                 dice += class_dice
                 metrics[f"dice {i}"] = class_dice
-            metrics["dice"] = dice_similarity_coefficient_with_logits(
+            metrics["dice"] = dice_similarity_coefficient_with_logits_clip(
                 self.logitfy_no_grad(outputs), self.logitfy_no_grad(labels)
             ).item()
             return c, metrics
 
 
 class DiceCELossWithLogits(_SegmentationLoss):
     def __init__(self, num_classes: int, *, lambda_ce: float = 1, lambda_soft_dice: float = 1,
-                 smooth: float = 1e-5, include_background: bool = True,
-                 min_percentage_per_class: float | None = None) -> None:
+                 smooth: float = 1e-5, include_background: bool = True) -> None:
         super().__init__(num_classes, include_background)
         self.lambda_ce: float = lambda_ce
         self.lambda_soft_dice: float = lambda_soft_dice
         self.smooth: float = smooth
-        self.min_percentage_per_class: float | None = min_percentage_per_class
 
     def _forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
         ce = nn.functional.cross_entropy(outputs, labels[:, 0].long())
@@ -79,8 +77,7 @@ def _forward(self, outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.T
         if not self.include_background:
             outputs = outputs[:, 1:]
             labels = labels[:, 1:]
-        soft_dice = soft_dice_coefficient(outputs, labels, smooth=self.smooth,
-                                          min_percentage_per_class=self.min_percentage_per_class)
+        soft_dice = soft_dice_coefficient(outputs, labels, smooth=self.smooth)
         metrics = {"soft dice": soft_dice.item(), "ce loss": ce.item()}
         c = self.lambda_ce * ce + self.lambda_soft_dice * (1 - soft_dice)
         return c, metrics

mipcandy/metrics.py

@@ -67,21 +67,33 @@ def dice_similarity_coefficient_multiclass(outputs: torch.Tensor, labels: torch.
     return apply_multiclass_to_binary(dice_similarity_coefficient_binary, outputs, labels, num_classes, if_empty)
 
 
-def dice_similarity_coefficient_with_logits(outputs: torch.Tensor, labels: torch.Tensor, *,
-                                            if_empty: float = 1) -> torch.Tensor:
+def _dice_with_logits(outputs: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     _args_check(outputs, labels, dtype=torch.float)
     axes = tuple(range(2, outputs.ndim))
     tp = (outputs * labels).sum(axes)
     fp = (outputs * (1 - labels)).sum(axes)
     fn = ((1 - outputs) * labels).sum(axes)
-    volume_sum = 2 * tp + fp + fn
-    if volume_sum == 0:
+    return tp, 2 * tp + fp + fn
+
+
+def dice_similarity_coefficient_with_logits(outputs: torch.Tensor, labels: torch.Tensor, *,
+                                            if_empty: float = 1) -> torch.Tensor:
+    tp, volume_sum = _dice_with_logits(outputs, labels)
+    if (volume_sum == 0).any():
         return torch.tensor(if_empty, dtype=torch.float)
-    return 2 * tp / volume_sum
+    dice = 2 * tp / volume_sum
+    return dice.mean()
+
+
+def dice_similarity_coefficient_with_logits_clip(outputs: torch.Tensor, labels: torch.Tensor, *,
+                                                 clip_min: float = 1e-8) -> torch.Tensor:
+    tp, volume_sum = _dice_with_logits(outputs, labels)
+    dice = 2 * tp / torch.clip(volume_sum, clip_min)
+    return dice.mean()
 
 
-def soft_dice_coefficient(outputs: torch.Tensor, labels: torch.Tensor, *, smooth: float = 1, batch_dice: bool = True,
-                          min_percentage_per_class: float | None = None) -> torch.Tensor:
+def soft_dice_coefficient(outputs: torch.Tensor, labels: torch.Tensor, *, smooth: float = 1, clip_min: float = 1e-8,
+                          batch_dice: bool = True) -> torch.Tensor:
     _args_check(outputs, labels)
     axes = tuple(range(2, outputs.ndim))
     if batch_dice:
@@ -93,16 +105,7 @@ def soft_dice_coefficient(outputs: torch.Tensor, labels: torch.Tensor, *, smooth
         intersection = intersection.sum(0)
         output_sum = output_sum.sum(0)
         label_sum = label_sum.sum(0)
-    dice = (2 * intersection + smooth) / (torch.clip(label_sum + output_sum + smooth, 1e-8))
-    if min_percentage_per_class:
-        total = label_sum.sum()
-        if total == 0:
-            return torch.tensor(1, device=outputs.device, dtype=outputs.dtype)
-        min_voxels = total * min_percentage_per_class
-        valid = label_sum >= min_voxels
-        if valid.any():
-            return dice[valid].mean()
-        return torch.tensor(1, device=outputs.device, dtype=outputs.dtype)
+    dice = (2 * intersection + smooth) / torch.clip(label_sum + output_sum + smooth, clip_min)
     return dice.mean()
 
 

mipcandy/presets/segmentation.py

@@ -94,10 +94,9 @@ def build_criterion(self) -> nn.Module:
         if self.num_classes < 2:
             if not self.include_background:
                 raise ValueError("Binary segmentation models must include background class")
-            loss = DiceBCELossWithLogits(min_percentage_per_class=1e-5)
+            loss = DiceBCELossWithLogits()
         else:
-            loss = DiceCELossWithLogits(self.num_classes, include_background=self.include_background,
-                                        min_percentage_per_class=1e-5)
+            loss = DiceCELossWithLogits(self.num_classes, include_background=self.include_background)
         if self.deep_supervision:
             if not self.deep_supervision_weights and self.deep_supervision_scales:
                 weights = np.array([1 / (2 ** i) for i in range(len(self.deep_supervision_scales))])