Generalized Dice Loss - Reduction="none" + batch=True

[feliksh]

Describe the bug
GeneralizedDiceLoss(..., reduction="none", batch=True) does not return the expected shape.

To Reproduce
Steps to reproduce the behavior:

Given tensors of shape:
  network_prediction -> B x C x H x W x D
  labels -> B x 1 x H x W x D
  with B=Batch_size, C=Channels/Classes, Height, Width, Depth

Correct behaviour, since I am asking for no reduction and batch=True:

loss_fun1 = DiceLoss(to_onehot_y=True, softmax=True, reduction="none", batch=True)
loss1 = loss_fun1(network_prediction, labels)

loss1 -> [C x 1 x 1 x 1]

While if you do the same with GeneralizedDiceLoss:

loss_fun2 = GeneralizedDiceLoss(to_onehot_y=True, softmax=True, reduction="none", batch=True)
loss2 = loss_fun1(network_prediction, labels)

loss2 -> 1 x 1 x 1 x 1

which is not correct, since again you are asking for no reduction and batch=True.

The relevant piece of code which seems to lose the correct shape is in
https://github.com/Project-MONAI/MONAI/blob/dev/monai/losses/dice.py#L360

final_reduce_dim = 0 if self.batch else 1
numer = 2.0 * (intersection * w).sum(final_reduce_dim, keepdim=True) + self.smooth_nr
denom = (denominator * w).sum(final_reduce_dim, keepdim=True) + self.smooth_dr
f: torch.Tensor = 1.0 - (numer / denom)

where numer and denom are collapsed to shape (1,) since both intersection,denominator and w are of shape (C,),
thus .sum(final_reduce_dim, keepdim=True) does collapse the shape (C,) into (1,).

I believe the solution would be to not let those tensors be of singleton shape (C,) but instead keep them at (1, C) or (C, 1)

Environment

Python 3.7.9,
Monai 0.9.1,
PyTorch 1.11.0

[KumoLiu]

Hi @feliksh, thanks for pointing it out, I just create a PR #5468 to fix it.

[feliksh]

That's fantastic, impressed by the quickness. Well done, thank you!