diff --git a/docs/source/inferers.rst b/docs/source/inferers.rst index 326f56e96c9..e1ba7f2edc5 100644 --- a/docs/source/inferers.rst +++ b/docs/source/inferers.rst @@ -25,6 +25,12 @@ Inferers :members: :special-members: __call__ +`MCDropoutInferer` +~~~~~~~~~~~~~~~~~~ +.. autoclass:: MCDropoutInferer + :members: + :special-members: __call__ + `SlidingWindowInferer` ~~~~~~~~~~~~~~~~~~~~~~ .. autoclass:: SlidingWindowInferer diff --git a/docs/source/metrics.rst b/docs/source/metrics.rst index ebd745e795f..b2442923d1a 100644 --- a/docs/source/metrics.rst +++ b/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 diff --git a/monai/inferers/__init__.py b/monai/inferers/__init__.py index fc78b9f7c4f..d729f13f3ef 100644 --- a/monai/inferers/__init__.py +++ b/monai/inferers/__init__.py @@ -17,6 +17,7 @@ DiffusionInferer, Inferer, LatentDiffusionInferer, + MCDropoutInferer, PatchInferer, SaliencyInferer, SimpleInferer, diff --git a/monai/inferers/inferer.py b/monai/inferers/inferer.py index 6d0f6838623..bbbc7394cee 100644 --- a/monai/inferers/inferer.py +++ b/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, diff --git a/monai/metrics/__init__.py b/monai/metrics/__init__.py index f55f92db1b8..f7cf0081d32 100644 --- a/monai/metrics/__init__.py +++ b/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 diff --git a/monai/metrics/active_learning_metrics.py b/monai/metrics/active_learning_metrics.py index a31d2fc1500..9dda723cc36 100644 --- a/monai/metrics/active_learning_metrics.py +++ b/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, diff --git a/tests/inferers/test_mc_dropout_inferer.py b/tests/inferers/test_mc_dropout_inferer.py new file mode 100644 index 00000000000..81faf8b29c1 --- /dev/null +++ b/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() diff --git a/tests/metrics/test_compute_uncertainty_decomposition.py b/tests/metrics/test_compute_uncertainty_decomposition.py new file mode 100644 index 00000000000..a5cfaf72e8c --- /dev/null +++ b/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()