Repository: Project-MONAI/MONAI

Issue:
Update the Generalized Wasserstein Dice loss
The implementation of the Generalized Wasserstein Dice loss that has been merged in MONAI does not correspond exactly to the loss proposed in https://link.springer.com/chapter/10.1007/978-3-319-75238-9_6

I will do the appropriate changes in the MONAI implementation and in the unit tests based on:
https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss

I have used the current implementation from my repo in the BraTS challenge 2020, as an additional test, and it performed very well (and better than the implementation currently in MONAI)
https://arxiv.org/abs/2011.01614
 




Discussion:
Hi @wyli and @ericspod ,

Could you please help double confirm the implementation of our Generalized Wasserstein Dice loss?

Thanks.


Hi @Nic-Ma, 

Thank you for your reply.

I have updated the code (including new unit tests) and I think it is ready to be merged.
I would like to know if I can add the generalized Wasserstein Dice loss to the documentation at the same time?

Thanks
I think you mean the doc-string of this loss function. Please also updated it in your PR.

Thanks.

If you have a PR ready for MONAI with the changes we'll look at that, that way it would be clearer how you think our implementation varies.
Hi @ericspod, thank you for offering to review the changes.

Just to clarify, the current implementation in MONAI is also based on my git repo https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss (please see #601).
In the meantime, I have done some changes in my repo and I would like the changes to be reflected in MONAI.

Sorry for the confusion.
No worries, we'll take a look when you're ready. If your code behaviour differs from what's published we need to make a note of it in the documentation of course.

Resolving patch:
diff --git a/monai/losses/__init__.py b/monai/losses/__init__.py
--- a/monai/losses/__init__.py
+++ b/monai/losses/__init__.py
@@ -17,6 +17,7 @@
     MaskedDiceLoss,
     dice,
     generalized_dice,
+    generalized_wasserstein_dice,
 )
 from .focal_loss import FocalLoss
 from .tversky import TverskyLoss
