Deep supervision loss wrapper class (Project-MONAI#5338)

This adds a DeepSupervisionLoss wrapper class around the main loss function to accept a list of tensors returned from a deeply supervised networks. The final loss is computed as the sum of weighted losses for each of deep supervision levels (accounting for potential differences in shapes between targets and ds outputs) The wrapper class is designed to work with arbitrary existing loss,e.g. ``` loss = DiceCELoss(to_onehot_y=True, softmax=True) ds_loss = DeepSupervisionLoss(loss) ``` Whereas the existing loss accepts the input as a single Tensor, ds_loss accepts the input as a list of Tensors (for each output of a deeply supervised network). If only a simple Tensor input is provided, ds_loss behaves exactly the same as the underlying loss. I added unit tests too. ### 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>

Author: myron

monai/losses/__init__.py

@@ -26,6 +26,7 @@
  generalized_dice_focal,
  generalized_wasserstein_dice,
 )
+from .ds_loss import DeepSupervisionLoss
 from .focal_loss import FocalLoss
 from .giou_loss import BoxGIoULoss, giou
 from .image_dissimilarity import GlobalMutualInformationLoss, LocalNormalizedCrossCorrelationLoss

monai/losses/ds_loss.py

@@ -0,0 +1,81 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch.nn.modules.loss import _Loss
+
+from monai.utils import pytorch_after
+
+
+class DeepSupervisionLoss(_Loss):
+ """
+ Wrapper class around the main loss function to accept a list of tensors returned from a deeply
+ supervised networks. The final loss is computed as the sum of weighted losses for each of deep supervision levels.
+ """
+
+ def __init__(self, loss: _Loss, weight_mode: str = "exp", weights: Optional[List[float]] = None) -> None:
+ """
+ Args:
+ loss: main loss instance, e.g DiceLoss().
+ weight_mode: {``"same"``, ``"exp"``, ``"two"``}
+ Specifies the weights calculation for each image level. Defaults to ``"exp"``.
+ - ``"same"``: all weights are equal to 1.
+ - ``"exp"``: exponentially decreasing weights by a power of 2: 0, 0.5, 0.25, 0.125, etc .
+ - ``"two"``: equal smaller weights for lower levels: 1, 0.5, 0.5, 0.5, 0.5, etc
+ weights: a list of weights to apply to each deeply supervised sub-loss, if provided, this will be used
+ regardless of the weight_mode
+ """
+ super().__init__()
+ self.loss = loss
+ self.weight_mode = weight_mode
+ self.weights = weights
+ self.interp_mode = "nearest-exact" if pytorch_after(1, 11) else "nearest"
+
+ def get_weight(self, level: int = 0) -> float:
+ """
+ Calculates a weight constant for a given image scale level
+ """
+ weight = 1.0
+ if self.weights is not None and len(self.weights) > level:
+ weight = self.weights[level]
+ elif self.weight_mode == "same":
+ weight = 1.0
+ elif self.weight_mode == "exp":
+ weight = max(0.5**level, 0.0625)
+ elif self.weight_mode == "two":
+ weight = 1.0 if level == 0 else 0.5
+
+ return weight
+
+ def get_loss(self, input: torch.Tensor, target: torch.Tensor):
+ """
+ Calculates a loss output accounting for differences in shapes,
+ and downsizing targets if necessary (using nearest neighbor interpolation)
+ Generally downsizing occurs for all level, except for the first (level==0)
+ """
+ if input.shape[2:] != target.shape[2:]:
+ target = F.interpolate(target, size=input.shape[2:], mode=self.interp_mode)
+ return self.loss(input, target)
+
+ def forward(self, input: Union[torch.Tensor, List[torch.Tensor]], target: torch.Tensor):
+
+ if isinstance(input, (list, tuple)):
+ loss = torch.zeros(1, dtype=torch.float, device=target.device)
+ for l in range(len(input)):
+ loss += self.get_loss(input[l].float(), target) * self.get_weight(l)
+ return loss
+ return self.loss(input.float(), target)
+
+
+ds_loss = DeepSupervisionLoss

tests/test_ds_loss.py

