Add MCDropoutInferer and uncertainty decomposition

docs/source/inferers.rst

@@ -25,6 +25,12 @@ Inferers
     :members:
     :special-members: __call__
 
+`MCDropoutInferer`
+~~~~~~~~~~~~~~~~~~
+.. autoclass:: MCDropoutInferer
+    :members:
+    :special-members: __call__
+
 `SlidingWindowInferer`
 ~~~~~~~~~~~~~~~~~~~~~~
 .. autoclass:: SlidingWindowInferer

docs/source/metrics.rst

@@ -199,6 +199,13 @@ Metrics
 .. autoclass:: MetricsReloadedCategorical
     :members:
 
+`Uncertainty decomposition`
+----------------------------
+.. autofunction:: compute_uncertainty_decomposition
+
+.. autoclass:: UncertaintyDecompositionMetric
+    :members:
+
 `Calibration Error`
 -------------------
 .. autofunction:: calibration_binning

monai/inferers/__init__.py

@@ -17,6 +17,7 @@
     DiffusionInferer,
     Inferer,
     LatentDiffusionInferer,
+    MCDropoutInferer,
     PatchInferer,
     SaliencyInferer,
     SimpleInferer,

monai/inferers/inferer.py

@@ -444,6 +444,64 @@ def __call__(
         return network(inputs, *args, **kwargs)
 
 
+class MCDropoutInferer(Inferer):
+    """
+    Monte Carlo dropout inference (Gal & Ghahramani, "Dropout as a Bayesian Approximation", ICML 2016): runs ``inferer``
+    ``num_samples`` times with the dropout layers of ``network`` active, giving a sample of predictions from which an
+    uncertainty can be estimated, for example with :py:class:`monai.metrics.UncertaintyDecompositionMetric` or
+    :py:class:`monai.metrics.VarianceMetric`.
+
+    Only dropout layers are switched to training mode, so normalization layers keep their evaluation behavior. The
+    original training/evaluation mode of every dropout layer is restored afterwards. ``network`` must be a
+    ``torch.nn.Module`` containing at least one dropout layer.
+
+    .. code-block:: python
+
+        inferer = MCDropoutInferer(SlidingWindowInferer(roi_size=(96, 96, 96)), num_samples=10)
+        samples = inferer(image, model)  # [num_samples, B, C, H, W, D]
+        probs = torch.softmax(samples, dim=2)
+        uncertainty = UncertaintyDecompositionMetric(spatial_map=True)(probs[:, 0])  # first image of the batch
+
+    Args:
+        inferer: the inferer used for each stochastic forward pass, defaults to :py:class:`SimpleInferer`.
+        num_samples: number of stochastic forward passes, must be at least 1.
+        return_samples: if True (default), return the stacked samples with a new leading dimension of size
+            ``num_samples``; otherwise return their mean, so the inferer can be used wherever a single prediction is
+            expected (for example as ``inferer`` of an evaluator).
+
+    """
+
+    def __init__(self, inferer: Inferer | None = None, num_samples: int = 10, return_samples: bool = True) -> None:
+        super().__init__()
+        if num_samples < 1:
+            raise ValueError(f"num_samples must be at least 1, got {num_samples}.")
+        self.inferer = SimpleInferer() if inferer is None else inferer
+        self.num_samples = num_samples
+        self.return_samples = return_samples
+
+    def __call__(self, inputs: torch.Tensor, network: nn.Module, *args: Any, **kwargs: Any) -> torch.Tensor:
+        """
+        Args:
+            inputs: model input data for inference.
+            network: target model to execute inference, a ``torch.nn.Module`` with dropout layers.
+            args: optional args to be passed to ``network``.
+            kwargs: optional keyword args to be passed to ``network``.
+
+        """
+        dropout_layers = [m for m in network.modules() if isinstance(m, nn.modules.dropout._DropoutNd)]
+        if not dropout_layers:
+            raise ValueError("network has no dropout layers, so all Monte Carlo dropout samples would be identical.")
+        original_modes = [m.training for m in dropout_layers]
+        try:
+            for m in dropout_layers:
+                m.train()
+            samples = torch.stack([self.inferer(inputs, network, *args, **kwargs) for _ in range(self.num_samples)])
+        finally:
+            for m, mode in zip(dropout_layers, original_modes):
+                m.train(mode)
+        return samples if self.return_samples else samples.mean(dim=0)
+
+
 class SlidingWindowInferer(Inferer):
     """
     Sliding window method for model inference,

monai/metrics/__init__.py

@@ -12,7 +12,14 @@
 from __future__ import annotations
 
 from .absolute_volume_difference import AbsoluteVolumeDifferenceMetric, compute_absolute_volume_difference
-from .active_learning_metrics import LabelQualityScore, VarianceMetric, compute_variance, label_quality_score
+from .active_learning_metrics import (
+    LabelQualityScore,
+    UncertaintyDecompositionMetric,
+    VarianceMetric,
+    compute_uncertainty_decomposition,
+    compute_variance,
+    label_quality_score,
+)
 from .average_precision import AveragePrecisionMetric, compute_average_precision
 from .calibration import CalibrationErrorMetric, CalibrationReduction, calibration_binning
 from .confusion_matrix import ConfusionMatrixMetric, compute_confusion_matrix_metric, get_confusion_matrix

monai/metrics/active_learning_metrics.py

@@ -105,6 +105,85 @@ def __call__(self, y_pred: Any, y: Any) -> torch.Tensor | None:
         )
 
 
+class UncertaintyDecompositionMetric(Metric):
+    """
+    Decompose the predictive uncertainty of a given T-repeats N-dimensional probability array/tensor into its total,
+    aleatoric and epistemic parts, see :py:func:`compute_uncertainty_decomposition`.
+
+    The repeats are typically the outputs of Monte Carlo dropout (:py:class:`monai.inferers.MCDropoutInferer`), of test
+    time augmentation (:py:class:`monai.data.TestTimeAugmentation` with ``return_full_data=True``) or of an ensemble.
+
+    Args:
+        spatial_map: Boolean, if set to True, spatial maps of the uncertainties will be returned corresponding to i/p
+            image dimensions
+        scalar_reduction: reduction type of the metric, either 'sum' or 'mean' can be used
+        eps: small value added inside the logarithms to avoid ``log(0)``
+
+    """
+
+    def __init__(self, spatial_map: bool = False, scalar_reduction: str = "mean", eps: float = 1e-8) -> None:
+        super().__init__()
+        self.spatial_map = spatial_map
+        self.scalar_reduction = scalar_reduction
+        self.eps = eps
+
+    def __call__(self, y_pred: Any) -> dict[str, torch.Tensor]:
+        """
+        Args:
+            y_pred: class probabilities, typically softmax of the model output.
+                It must be N-repeats, repeat-first tensor [N,C,H,W,D].
+
+        Returns:
+            Dictionary with keys ``"total"``, ``"aleatoric"`` and ``"epistemic"``, each a scalar or a spatial map
+
+        """
+        return compute_uncertainty_decomposition(
+            y_pred=y_pred, spatial_map=self.spatial_map, scalar_reduction=self.scalar_reduction, eps=self.eps
+        )
+
+
+def compute_uncertainty_decomposition(
+    y_pred: torch.Tensor, spatial_map: bool = False, scalar_reduction: str = "mean", eps: float = 1e-8
+) -> dict[str, torch.Tensor]:
+    """
+    Decompose predictive uncertainty into total, aleatoric and epistemic parts with entropies:
+
+    - total: entropy of the mean prediction (predictive entropy), ``H[E_n p_n]``
+    - aleatoric: mean entropy of the individual predictions (expected entropy), ``E_n H[p_n]``
+    - epistemic: their difference, the mutual information between the prediction and the model, ``total - aleatoric``
+
+    Entropies are in nats. See Depeweg et al., "Decomposition of Uncertainty in Bayesian Deep Learning for Efficient
+    and Risk-sensitive Learning", ICML 2018.
+
+    Args:
+        y_pred: [N, C, H, W, D] or [N, C, H, W] or [N, C] class probabilities where N is repeats, C is channels and
+            H, W, D stand for Height, Width & Depth. Probabilities must sum to one over C.
+        spatial_map: Boolean, if set to True, spatial maps of the uncertainties will be returned corresponding to i/p
+            image dimensions
+        scalar_reduction: reduction type of the metric, either 'sum' or 'mean' can be used
+        eps: small value added inside the logarithms to avoid ``log(0)``
+    Returns:
+        Dictionary with keys ``"total"``, ``"aleatoric"`` and ``"epistemic"``, each a scalar or a spatial map
+    """
+    if y_pred.ndim < 2:
+        raise ValueError(f"y_pred must have at least 2 dimensions [N, C, ...], got shape {tuple(y_pred.shape)}.")
+    y_pred = y_pred.float()
+
+    mean_pred = y_pred.mean(dim=0)
+    total = -(mean_pred * torch.log(mean_pred + eps)).sum(dim=0)
+    aleatoric = -(y_pred * torch.log(y_pred + eps)).sum(dim=1).mean(dim=0)
+    # clamp the small negative values that rounding can produce
+    uncertainties = {"total": total, "aleatoric": aleatoric, "epistemic": (total - aleatoric).clamp_min(0)}
+
+    if spatial_map:
+        return uncertainties
+    if scalar_reduction == MetricReduction.MEAN:
+        return {k: torch.mean(v) for k, v in uncertainties.items()}
+    if scalar_reduction == MetricReduction.SUM:
+        return {k: torch.sum(v) for k, v in uncertainties.items()}
+    raise ValueError(f"scalar_reduction={scalar_reduction} not supported.")
+
+
 def compute_variance(
     y_pred: torch.Tensor,
     include_background: bool = True,

tests/inferers/test_mc_dropout_inferer.py

@@ -0,0 +1,83 @@
+# 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 __future__ import annotations
+
+import unittest
+
+import torch
+import torch.nn as nn
+
+from monai.inferers import MCDropoutInferer, SimpleInferer, SlidingWindowInferer
+
+
+def _net(with_dropout: bool = True) -> nn.Module:
+    layers = [nn.Conv2d(1, 4, 3, padding=1), nn.BatchNorm2d(4), nn.ReLU()]
+    if with_dropout:
+        layers.append(nn.Dropout2d(0.5))
+    layers.append(nn.Conv2d(4, 2, 1))
+    return nn.Sequential(*layers).eval()
+
+
+class TestMCDropoutInferer(unittest.TestCase):
+
+    def setUp(self):
+        torch.manual_seed(0)
+        self.inputs = torch.rand(2, 1, 16, 16)
+
+    def test_samples_vary_and_have_leading_sample_dim(self):
+        net = _net()
+        samples = MCDropoutInferer(num_samples=5)(self.inputs, net)
+        self.assertEqual(samples.shape, (5, 2, 2, 16, 16))
+        self.assertFalse(torch.allclose(samples[0], samples[1]))
+
+    def test_modes_are_restored_and_batchnorm_stays_in_eval(self):
+        net = _net()
+        bn = net[1]
+        running_mean = bn.running_mean.clone()
+        MCDropoutInferer(num_samples=3)(self.inputs, net)
+        self.assertFalse(any(m.training for m in net.modules()))
+        torch.testing.assert_close(bn.running_mean, running_mean)
+
+    def test_modes_are_restored_on_error(self):
+        net = _net()
+
+        def failing_network(x):
+            raise RuntimeError("boom")
+
+        failing_network.modules = net.modules
+        with self.assertRaises(RuntimeError):
+            MCDropoutInferer(num_samples=2)(self.inputs, failing_network)
+        self.assertFalse(any(m.training for m in net.modules()))
+
+    def test_return_mean(self):
+        net = _net()
+        torch.manual_seed(1)
+        samples = MCDropoutInferer(num_samples=4)(self.inputs, net)
+        torch.manual_seed(1)
+        mean = MCDropoutInferer(num_samples=4, return_samples=False)(self.inputs, net)
+        torch.testing.assert_close(mean, samples.mean(dim=0))
+
+    def test_wraps_sliding_window_inferer(self):
+        net = _net()
+        inferer = MCDropoutInferer(SlidingWindowInferer(roi_size=(8, 8), sw_batch_size=2), num_samples=3)
+        samples = inferer(self.inputs, net)
+        self.assertEqual(samples.shape, (3, 2, 2, 16, 16))
+
+    def test_invalid_arguments(self):
+        with self.assertRaises(ValueError):
+            MCDropoutInferer(num_samples=0)
+        with self.assertRaises(ValueError):
+            MCDropoutInferer(SimpleInferer())(self.inputs, _net(with_dropout=False))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/metrics/test_compute_uncertainty_decomposition.py

@@ -0,0 +1,76 @@
+# 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 __future__ import annotations
+
+import math
+import unittest
+
+import torch
+
+from monai.metrics import UncertaintyDecompositionMetric, compute_uncertainty_decomposition
+
+LN2 = math.log(2)
+
+
+class TestComputeUncertaintyDecomposition(unittest.TestCase):
+
+    def test_identical_confident_samples_have_no_uncertainty(self):
+        y_pred = torch.zeros(4, 2, 3, 3)
+        y_pred[:, 0] = 1.0
+        result = compute_uncertainty_decomposition(y_pred, spatial_map=True)
+        for key in ("total", "aleatoric", "epistemic"):
+            self.assertEqual(result[key].shape, (3, 3))
+            torch.testing.assert_close(result[key], torch.zeros(3, 3), atol=1e-6, rtol=0)
+
+    def test_disagreeing_confident_samples_are_epistemic(self):
+        y_pred = torch.zeros(2, 2, 4)
+        y_pred[0, 0] = 1.0
+        y_pred[1, 1] = 1.0
+        result = compute_uncertainty_decomposition(y_pred, spatial_map=True)
+        torch.testing.assert_close(result["total"], torch.full((4,), LN2), atol=1e-5, rtol=0)
+        torch.testing.assert_close(result["aleatoric"], torch.zeros(4), atol=1e-5, rtol=0)
+        torch.testing.assert_close(result["epistemic"], torch.full((4,), LN2), atol=1e-5, rtol=0)
+
+    def test_agreeing_uniform_samples_are_aleatoric(self):
+        y_pred = torch.full((5, 2, 4), 0.5)
+        result = compute_uncertainty_decomposition(y_pred, spatial_map=True)
+        torch.testing.assert_close(result["total"], torch.full((4,), LN2), atol=1e-5, rtol=0)
+        torch.testing.assert_close(result["aleatoric"], torch.full((4,), LN2), atol=1e-5, rtol=0)
+        torch.testing.assert_close(result["epistemic"], torch.zeros(4), atol=1e-5, rtol=0)
+
+    def test_scalar_reductions_and_total_is_sum_of_parts(self):
+        y_pred = torch.softmax(torch.randn(6, 3, 5, 5, generator=torch.Generator().manual_seed(0)), dim=1)
+        maps = compute_uncertainty_decomposition(y_pred, spatial_map=True)
+        torch.testing.assert_close(maps["total"], maps["aleatoric"] + maps["epistemic"], atol=1e-5, rtol=0)
+        self.assertTrue(bool((maps["epistemic"] >= 0).all()))
+        means = compute_uncertainty_decomposition(y_pred, scalar_reduction="mean")
+        sums = compute_uncertainty_decomposition(y_pred, scalar_reduction="sum")
+        for key in maps:
+            torch.testing.assert_close(means[key], maps[key].mean())
+            torch.testing.assert_close(sums[key], maps[key].sum())
+
+    def test_metric_class_matches_function(self):
+        y_pred = torch.softmax(torch.randn(4, 2, 3, 3, 3, generator=torch.Generator().manual_seed(1)), dim=1)
+        expected = compute_uncertainty_decomposition(y_pred, spatial_map=True)
+        result = UncertaintyDecompositionMetric(spatial_map=True)(y_pred)
+        for key in expected:
+            torch.testing.assert_close(result[key], expected[key])
+
+    def test_invalid_inputs(self):
+        with self.assertRaises(ValueError):
+            compute_uncertainty_decomposition(torch.rand(4))
+        with self.assertRaises(ValueError):
+            compute_uncertainty_decomposition(torch.rand(4, 2), scalar_reduction="max")
+
+
+if __name__ == "__main__":
+    unittest.main()