DiceBCELossWithLogits is expected to support multiclass (#110)

* Updated loss computation to handle multi-class logit conversion. (#109)

* Fixed label conversion order in loss calculation to ensure compatible data types. (#109)

* Ensured tensor contiguity in one-hot encoding conversion. (#109)

Author: ATATC

mipcandy/common/optim/loss.py

@@ -3,6 +3,7 @@
 import torch
 from torch import nn
 
+from mipcandy.data import convert_ids_to_logits
 from mipcandy.metrics import do_reduction, soft_dice_coefficient
 
 
@@ -33,6 +34,11 @@ def __init__(self, num_classes: int, *, lambda_bce: float = .5, lambda_soft_dice
         self.smooth: float = smooth
 
     def forward(self, masks: torch.Tensor, labels: torch.Tensor) -> tuple[torch.Tensor, dict[str, float]]:
+        if self.num_classes != 1 and labels.shape[1] == 1:
+            d = labels.ndim - 2
+            if d not in (1, 2, 3):
+                raise ValueError(f"Expected labels to be 1D, 2D, or 3D, got {d} spatial dimensions")
+            labels = convert_ids_to_logits(labels.int(), d, self.num_classes)
         labels = labels.float()
         bce = nn.functional.binary_cross_entropy_with_logits(masks, labels)
         masks = masks.sigmoid()

mipcandy/data/convertion.py

@@ -15,7 +15,7 @@ def convert_ids_to_logits(ids: torch.Tensor, d: Literal[1, 2, 3], num_classes: i
             ids = ids.squeeze(1)
         else:
             raise ValueError(f"`ids` should be {d} dimensional or {d + 1} dimensional with single channel")
-    return nn.functional.one_hot(ids.long(), num_classes).movedim(-1, 1).float()
+    return nn.functional.one_hot(ids.long(), num_classes).movedim(-1, 1).contiguous().float()
 
 
 def convert_logits_to_ids(logits: torch.Tensor, *, channel_dim: int = 1) -> torch.Tensor: