Nic-Ma: Hi @wyli ,

Thanks for your review.
I update the code according to your suggestions.
Could you please help review it again?
Thanks.

diff --git a/monai/losses/dice.py b/monai/losses/dice.py
index 10dc15ad..8e8ff20f 100644
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -23,63 +23,93 @@ from monai.utils.aliases import alias
 @alias("dice", "Dice")
 class DiceLoss(_Loss):
     """
-    Multiclass dice loss. Input logits `pred` (BNHW[D] where N is number of classes) is compared with ground truth
-    `ground' (B1HW[D]). Axis N of `pred` is expected to have logit predictions for each class rather than being image
-    channels, while the same axis of `ground` should be 1. If the N channel of `pred` is 1 binary dice loss will be
-    calculated. The `smooth` parameter is a value added to the intersection and union components of the inter-over-union
-    calculation to smooth results and prevent divide-by-0, this value should be small. The `include_background` class
-    attribute can be set to False for an instance of DiceLoss to exclude the first category (channel index 0) which is
-    by convention assumed to be background. If the non-background segmentations are small compared to the total image
-    size they can get overwhelmed by the signal from the background so excluding it in such cases helps convergence.
+    Compute average Dice loss between two tensors. It can support both multi-classes and multi-labels tasks.
+    Input logits `pred` (BNHW[D] where N is number of classes) is compared with ground truth `ground' (BNHW[D]).
+    Axis N of `pred` is expected to have logit predictions for each class rather than being image channels,
+    while the same axis of `ground` can be 1 or N(one-hot format). The `smooth` parameter is a value added to the
+    intersection and union components of the inter-over-union calculation to smooth results and prevent divide by 0,
+    this value should be small. The `include_background` class attribute can be set to False for an instance of
+    DiceLoss to exclude the first category (channel index 0) which is by convention assumed to be background.
+    If the non-background segmentations are small compared to the total image size they can get overwhelmed by
+    the signal from the background so excluding it in such cases helps convergence.
     """
 
-    def __init__(self, include_background=True, do_sigmoid=False, do_softmax=False):
+    def __init__(
+        self,
+        include_background=True,
+        to_onehot_y=False,
+        do_sigmoid=False,
+        do_softmax=False,
+        squared_pred=False,
+        jaccard=False
+    ):
         """
         Args:
             include_background (bool): If False channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y (bool): whether to convert `y` into the one-hot format. Defaults to False.
             do_sigmoid (bool): If True, apply a sigmoid function to the prediction.
             do_softmax (bool): If True, apply a softmax function to the prediction.
+            squared_pred (bool): use squared versions of targets and predictions in the denominator or not.
+            jaccard (bool): compute Jaccard Index (soft IoU) instead of dice or not.
+
         """
         super().__init__()
         self.include_background = include_background
+        self.to_onehot_y = to_onehot_y
         if do_sigmoid and do_softmax:
             raise ValueError('do_sigmoid=True and do_softmax=Ture are not compatible.')
         self.do_sigmoid = do_sigmoid
         self.do_softmax = do_softmax
+        self.squared_pred = squared_pred
+        self.jaccard = jaccard
 
     def forward(self, pred, ground, smooth=1e-5):
-        if ground.shape[1] != 1:
-            raise ValueError("Ground truth should have only a single channel, shape is " + str(ground.shape))
-
-        psum = pred.float()
+        """
+        Args:
+            pred (tensor): the shape should be BNH[WD].
+            ground (tensor): the shape should be BNH[WD].
+            smooth (float): a small constant to avoid nan.
+        """
         if self.do_sigmoid:
-            psum = psum.sigmoid()  # use sigmoid activation
-        if pred.shape[1] == 1:
+            pred = torch.sigmoid(pred)
+        n_pred_ch = pred.shape[1]
+        if n_pred_ch == 1:
             if self.do_softmax:
-                raise ValueError('do_softmax is not compatible with single channel prediction.')
+                warnings.warn('single channel prediction, `do_softmax=True` ignored.')
+            if self.to_onehot_y:
+                warnings.warn('single channel prediction, `to_onehot_y=True` ignored.')
             if not self.include_background:
                 warnings.warn('single channel prediction, `include_background=False` ignored.')
-            tsum = ground
-        else:  # multiclass dice loss
+        else:
             if self.do_softmax:
