fixes DiceCELoss for multichannel targets (#5292)

Fixes DiceCELoss for multichannel targets.
Currently if "target" (ground truth label) is provided as a multichannel
data (each channel is binary or float), then current DiceCELoss attempts
to convert it to 1-channel using argmax (which could be impossible with
overlapping labels). There is no need for argmax, since pytorch's cross
entropy can handle multi-channel targets already.


### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/`
folder.

Signed-off-by: myron <[email protected]>
Signed-off-by: monai-bot <[email protected]>
Co-authored-by: monai-bot <[email protected]>

Author: myron

monai/data/box_utils.py

@@ -657,7 +657,7 @@ def boxes_center_distance(
     center2 = box_centers(boxes2_t.to(COMPUTE_DTYPE))  # (M, spatial_dims)
 
     if euclidean:
-        dists = (center1[:, None] - center2[None]).pow(2).sum(-1).sqrt()
+        dists = (center1[:, None] - center2[None]).pow(2).sum(-1).sqrt()  # type: ignore
     else:
         # before sum: (N, M, spatial_dims)
         dists = (center1[:, None] - center2[None]).sum(-1)

monai/losses/dice.py

@@ -21,7 +21,7 @@
 from monai.losses.focal_loss import FocalLoss
 from monai.losses.spatial_mask import MaskedLoss
 from monai.networks import one_hot
-from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option
+from monai.utils import DiceCEReduction, LossReduction, Weight, look_up_option, pytorch_after
 
 
 class DiceLoss(_Loss):
@@ -692,6 +692,7 @@ def __init__(
             raise ValueError("lambda_ce should be no less than 0.0.")
         self.lambda_dice = lambda_dice
         self.lambda_ce = lambda_ce
+        self.old_pt_ver = not pytorch_after(1, 10)
 
     def ce(self, input: torch.Tensor, target: torch.Tensor):
         """
@@ -701,12 +702,16 @@ def ce(self, input: torch.Tensor, target: torch.Tensor):
 
         """
         n_pred_ch, n_target_ch = input.shape[1], target.shape[1]
-        if n_pred_ch == n_target_ch:
-            # target is in the one-hot format, convert to BH[WD] format to calculate ce loss
-            target = torch.argmax(target, dim=1)
-        else:
+        if n_pred_ch != n_target_ch and n_target_ch == 1:
             target = torch.squeeze(target, dim=1)
-        target = target.long()
+            target = target.long()
+        elif self.old_pt_ver:
+            warnings.warn(
+                f"Multichannel targets are not supported in this older Pytorch version {torch.__version__}. "
+                "Using argmax (as a workaround) to convert target to a single channel."
+            )
+            target = torch.argmax(target, dim=1)
+
         return self.cross_entropy(input, target)
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: