[eliem08]

Part of #5207. Opened as a draft to get feedback on the API before polishing; happy to rename or move things.

Description

The 2022 discussion in #5207 concluded that test-time augmentation needs an Engine, because it re-runs transforms per subject. Monte Carlo dropout does not, so this adds it as a thin inferer, plus a way to decompose the resulting uncertainty:

monai.inferers.MCDropoutInferer(inferer=None, num_samples=10, return_samples=True)

• wraps any inferer (default SimpleInferer, composes with SlidingWindowInferer);
• switches only dropout layers (nn.modules.dropout._DropoutNd) to train mode, so normalization layers keep their eval behaviour, and restores every dropout layer's mode afterwards, also on error;
• returns the samples stacked on a new leading dim, or their mean with return_samples=False so it can be an evaluator's inferer;
• raises if the network has no dropout layers (all samples would be identical).

monai.metrics.compute_uncertainty_decomposition / UncertaintyDecompositionMetric, next to VarianceMetric and with the same repeat-first [N, C, ...] input, spatial_map and scalar_reduction conventions, return

• total: predictive entropy H[E_n p_n]
• aleatoric: expected entropy E_n H[p_n]
• epistemic: mutual information, total - aleatoric (Depeweg et al., ICML 2018).

They work on samples from MC dropout, TestTimeAugmentation(return_full_data=True) or an ensemble.

inferer = MCDropoutInferer(SlidingWindowInferer(roi_size=(96, 96, 96)), num_samples=10)
samples = inferer(image, model)  # [10, B, C, H, W, D]
maps = UncertaintyDecompositionMetric(spatial_map=True)(torch.softmax(samples, dim=2)[:, 0])

Checked end to end with UNet(..., dropout=0.2) and SlidingWindowInferer on a 48³ volume: samples (8, 1, 3, 48, 48, 48), per-voxel maps (48, 48, 48), all modules back in eval mode afterwards.

Questions for reviewers:

1. Is active_learning_metrics.py the right home for the decomposition, given VarianceMetric is there?
2. Should a follow-up add a tutorial combining TTA and MC dropout (aleatoric vs epistemic maps)?

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• New tests added to cover the changes.
• Documentation updated (docs/source/inferers.rst, docs/source/metrics.rst).

Tests

• tests/metrics/test_compute_uncertainty_decomposition.py: analytic cases (identical confident samples give zero uncertainty everywhere; disagreeing confident samples give total = epistemic = ln 2, aleatoric 0; uniform samples give total = aleatoric = ln 2, epistemic 0), total = aleatoric + epistemic, reductions, metric class vs function, invalid inputs.
• tests/inferers/test_mc_dropout_inferer.py: samples vary, BatchNorm running stats unchanged, modes restored (also on error), mean mode, SlidingWindowInferer wrapping, invalid arguments.
• With the existing test_compute_variance, test_label_quality_score, saliency and slice inferer tests: 48 passed. black, isort and ruff pass.

I prepared this with AI assistance and have reviewed it; I work on uncertainty estimation for breast cancer pathology.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}