fixes DiceCELoss for multichannel targets (Project-MONAI#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 <amyronenko@nvidia.com> Signed-off-by: monai-bot <monai.miccai2019@gmail.com> Co-authored-by: monai-bot <monai.miccai2019@gmail.com>

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: