From c8acc62173bac41b55e2f69be2f7c38796ba0094 Mon Sep 17 00:00:00 2001 From: Josafat-Mattias Burmeister Date: Sun, 9 Jan 2022 14:48:04 +0100 Subject: [PATCH 1/3] fix differentiability of generalized Dice loss Signed-off-by: Josafat-Mattias Burmeister --- monai/losses/dice.py | 8 ++++---- tests/test_generalized_dice_loss.py | 9 +++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/monai/losses/dice.py b/monai/losses/dice.py index 7f2037ca54..86c145ba28 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -348,10 +348,10 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: denominator = ground_o + pred_o w = self.w_func(ground_o.float()) - for b in w: - infs = torch.isinf(b) - b[infs] = 0.0 - b[infs] = torch.max(b) + infs = torch.isinf(w) + w[infs] = 0.0 + max_values = torch.max(w, dim=1)[0].unsqueeze(dim=1) + w = w + infs * max_values final_reduce_dim = 0 if self.batch else 1 numer = 2.0 * (intersection * w).sum(final_reduce_dim, keepdim=True) + self.smooth_nr diff --git a/tests/test_generalized_dice_loss.py b/tests/test_generalized_dice_loss.py index f93878a683..7bac3f1bb8 100644 --- a/tests/test_generalized_dice_loss.py +++ b/tests/test_generalized_dice_loss.py @@ -173,6 +173,15 @@ def test_input_warnings(self): loss = GeneralizedDiceLoss(to_onehot_y=True) loss.forward(chn_input, chn_target) + def test_differentiability(self): + input = torch.ones((1, 1, 1, 3)) + target = torch.ones((1, 1, 1, 3)) + input.requires_grad = True + target.requires_grad = True + loss_module = GeneralizedDiceLoss() + loss = loss_module(input, target) + self.assertNotEqual(loss.grad_fn, None) + @SkipIfBeforePyTorchVersion((1, 7, 0)) def test_script(self): loss = GeneralizedDiceLoss() From d9a89727428a033beef872c5a640a66da64e8888 Mon Sep 17 00:00:00 2001 From: Josafat-Mattias Burmeister Date: Sun, 9 Jan 2022 18:24:49 +0100 Subject: [PATCH 2/3] fix weight clamping for the case that is True Signed-off-by: Josafat-Mattias Burmeister --- monai/losses/dice.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/monai/losses/dice.py b/monai/losses/dice.py index 86c145ba28..cfe7fb4b93 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -349,9 +349,12 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: w = self.w_func(ground_o.float()) infs = torch.isinf(w) - w[infs] = 0.0 - max_values = torch.max(w, dim=1)[0].unsqueeze(dim=1) - w = w + infs * max_values + if self.batch: + w[infs] = torch.max(w) + else: + w[infs] = 0.0 + max_values = torch.max(w, dim=1)[0].unsqueeze(dim=1) + w = w + infs * max_values final_reduce_dim = 0 if self.batch else 1 numer = 2.0 * (intersection * w).sum(final_reduce_dim, keepdim=True) + self.smooth_nr From f8369c3b2ffec5d6b559c70f2d3bf3634b35443a Mon Sep 17 00:00:00 2001 From: Josafat-Mattias Burmeister Date: Sun, 9 Jan 2022 21:13:56 +0100 Subject: [PATCH 3/3] add test case for batch mode --- tests/test_generalized_dice_loss.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/test_generalized_dice_loss.py b/tests/test_generalized_dice_loss.py index 7bac3f1bb8..fa301201e4 100644 --- a/tests/test_generalized_dice_loss.py +++ b/tests/test_generalized_dice_loss.py @@ -174,12 +174,23 @@ def test_input_warnings(self): loss.forward(chn_input, chn_target) def test_differentiability(self): - input = torch.ones((1, 1, 1, 3)) + prediction = torch.ones((1, 1, 1, 3)) target = torch.ones((1, 1, 1, 3)) - input.requires_grad = True + prediction.requires_grad = True target.requires_grad = True - loss_module = GeneralizedDiceLoss() - loss = loss_module(input, target) + + generalized_dice_loss = GeneralizedDiceLoss() + loss = generalized_dice_loss(prediction, target) + self.assertNotEqual(loss.grad_fn, None) + + def test_batch(self): + prediction = torch.zeros(2, 3, 3, 3) + target = torch.zeros(2, 3, 3, 3) + prediction.requires_grad = True + target.requires_grad = True + + generalized_dice_loss = GeneralizedDiceLoss(batch=True) + loss = generalized_dice_loss(prediction, target) self.assertNotEqual(loss.grad_fn, None) @SkipIfBeforePyTorchVersion((1, 7, 0))