-                psum = torch.softmax(pred, 1)
-            tsum = one_hot(ground, pred.shape[1])  # B1HW(D) -> BNHW(D)
-            # exclude background category so that it doesn't overwhelm the other segmentations if they are small
+                pred = torch.softmax(pred, 1)
+            if self.to_onehot_y:
+                ground = one_hot(ground, n_pred_ch)
             if not self.include_background:
-                tsum = tsum[:, 1:]
-                psum = psum[:, 1:]
-        assert tsum.shape == psum.shape, ("Ground truth one-hot has differing shape (%r) from source (%r)" %
-                                          (tsum.shape, psum.shape))
+                # if skipping background, removing first channel
+                ground = ground[:, 1:]
+                pred = pred[:, 1:]
+                assert ground.shape == pred.shape, ('ground truth one-hot has differing shape (%r) from pred (%r)' %
+                                                    (ground.shape, pred.shape))
 
-        batchsize, n_classes = tsum.shape[:2]
-        tsum = tsum.float().view(batchsize, n_classes, -1)
-        psum = psum.view(batchsize, n_classes, -1)
+        # reducing only spatial dimensions (not batch nor channels)
+        reduce_axis = list(range(2, len(pred.shape)))
+        intersection = torch.sum(ground * pred, reduce_axis)
 
-        intersection = psum * tsum
-        sums = psum + tsum
+        if self.squared_pred:
+            ground = torch.pow(ground, 2)
+            pred = torch.pow(pred, 2)
 
-        score = (2.0 * intersection.sum(2) + smooth) / (sums.sum(2) + smooth)
-        return 1 - score.mean()
+        ground_o = torch.sum(ground, reduce_axis)
+        pred_o = torch.sum(pred, reduce_axis)
+
+        denominator = ground_o + pred_o
+
+        if self.jaccard:
+            denominator -= intersection
+
+        f = (2.0 * intersection + smooth) / (denominator + smooth)
+        return 1.0 - f.mean()  # final reduce_mean across batches and channels
 
 
 @alias("generalized_dice", "generalised_dice")
@@ -94,24 +124,33 @@ class GeneralizedDiceLoss(_Loss):
         https://github.com/NifTK/NiftyNet/blob/v0.6.0/niftynet/layer/loss_segmentation.py#L279
     """
 
-    def __init__(self, include_background=True, do_sigmoid=False, do_softmax=False, w_type='square'):
+    def __init__(
+        self,
+        include_background=True,
+        to_onehot_y=False,
+        do_sigmoid=False,
+        do_softmax=False,
+        w_type='square'
+    ):
         """
         Args:
             include_background (bool): If False channel index 0 (background category) is excluded from the calculation.
+            to_onehot_y (bool): whether to convert `y` into the one-hot format. Defaults to False.
             do_sigmoid (bool): If True, apply a sigmoid function to the prediction.
             do_softmax (bool): If True, apply a softmax function to the prediction.
             w_type ('square'|'simple'|'uniform'): type of function to transform ground truth volume to a weight factor.
+
         """
         super().__init__()
         self.include_background = include_background
+        self.to_onehot_y = to_onehot_y
         if do_sigmoid and do_softmax:
             raise ValueError('do_sigmoid=True and do_softmax=Ture are not compatible.')
         self.do_sigmoid = do_sigmoid
         self.do_softmax = do_softmax
-
         self.w_func = torch.ones_like
         if w_type == 'simple':
-            self.w_func = lambda x: torch.reciprocal(x)
+            self.w_func = torch.reciprocal
         elif w_type == 'square':
             self.w_func = lambda x: torch.reciprocal(x * x)
         else:
@@ -121,44 +160,45 @@ class GeneralizedDiceLoss(_Loss):
         """
         Args:
             pred (tensor): the shape should be BNH[WD].