diff --git a/monai/losses/dice.py b/monai/losses/dice.py
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -360,70 +360,88 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
 class GeneralizedWassersteinDiceLoss(_Loss):
     """
-    Generalized Wasserstein Dice Loss [1] in PyTorch.
-    Compared to [1] we used a weighting method similar to the one
-    used in the generalized Dice Loss [2].
-
-    References:
-    ===========
-    [1] "Generalised Wasserstein Dice Score for Imbalanced Multi-class
-        Segmentation using Holistic Convolutional Networks",
-        Fidon L. et al. MICCAI BrainLes 2017.
-    [2] "Generalised dice overlap as a deep learning loss function
-        for highly unbalanced segmentations",
-        Sudre C., et al. MICCAI DLMIA 2017.
-
-    wasserstein_distance_map:
-    Compute the voxel-wise Wasserstein distance (eq. 6 in [1]) between the
-    flattened prediction and the flattened labels (ground_truth) with respect
-    to the distance matrix on the label space M.
-    References:
-    [1] "Generalised Wasserstein Dice Score for Imbalanced Multi-class
-        Segmentation using Holistic Convolutional Networks",
-        Fidon L. et al. MICCAI BrainLes 2017
-
-    compute_weights_generalized_true_positives:
-    Compute the weights \alpha_l of eq. 9 in [1] but using the weighting
-    method proposed in the generalized Dice Loss [2].
-    References:
-    [1] "Generalised Wasserstein Dice Score for Imbalanced Multi-class
-        Segmentation using Holistic Convolutional Networks",
-        Fidon L. et al. MICCAI BrainLes 2017
-    [2] "Generalised dice overlap as a deep learning loss function
-        for highly unbalanced segmentations." Sudre C., et al.
-        MICCAI DLMIA 2017.
+    Compute the generalized Wasserstein Dice Loss defined in:
+
+        Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
+        Segmentation using Holistic Convolutional Networks. BrainLes 2017.
+
+    Or its variant (use the option weighting_mode="GDL") defined in the Appendix of:
+
+        Tilborghs, S. et al. (2020) Comparative study of deep learning methods for the automatic
+        segmentation of lung, lesion and lesion type in CT scans of COVID-19 patients.
+        arXiv preprint arXiv:2007.15546
+
+    Adapted from:
+        https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss
     """
 
     def __init__(
         self,
         dist_matrix: Union[np.ndarray, torch.Tensor],
+        weighting_mode: str = "default",
         reduction: Union[LossReduction, str] = LossReduction.MEAN,
         smooth_nr: float = 1e-5,
         smooth_dr: float = 1e-5,
     ) -> None:
         """
         Args:
-            dist_matrix: 2d tensor or 2d numpy array; matrix of distances
-                between the classes. It must have dimension C x C where C is the
-                number of classes.
-            reduction: str; reduction mode.
+            dist_matrix: 2d tensor or 2d numpy array; matrix of distances between the classes.
+            It must have dimension C x C where C is the number of classes.
+            weighting_mode: {``"default"``, ``"GDL"``}
+                Specifies how to weight the class-specific sum of errors.
+                Default to ``"default"``.
+
+                - ``"default"``: (recommended) use the original weighting method as in:
+                    Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
+                    Segmentation using Holistic Convolutional Networks. BrainLes 2017.
+                - ``"GDL"``: use a GDL-like weighting method as in the Appendix of:
+                    Tilborghs, S. et al. (2020) Comparative study of deep learning methods for the automatic
+                    segmentation of lung, lesion and lesion type in CT scans of COVID-19 patients.
+                    arXiv preprint arXiv:2007.15546
+            reduction: {``"none"``, ``"mean"``, ``"sum"``}
+                Specifies the reduction to apply to the output. Defaults to ``"mean"``.
+
+                - ``"none"``: no reduction will be applied.
+                - ``"mean"``: the sum of the output will be divided by the number of elements in the output.
+                - ``"sum"``: the output will be summed.
             smooth_nr: a small constant added to the numerator to avoid zero.
             smooth_dr: a small constant added to the denominator to avoid nan.
 
         Raises:
             ValueError: When ``dist_matrix`` is not a square matrix.
 
+        Example:
+            .. code-block:: python
+
+                import torch
+                import numpy as np
+                from monai.losses import GeneralizedWassersteinDiceLoss
+
+                # Example with 3 classes (including the background: label 0).
+                # The distance between the background class (label 0) and the other classes is the maximum, equal to 1.
+                # The distance between class 1 and class 2 is 0.5.
+                dist_mat = np.array([[0.0, 1.0, 1.0], [1.0, 0.0, 0.5], [1.0, 0.5, 0.0]], dtype=np.float32)
+                wass_loss = GeneralizedWassersteinDiceLoss(dist_matrix=dist_mat)
+
+                pred_score = torch.tensor([[1000, 0, 0], [0, 1000, 0], [0, 0, 1000]], dtype=torch.float32)
+                grnd = torch.tensor([0, 1, 2], dtype=torch.int64)
+                wass_loss(pred_score, grnd)  # 0
+
         """
         super(GeneralizedWassersteinDiceLoss, self).__init__(reduction=LossReduction(reduction).value)
 
         if dist_matrix.shape[0] != dist_matrix.shape[1]:
             raise ValueError(f"dist_matrix must be C x C, got {dist_matrix.shape[0]} x {dist_matrix.shape[1]}.")
 
+        if weighting_mode not in ["default", "GDL"]:
+            raise ValueError("weighting_mode must be either 'default' or 'GDL, got %s." % weighting_mode)
+
         self.m = dist_matrix
         if isinstance(self.m, np.ndarray):
             self.m = torch.from_numpy(self.m)
         if torch.max(self.m) != 1:
             self.m = self.m / torch.max(self.m)
+        self.alpha_mode = weighting_mode
         self.num_classes = self.m.size(0)
         self.smooth_nr = float(smooth_nr)
         self.smooth_dr = float(smooth_dr)
@@ -445,18 +463,47 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         # Compute the Wasserstein distance map
         wass_dist_map = self.wasserstein_distance_map(probs, flat_target)
 
