Introduction
TorchUncertainty is a PyTorch-based framework developed by ENSTA Paris (U2IS) to facilitate uncertainty quantification in deep learning models. Released in 2023 under the Apache 2.0 license, the framework integrates a wide range of techniques for uncertainty-aware deep learning, including ensembles, Bayesian layers, post-hoc calibration, and out-of-distribution (OOD) detection. It is actively maintained, with contributions from the research community and regular updates to incorporate new methods and improvements.
Its design emphasizes collaboration, aiming to consolidate multiple uncertainty quantification techniques into a single, accessible toolkit for researchers and practitioners.
TorchUncertainty supports Python 3.10 or higher and is compatible with PyTorch. Installation is straightforward via PyPI, with additional setup options for contributors, including a Docker image to streamline experimentation on cloud-based GPU instances.
How It Works
TorchUncertainty provides a multi-level API designed to simplify the integration of uncertainty quantification techniques into deep learning workflows. The framework is structured to support four primary tasks: classification, probabilistic and pointwise regression, segmentation, and pixelwise regression (e.g., monocular depth estimation).
At its core, TorchUncertainty offers pre-built routines for uncertainty-aware training and evaluation, automated model assessment with proper scoring metrics, and a collection of layers, models, metrics, and losses that can be seamlessly integrated into existing networks.
The framework’s modular design allows users to select and combine methods based on their specific requirements. For example, users can apply Bayesian neural networks for probabilistic predictions, ensemble methods like Deep Ensembles or BatchEnsemble for improved reliability, or post-hoc calibration techniques such as temperature scaling to refine model outputs.
Use Cases
TorchUncertainty is particularly valuable in applications where reliable uncertainty estimation is critical. Its capabilities span multiple domains, including risk-sensitive predictions, anomaly detection, forecasting, scientific analysis, image segmentation, and decision-making systems.
- Risk-sensitive predictions: The framework enables models to provide uncertainty estimates alongside predictions, which is essential in high-stakes domains such as healthcare, finance, and autonomous systems. Techniques like selective classification and OOD detection help models abstain from predictions when uncertainty is high, reducing the risk of erroneous decisions.
- Anomaly detection: By quantifying uncertainty, TorchUncertainty can identify outliers or anomalous inputs where the model’s confidence is low.
- Forecasting and scientific analysis: For tasks involving continuous-valued outputs, such as time-series forecasting or physical simulations, TorchUncertainty supports probabilistic regression and quantile regression. These methods provide uncertainty bounds that are essential for assessing the reliability of predictions in fields like meteorology, economics, and engineering.
- Image segmentation: The framework includes tools for pixelwise regression and segmentation, enabling models to produce uncertainty maps alongside segmentation outputs. This is useful in medical imaging, autonomous driving, and remote sensing, where understanding the confidence of each pixel’s classification is vital for downstream decision-making.
- Decision-making systems: In applications where models inform critical decisions, such as autonomous vehicles or adaptive control systems, uncertainty quantification ensures that predictions are accompanied by measures of confidence.
Key Features & Capabilities
TorchUncertainty implements a comprehensive suite of uncertainty quantification methods, categorized into uncertainty quantification models, augmentation methods, and post-processing techniques. Below is a detailed breakdown of the supported methods:
- Uncertainty Quantification Models:
- Deep Ensembles, BatchEnsemble, Masksembles, and MIMO: Techniques for aggregating predictions from multiple models to improve reliability and reduce variance.
- MC-Dropout: A method for approximating Bayesian neural networks by applying dropout at test time to estimate uncertainty.
- Packed-Ensembles: An efficient ensemble method that reduces computational overhead while maintaining uncertainty estimation capabilities.
- Mixup Ensembles and MixupMP: Ensembles that incorporate data augmentation techniques like Mixup to improve generalization and uncertainty estimates.
- Variational Bayesian Neural Networks: Models that treat network parameters as probabilistic distributions for inherent uncertainty quantification.
- Checkpoint Ensembles and Snapshot Ensembles: Methods for generating diverse model checkpoints to approximate ensembles without additional training.
- Deep Evidential Classification and Regression: Methods that provide uncertainty estimates by modeling belief and disbelief in predictions.
- Quantile Regression with Pinball Loss and RQR: Techniques for estimating prediction intervals and quantiles to capture uncertainty in continuous outputs.
- Augmentation Methods:
- Mixup, MixupIO, RegMixup, and WarpingMixup: Data augmentation techniques that improve model robustness and uncertainty estimation by generating synthetic training examples.
- Post-processing Methods:
- Histogram binning, isotonic regression, and Bayesian binning into quantiles: Methods for refining uncertainty estimates post-hoc.
- DEUP (Direct Epistemic Uncertainty Prediction): A post-hoc method for estimating epistemic uncertainty by predicting errors directly.
- Metrics and Evaluation:
TorchUncertainty includes a wide range of metrics for evaluating uncertainty quantification performance across tasks.
- Out-of-distribution detection: AUROC, AUPR, and FPR95 for evaluating the ability to detect inputs from unseen distributions.
- Prediction diversity: Disagreement, Entropy, MutualInformation, and VariationRatio for assessing the diversity of ensemble predictions.
- Pixelwise-regression and segmentation metrics: Task-specific metrics for evaluating uncertainty in spatial predictions.
The framework also includes official implementations of several influential papers, such as Packed-Ensembles for Efficient Uncertainty Estimation (ICLR 2023), LP-BNN for encoding latent posteriors in Bayesian neural networks (IEEE TPAMI 2023), and MUAD for benchmarking multiple uncertainty types in autonomous driving (BMVC 2022). These implementations provide validated baselines for researchers and practitioners.
TorchUncertainty also supports experiment tracking via MLflow, allowing users to log and compare runs, inspect metric curves, and manage artifacts. This integration facilitates reproducibility and collaboration in research settings.
Official Implementations and Tutorials
TorchUncertainty includes official implementations of several key research papers, providing validated and reproducible code for advanced uncertainty quantification techniques. These implementations serve as benchmarks and starting points for researchers:
- Packed-Ensembles for Efficient Uncertainty Estimation (ICLR 2023): An ensemble method designed to reduce computational overhead while maintaining uncertainty estimation accuracy.
Additionally, the project maintains configuration files for experiments, with results tracked using MLflow for easy comparison and analysis.
Projects utilizing TorchUncertainty include research on uncertainty quantification in text-to-image generation (CVPR 2025), the impact of label smoothing on selective classification (ICLR 2025), and exploring Bayesian neural network posteriors (ICLR 2024).