-            ground (tensor): the shape should be B1H[WD].
+            ground (tensor): the shape should be BNH[WD].
             smooth (float): a small constant to avoid nan.
         """
-        if ground.shape[1] != 1:
-            raise ValueError("Ground truth should have only a single channel, shape is " + str(ground.shape))
-
-        psum = pred.float()
         if self.do_sigmoid:
-            psum = psum.sigmoid()  # use sigmoid activation
-        if pred.shape[1] == 1:
+            pred = torch.sigmoid(pred)
+        n_pred_ch = pred.shape[1]
+        if n_pred_ch == 1:
             if self.do_softmax:
-                raise ValueError('do_softmax is not compatible with single channel prediction.')
+                warnings.warn('single channel prediction, `do_softmax=True` ignored.')
+            if self.to_onehot_y:
+                warnings.warn('single channel prediction, `to_onehot_y=True` ignored.')
             if not self.include_background:
                 warnings.warn('single channel prediction, `include_background=False` ignored.')
-            tsum = ground
-        else:  # multiclass dice loss
+        else:
             if self.do_softmax:
-                psum = torch.softmax(pred, 1)
-            tsum = one_hot(ground, pred.shape[1])  # B1HW(D) -> BNHW(D)
-            # exclude background category so that it doesn't overwhelm the other segmentations if they are small
+                pred = torch.softmax(pred, 1)
+            if self.to_onehot_y:
+                ground = one_hot(ground, n_pred_ch)
             if not self.include_background:
-                tsum = tsum[:, 1:]
-                psum = psum[:, 1:]
-        assert tsum.shape == psum.shape, ("Ground truth one-hot has differing shape (%r) from source (%r)" %
-                                          (tsum.shape, psum.shape))
+                # if skipping background, removing first channel
+                ground = ground[:, 1:]
+                pred = pred[:, 1:]
+                assert ground.shape == pred.shape, ('ground truth one-hot has differing shape (%r) from pred (%r)' %
+                                                    (ground.shape, pred.shape))
+
+        # reducing only spatial dimensions (not batch nor channels)
+        reduce_axis = list(range(2, len(pred.shape)))
+        intersection = torch.sum(ground * pred, reduce_axis)
 
-        batchsize, n_classes = tsum.shape[:2]
-        tsum = tsum.float().view(batchsize, n_classes, -1)
-        psum = psum.view(batchsize, n_classes, -1)
+        ground_o = torch.sum(ground, reduce_axis)
+        pred_o = torch.sum(pred, reduce_axis)
 
-        intersection = psum * tsum
-        sums = psum + tsum
+        denominator = ground_o + pred_o
 
-        w = self.w_func(tsum.sum(2))
+        w = self.w_func(ground_o.float())
         for b in w:
             infs = torch.isinf(b)
             b[infs] = 0.0
             b[infs] = torch.max(b)
 
-        score = (2.0 * intersection.sum(2) * w + smooth) / (sums.sum(2) * w + smooth)
-        return 1 - score.mean()
+        f = (2.0 * intersection * w + smooth) / (denominator * w + smooth)
+        return 1.0 - f.mean()  # final reduce_mean across batches and channels
diff --git a/monai/networks/utils.py b/monai/networks/utils.py
index 628e4ea7..9d60cfc1 100644
--- a/monai/networks/utils.py
+++ b/monai/networks/utils.py
@@ -46,12 +46,22 @@ def slice_channels(tensor, *slicevals):
     return tensor[slices]
 
 
-def predict_segmentation(logits):
+def predict_segmentation(logits, mutually_exclusive=False, threshold=0):
     """