-        # Compute the generalised number of true positives
-        alpha = self.compute_weights_generalized_true_positives(flat_target)
-        true_pos = self.compute_generalized_true_positive(alpha, flat_target, wass_dist_map)
-        denom = self.compute_denominator(alpha, flat_target, wass_dist_map)
-
-        # Compute and return the final loss
+        # Compute the values of alpha to use
+        alpha = self._compute_alpha_generalized_true_positives(flat_target)
+
+        # Compute the nemerator and denominator of the generalized Wasserstein Dice loss
+        if self.alpha_mode == "GDL":
+            # use GDL-style alpha weights (i.e. normalize by the volume of each class)
+            # contrary to the original definition we also use alpha in the "generalized all error".
+            true_pos = self._compute_generalized_true_positive(alpha, flat_target, wass_dist_map)
+            denom = self._compute_denominator(alpha, flat_target, wass_dist_map)
+        else:  # default: as in the original paper
+            # (i.e. alpha=1 for all foreground classes and 0 for the background).
+            # Compute the generalised number of true positives
+            true_pos = self._compute_generalized_true_positive(alpha, flat_target, wass_dist_map)
+            all_error = torch.sum(wass_dist_map, dim=1)
+            denom = 2 * true_pos + all_error
+
+        # Compute the final loss
         wass_dice: torch.Tensor = (2.0 * true_pos + self.smooth_nr) / (denom + self.smooth_dr)
         wass_dice_loss: torch.Tensor = 1.0 - wass_dice
-        return wass_dice_loss.mean()
+
+        if self.reduction == LossReduction.MEAN.value:
+            wass_dice_loss = torch.mean(wass_dice_loss)  # the batch and channel average
+        elif self.reduction == LossReduction.SUM.value:
+            wass_dice_loss = torch.sum(wass_dice_loss)  # sum over the batch and channel dims
+        elif self.reduction == LossReduction.NONE.value:
+            pass  # returns [N, n_classes] losses
+        else:
+            raise ValueError(f'Unsupported reduction: {self.reduction}, available options are ["mean", "sum", "none"].')
+
+        return wass_dice_loss
 
     def wasserstein_distance_map(self, flat_proba: torch.Tensor, flat_target: torch.Tensor) -> torch.Tensor:
         """
+        Compute the voxel-wise Wasserstein distance between the
+        flattened prediction and the flattened labels (ground_truth) with respect
+        to the distance matrix on the label space M.
+        This corresponds to eq. 6 in:
+
+            Fidon L. et al. (2017) Generalised Wasserstein Dice Score for Imbalanced Multi-class
+            Segmentation using Holistic Convolutional Networks. BrainLes 2017.
+
         Args:
             flat_proba: the probabilities of input(predicted) tensor.
             flat_target: the target tensor.
@@ -485,7 +532,7 @@ def wasserstein_distance_map(self, flat_proba: torch.Tensor, flat_target: torch.
         wasserstein_map = torch.sum(wasserstein_map, dim=1)
         return wasserstein_map
 
-    def compute_generalized_true_positive(
+    def _compute_generalized_true_positive(
         self, alpha: torch.Tensor, flat_target: torch.Tensor, wasserstein_distance_map: torch.Tensor
     ) -> torch.Tensor:
         """
@@ -507,7 +554,7 @@ def compute_generalized_true_positive(
         )
         return generalized_true_pos
 