@@ -0,0 +1,188 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+from parameterized import parameterized
+
+from monai.losses import DeepSupervisionLoss, DiceCELoss, DiceFocalLoss, DiceLoss
+from tests.utils import SkipIfBeforePyTorchVersion, test_script_save
+
+TEST_CASES_DICECE = [
+ [
+ {"to_onehot_y": True},
+ {},
+ {
+ "input": torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 0.606557,
+ ]
+]
+
+TEST_CASES_DICECE2 = [
+ [
+ {"to_onehot_y": True},
+ {},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 1.78144,
+ ],
+ [
+ {"to_onehot_y": True},
+ {"weight_mode": "same"},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 3.5529,
+ ],
+ [
+ {"to_onehot_y": True},
+ {"weight_mode": "two"},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 2.07973,
+ ],
+ [
+ {"to_onehot_y": True},
+ {"weights": [0.1, 0.2, 0.3]},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 0.76924,
+ ],
+]
+
+
+TEST_CASES_DICE = [
+ [
+ {"to_onehot_y": True},
+ {
+ "input": torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 0.166666, # the result equals to -1 + np.log(1 + np.exp(1))
+ ],
+ [
+ {"to_onehot_y": True},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 0.666665,
+ ],
+]
+
+TEST_CASES_DICEFOCAL = [
+ [
+ {"to_onehot_y": True},
+ {
+ "input": torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 0.32124, # the result equals to -1 + np.log(1 + np.exp(1))
+ ],
+ [
+ {"to_onehot_y": True},
+ {
+ "input": [
+ torch.tensor([[[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]], [[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ torch.tensor([[[[1.0, 1.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]]]]),
+ torch.tensor([[[[1.0], [0.0]], [[1.0], [0.0]]]]),
+ ],
+ "target": torch.tensor([[[[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]]]]),
+ },
+ 1.06452,
+ ],
+]
+
+
+class TestDSLossDiceCE(unittest.TestCase):
+ @parameterized.expand(TEST_CASES_DICECE)
+ def test_result(self, input_param, input_param2, input_data, expected_val):
+ diceceloss = DeepSupervisionLoss(DiceCELoss(**input_param), **input_param2)
+ result = diceceloss(**input_data)
+ np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+
+ def test_ill_shape(self):
+ loss = DeepSupervisionLoss(DiceCELoss())
+ with self.assertRaisesRegex(ValueError, ""):
+ loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3)))
+
+ def test_ill_reduction(self):
+ with self.assertRaisesRegex(ValueError, ""):
+ loss = DeepSupervisionLoss(DiceCELoss(reduction="none"))
+ loss(torch.ones((1, 2, 3)), torch.ones((1, 1, 2, 3)))
+
+ @SkipIfBeforePyTorchVersion((1, 10))
+ def test_script(self):
+ loss = DeepSupervisionLoss(DiceCELoss())
+ test_input = torch.ones(2, 1, 8, 8)
+ test_script_save(loss, test_input, test_input)
+
+
+@SkipIfBeforePyTorchVersion((1, 11))
+class TestDSLossDiceCE2(unittest.TestCase):
+ @parameterized.expand(TEST_CASES_DICECE2)
+ def test_result(self, input_param, input_param2, input_data, expected_val):
+ diceceloss = DeepSupervisionLoss(DiceCELoss(**input_param), **input_param2)
+ result = diceceloss(**input_data)
+ np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+
+
+@SkipIfBeforePyTorchVersion((1, 11))
+class TestDSLossDice(unittest.TestCase):
+ @parameterized.expand(TEST_CASES_DICE)
+ def test_result(self, input_param, input_data, expected_val):
+ loss = DeepSupervisionLoss(DiceLoss(**input_param))
+ result = loss(**input_data)
+ np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+
+
+@SkipIfBeforePyTorchVersion((1, 11))
+class TestDSLossDiceFocal(unittest.TestCase):
+ @parameterized.expand(TEST_CASES_DICEFOCAL)
+ def test_result(self, input_param, input_data, expected_val):
+ loss = DeepSupervisionLoss(DiceFocalLoss(**input_param))
+ result = loss(**input_data)
+ np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)
+
+
+if __name__ == "__main__":
+ unittest.main()