-    Given the logits from a network, computing the segmentation by thresholding all values above 0 if `logits` has one
-    channel, or computing the `argmax` along the channel axis otherwise, logits has shape `BCHW[D]`
+    Given the logits from a network, computing the segmentation by thresholding all values above 0
+    if multi-labels task, computing the `argmax` along the channel axis if multi-classes task,
+    logits has shape `BCHW[D]`.
+
+    Args:
+        logits (Tensor): raw data of model output.
+        mutually_exclusive (bool): if True, `logits` will be converted into a binary matrix using
+            a combination of argmax, which is suitable for multi-classes task. Defaults to False.
+        threshold (float): thresholding the prediction values if multi-labels task.
     """
-    if logits.shape[1] == 1:
-        return (logits >= 0).int()  # for binary segmentation threshold on channel 0
+    if not mutually_exclusive:
+        return (logits >= threshold).int()
     else:
-        return logits.argmax(1).unsqueeze(1)  # take the index of the max value along dimension 1
+        if logits.shape[1] == 1:
+            warnings.warn('single channel prediction, `mutually_exclusive=True` ignored, use threshold instead.')
+            return (logits >= threshold).int()
+        return logits.argmax(1, keepdim=True)

dice loss support multi-labels task
**Is your feature request related to a problem? Please describe.**
The current implementation of DiceLoss can only support Binary and Multi-classes tasks.
Need to also support Multi-labels task.

diff --git a/tests/test_dice_loss.py b/tests/test_dice_loss.py
index e937185f..c0f12332 100644
--- a/tests/test_dice_loss.py
+++ b/tests/test_dice_loss.py
@@ -45,6 +45,7 @@ TEST_CASE_2 = [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
 TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
@@ -57,6 +58,7 @@ TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_sigmoid': True,
     },
     {
@@ -70,6 +72,7 @@ TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_softmax': True,
     },
     {
@@ -93,10 +96,39 @@ TEST_CASE_6 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
     0.307576,
 ]
 
+TEST_CASE_7 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'squared_pred': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-5,
+    },
+    0.178337,
+]
+
+TEST_CASE_8 = [  # shape: (1, 1, 2, 2), (1, 1, 2, 2)
+    {
+        'include_background': True,
+        'do_sigmoid': True,
+        'jaccard': True,
+    },
+    {
+        'pred': torch.tensor([[[[1., -1.], [-1., 1.]]]]),
+        'ground': torch.tensor([[[[1., 0.], [1., 1.]]]]),
+        'smooth': 1e-5,
+    },
+    -0.059094,
+]
+
 
 class TestDiceLoss(unittest.TestCase):
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4,
+                           TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8])
     def test_shape(self, input_param, input_data, expected_val):
         result = DiceLoss(**input_param).forward(**input_data)
         self.assertAlmostEqual(result.item(), expected_val, places=5)
diff --git a/tests/test_generalized_dice_loss.py b/tests/test_generalized_dice_loss.py
index b2ce9616..a3b328d1 100644
--- a/tests/test_generalized_dice_loss.py
+++ b/tests/test_generalized_dice_loss.py
@@ -45,6 +45,7 @@ TEST_CASE_1 = [  # shape: (2, 1, 2, 2), (2, 1, 2, 2)
 TEST_CASE_2 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),
@@ -57,6 +58,7 @@ TEST_CASE_2 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_sigmoid': True,
     },
     {
@@ -70,6 +72,7 @@ TEST_CASE_3 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': True,
+        'to_onehot_y': True,
         'do_softmax': True,
     },
     {
@@ -83,6 +86,7 @@ TEST_CASE_4 = [  # shape: (2, 2, 3), (2, 1, 3)
 TEST_CASE_5 = [  # shape: (2, 2, 3), (2, 1, 3)
     {
         'include_background': False,
+        'to_onehot_y': True,
     },
     {
         'pred': torch.tensor([[[1., 1., 0.], [0., 0., 1.]], [[1., 0., 1.], [0., 1., 0.]]]),

pip install -e .[all]
parse_log_pytest
requirements.txt
pytest black flake8 mypy pytype
apt-get update apt-get install -y gcc
3.6
requirements.txt
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
absl-py==1.4.0
attrs==22.2.0
black==22.8.0
cachetools==4.2.4
certifi==2021.5.30
charset-normalizer==2.0.12
click==8.0.4
coverage==6.2
cycler==0.11.0
dataclasses==0.8
decorator==4.4.2
flake8==5.0.4
google-auth==2.22.0
google-auth-oauthlib==0.4.6
grpcio==1.48.2
idna==3.10
imageio==2.15.0
importlab==0.8.1
importlib-metadata==4.8.3
iniconfig==1.1.1
kiwisolver==1.3.1
libcst==0.4.1
Markdown==3.3.7
matplotlib==3.3.4
mccabe==0.7.0
-e git+https://github.com/Project-MONAI/MONAI.git@62cd3784236421b1bea05fa7f925d60ad466ae8d#egg=monai
mypy==0.971
mypy-extensions==1.0.0
networkx==2.5.1
nibabel==3.2.2
ninja==1.11.1.1
numpy==1.19.5
oauthlib==3.2.2
packaging==21.3
parameterized==0.8.1
pathspec==0.9.0
Pillow==8.4.0
platformdirs==2.4.0
pluggy==1.0.0
protobuf==3.19.6
py==1.11.0
pyasn1==0.5.1
pyasn1-modules==0.3.0
pycodestyle==2.9.1
pyflakes==2.5.0
pyparsing==3.1.4
pytest==7.0.1
python-dateutil==2.9.0.post0
pytorch-ignite==0.3.0
pytype==2022.1.5
PyWavelets==1.1.1
PyYAML==6.0.1
requests==2.27.1
requests-oauthlib==2.0.0
rsa==4.9
scikit-image==0.17.2
scipy==1.5.4
six==1.17.0
tabulate==0.8.10
tensorboard==2.10.1
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tifffile==2020.9.3
toml==0.10.2
tomli==1.2.3
torch==1.10.2
typed-ast==1.5.5
typing-inspect==0.9.0
typing_extensions==4.1.1
urllib3==1.26.20
Werkzeug==2.0.3
zipp==3.6.0

name: MONAI
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - ca-certificates=2025.2.25=h06a4308_0
  - certifi=2021.5.30=py36h06a4308_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - libffi=3.3=he6710b0_2
  - libgcc-ng=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libstdcxx-ng=11.2.0=h1234567_1
  - ncurses=6.4=h6a678d5_0
  - openssl=1.1.1w=h7f8727e_0
  - pip=21.2.2=py36h06a4308_0
  - python=3.6.13=h12debd9_1
  - readline=8.2=h5eee18b_0
  - setuptools=58.0.4=py36h06a4308_0
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - wheel=0.37.1=pyhd3eb1b0_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - absl-py==1.4.0
      - attrs==22.2.0
      - black==22.8.0
      - cachetools==4.2.4
      - charset-normalizer==2.0.12
      - click==8.0.4
      - coverage==6.2
      - cycler==0.11.0
      - dataclasses==0.8
      - decorator==4.4.2
      - flake8==5.0.4
      - google-auth==2.22.0
      - google-auth-oauthlib==0.4.6
      - grpcio==1.48.2
      - idna==3.10
      - imageio==2.15.0
      - importlab==0.8.1
      - importlib-metadata==4.8.3
      - iniconfig==1.1.1
      - kiwisolver==1.3.1
      - libcst==0.4.1
      - markdown==3.3.7
      - matplotlib==3.3.4
      - mccabe==0.7.0
      - mypy==0.971
      - mypy-extensions==1.0.0
      - networkx==2.5.1
      - nibabel==3.2.2
      - ninja==1.11.1.1
      - numpy==1.19.5
      - oauthlib==3.2.2
      - packaging==21.3
      - parameterized==0.8.1
      - pathspec==0.9.0
      - pillow==8.4.0
      - platformdirs==2.4.0
      - pluggy==1.0.0
      - protobuf==3.19.6
      - py==1.11.0
      - pyasn1==0.5.1
      - pyasn1-modules==0.3.0
      - pycodestyle==2.9.1
      - pyflakes==2.5.0
      - pyparsing==3.1.4
      - pytest==7.0.1
      - python-dateutil==2.9.0.post0
      - pytorch-ignite==0.3.0
      - pytype==2022.1.5
      - pywavelets==1.1.1
      - pyyaml==6.0.1
      - requests==2.27.1
      - requests-oauthlib==2.0.0
      - rsa==4.9
      - scikit-image==0.17.2
      - scipy==1.5.4
      - six==1.17.0
      - tabulate==0.8.10
      - tensorboard==2.10.1
      - tensorboard-data-server==0.6.1
      - tensorboard-plugin-wit==1.8.1
      - tifffile==2020.9.3
      - toml==0.10.2
      - tomli==1.2.3
      - torch==1.10.2
      - typed-ast==1.5.5
      - typing-extensions==4.1.1
      - typing-inspect==0.9.0
      - urllib3==1.26.20
      - werkzeug==2.0.3
      - zipp==3.6.0
prefix: /opt/conda/envs/MONAI

tests/test_dice_loss.py::TestDiceLoss::test_shape_2 tests/test_dice_loss.py::TestDiceLoss::test_shape_3 tests/test_dice_loss.py::TestDiceLoss::test_shape_4 tests/test_dice_loss.py::TestDiceLoss::test_shape_6 tests/test_dice_loss.py::TestDiceLoss::test_shape_7 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_2 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_3 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_4 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_5
tests/test_dice_loss.py::TestDiceLoss::test_shape_0 tests/test_dice_loss.py::TestDiceLoss::test_shape_1 tests/test_dice_loss.py::TestDiceLoss::test_shape_5 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_0 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_1 tests/test_generalized_dice_loss.py::TestGeneralizedDiceLoss::test_shape_6
monai/losses/dice.py monai/networks/utils.py
monai/losses/dice.py monai/networks/utils.py