-    def compute_denominator(
+    def _compute_denominator(
         self, alpha: torch.Tensor, flat_target: torch.Tensor, wasserstein_distance_map: torch.Tensor
     ) -> torch.Tensor:
         """
@@ -529,14 +576,21 @@ def compute_denominator(
         )
         return generalized_true_pos
 
-    def compute_weights_generalized_true_positives(self, flat_target: torch.Tensor) -> torch.Tensor:
+    def _compute_alpha_generalized_true_positives(self, flat_target: torch.Tensor) -> torch.Tensor:
         """
         Args:
             flat_target: the target tensor.
         """
-        one_hot = F.one_hot(flat_target, num_classes=self.num_classes).permute(0, 2, 1).float()
-        volumes = torch.sum(one_hot, dim=2)
-        alpha: torch.Tensor = 1.0 / (volumes + 1.0)
+        alpha: torch.Tensor = torch.ones((flat_target.size(0), self.num_classes)).float().to(flat_target.device)
+        if self.alpha_mode == "GDL":  # GDL style
+            # Define alpha like in the generalized dice loss
+            # i.e. the inverse of the volume of each class.
+            one_hot = F.one_hot(flat_target, num_classes=self.num_classes).permute(0, 2, 1).float()
+            volumes = torch.sum(one_hot, dim=2)
+            alpha = 1.0 / (volumes + 1.0)
+        else:  # default, i.e. like in the original paper
+            # alpha weights are 0 for the background and 1 the other classes
+            alpha[:, 0] = 0.0
         return alpha
 
 


Test patch:
diff --git a/tests/test_generalized_wasserstein_dice_loss.py b/tests/test_generalized_wasserstein_dice_loss.py
--- a/tests/test_generalized_wasserstein_dice_loss.py
+++ b/tests/test_generalized_wasserstein_dice_loss.py
@@ -13,7 +13,9 @@
 
 import numpy as np
 import torch
+import torch.nn as nn
 import torch.nn.functional as F
+import torch.optim as optim
 
 from monai.losses import GeneralizedWassersteinDiceLoss
 
@@ -27,21 +29,53 @@ def test_bin_seg_2d(self):
         pred_very_good = 1000 * F.one_hot(target, num_classes=2).permute(0, 3, 1, 2).float()
         pred_very_poor = 1000 * F.one_hot(1 - target, num_classes=2).permute(0, 3, 1, 2).float()
 
-        # initialize the loss
-        loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]))
+        for weight_mode in ["default", "GDL"]:
+            # initialize the loss
+            loss = GeneralizedWassersteinDiceLoss(
+                dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode=weight_mode
+            )
+
+            # the loss for pred_very_good should be close to 0
+            loss_good = float(loss.forward(pred_very_good, target))
+            self.assertAlmostEqual(loss_good, 0.0, places=3)
+
+            # same test, but with target with a class dimension
+            target_4dim = target.unsqueeze(1)
+            loss_good = float(loss.forward(pred_very_good, target_4dim))
+            self.assertAlmostEqual(loss_good, 0.0, places=3)
+
+            # the loss for pred_very_poor should be close to 1
+            loss_poor = float(loss.forward(pred_very_poor, target))
+            self.assertAlmostEqual(loss_poor, 1.0, places=3)
+
+    def test_different_target_data_type(self):
+        """
+        Test if the loss is compatible with all the integer types
+        for the target segmentation.
+        """
+        # define 2d examples
+        target = torch.tensor([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]])
+
+        # add another dimension corresponding to the batch (batch size = 1 here)
+        target = target.unsqueeze(0)  # shape (1, H, W)
+        pred_very_good = 1000 * F.one_hot(target, num_classes=2).permute(0, 3, 1, 2).float()
 
-        # the loss for pred_very_good should be close to 0
-        loss_good = float(loss.forward(pred_very_good, target))
-        self.assertAlmostEqual(loss_good, 0.0, places=3)
+        target_uint8 = target.to(torch.uint8)
+        target_int8 = target.to(torch.int8)
+        target_short = target.short()
+        target_int = target.int()
+        target_long = target.long()
+        target_list = [target_uint8, target_int8, target_short, target_int, target_long]
 
-        # same test, but with target with a class dimension
-        target_4dim = target.unsqueeze(1)
-        loss_good = float(loss.forward(pred_very_good, target_4dim))
-        self.assertAlmostEqual(loss_good, 0.0, places=3)
+        for w_mode in ["default", "GDL"]:
+            # initialize the loss
+            loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode=w_mode)
 
-        # the loss for pred_very_poor should be close to 1
-        loss_poor = float(loss.forward(pred_very_poor, target))
-        self.assertAlmostEqual(loss_poor, 1.0, places=3)
+            # The test should pass irrespectively of the integer type used
+            for t in target_list:
+                # the loss for pred_very_good should be close to 0
+                loss_good = float(loss.forward(pred_very_good, t))
+                self.assertAlmostEqual(loss_good, 0.0, places=3)
 
     def test_empty_class_2d(self):
         num_classes = 2
@@ -52,16 +86,134 @@ def test_empty_class_2d(self):
         pred_very_good = 1000 * F.one_hot(target, num_classes=num_classes).permute(0, 3, 1, 2).float()
         pred_very_poor = 1000 * F.one_hot(1 - target, num_classes=num_classes).permute(0, 3, 1, 2).float()
 
-        # initialize the loss
-        loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]))
+        for w_mode in ["default", "GDL"]:
+            # initialize the loss
+            loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode=w_mode)
+
+            # loss for pred_very_good should be close to 0
+            loss_good = float(loss.forward(pred_very_good, target))
+            self.assertAlmostEqual(loss_good, 0.0, places=3)
+
+            # loss for pred_very_poor should be close to 1
+            loss_poor = float(loss.forward(pred_very_poor, target))
+            self.assertAlmostEqual(loss_poor, 1.0, places=3)
+
+    def test_bin_seg_3d(self):
+        # define 3d examples
+        target = torch.tensor(
+            [
+                # raw 0
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                # raw 1
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                # raw 2
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+            ]
+        )
 
-        # loss for pred_very_good should be close to 0
-        loss_good = float(loss.forward(pred_very_good, target))
-        self.assertAlmostEqual(loss_good, 0.0, places=3)
-
-        # loss for pred_very_poor should be close to 1
-        loss_poor = float(loss.forward(pred_very_poor, target))
-        self.assertAlmostEqual(loss_poor, 1.0, places=3)
+        # add another dimension corresponding to the batch (batch size = 1 here)
+        target = target.unsqueeze(0)  # shape (1, H, W, D)
+        pred_very_good = 1000 * F.one_hot(target, num_classes=2).permute(0, 4, 1, 2, 3).float()
+        pred_very_poor = 1000 * F.one_hot(1 - target, num_classes=2).permute(0, 4, 1, 2, 3).float()
+
+        for w_mode in ["default", "GDL"]:
+            # initialize the loss
+            loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode=w_mode)
+
+            # mean dice loss for pred_very_good should be close to 0
+            loss_good = float(loss.forward(pred_very_good, target))
+            self.assertAlmostEqual(loss_good, 0.0, places=3)
+
+            # mean dice loss for pred_very_poor should be close to 1
+            loss_poor = float(loss.forward(pred_very_poor, target))
+            self.assertAlmostEqual(loss_poor, 1.0, places=3)
+
+    def test_convergence(self):
+        """
+        The goal of this test is to assess if the gradient of the loss function
+        is correct by testing if we can train a one layer neural network
+        to segment one image.
+        We verify that the loss is decreasing in almost all SGD steps.
+        """
+        learning_rate = 0.01
+        max_iter = 50
+
+        # define a simple 3d example
+        target_seg = torch.tensor(
+            [
+                # raw 0
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                # raw 1
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+                # raw 2
+                [[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]],
+            ]
+        )
+        target_seg = torch.unsqueeze(target_seg, dim=0)
+        image = 12 * target_seg + 27  # dummy image to segment
+        image = image.float()
+        num_classes = 2
+        num_voxels = 3 * 4 * 4
+
+        # define a model with one layer
+        class OnelayerNet(nn.Module):
+            def __init__(self):
+                super(OnelayerNet, self).__init__()
+                self.layer = nn.Linear(num_voxels, num_voxels * num_classes)
+
+            def forward(self, x):
+                x = x.view(-1, num_voxels)
+                x = self.layer(x)
+                x = x.view(-1, num_classes, 3, 4, 4)
+                return x
+
+        for w_mode in ["default", "GDL"]:
+            # initialise the network
+            net = OnelayerNet()
+
+            # initialize the loss
+            loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode=w_mode)
+
+            # initialize an SGD
+            optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9)
+
+            # initial difference between pred and target
+            pred_start = torch.argmax(net(image), dim=1)
+            diff_start = torch.norm(pred_start.float() - target_seg.float())
+
+            loss_history = []
+            # train the network
+            for _ in range(max_iter):
+                # set the gradient to zero
+                optimizer.zero_grad()
+
+                # forward pass
+                output = net(image)
+                loss_val = loss(output, target_seg)
+
+                # backward pass
+                loss_val.backward()
+                optimizer.step()
+
+                # stats
+                loss_history.append(loss_val.item())
+
+            # difference between pred and target after training
+            pred_end = torch.argmax(net(image), dim=1)
+            diff_end = torch.norm(pred_end.float() - target_seg.float())
+
+            # count the number of SGD steps in which the loss decreases
+            num_decreasing_steps = 0
+            for i in range(len(loss_history) - 1):
+                if loss_history[i] > loss_history[i + 1]:
+                    num_decreasing_steps += 1
+            decreasing_steps_ratio = float(num_decreasing_steps) / (len(loss_history) - 1)
+
+            # verify that the loss is decreasing for sufficiently many SGD steps
+            self.assertTrue(decreasing_steps_ratio > 0.75)
+
+            # check that the predicted segmentation has improved
+            self.assertGreater(diff_start, diff_end)
 
 
 if __name__ == "__main__":