TorchUncertainty is a package designed to help leverage uncertainty quantification techniques to make deep neural networks more reliable. It aims at being collaborative and including as many methods as possible, so reach out to add yours!

🚧 TorchUncertainty is in early development 🚧 - expect changes, but reach out and contribute if you are interested in the project! Please raise an issue if you have any bugs or difficulties and join the discord server.

📚 Our webpage and documentation is available here: torch-uncertainty.github.io. 📚

TorchUncertainty contains the official implementations of multiple papers from major machine-learning and computer vision conferences and was featured in tutorials at WACV 2024, HAICON 2024 and ECCV 2024.

This package provides a multi-level API, including:

• easy-to-use ⚡ lightning uncertainty-aware training & evaluation routines for 4 tasks: classification, probabilistic and pointwise regression, and segmentation.
• fully automated evaluation of the performance of models with proper scores, selective classification, out-of-distribution detection and distribution shift performance metrics!
• ready-to-train baselines on research datasets, such as ImageNet and CIFAR
• layers, models, metrics, & losses available for your networks
• scikit-learn style post-processing methods such as Temperature Scaling.
• transformations and augmentations, including corruptions resulting in additional "corrupted datasets" available on HuggingFace

Have a look at the Reference page or the API reference for a more exhaustive list of the implemented methods, datasets, metrics, etc.

TorchUncertainty requires Python 3.10 or greater. Install the desired PyTorch version in your environment. Then, install the package from PyPI:

pip install torch-uncertainty

The installation procedure for contributors is different: have a look at the contribution page.

For contributors running experiments on cloud GPU instances, we provide a pre-built Docker image that includes all necessary dependencies and configurations and the Dockerfile for building your custom Docker images. This allows you to quickly launch an experiment-ready container with minimal setup. Please refer to DOCKER.md for further details.

We make a quickstart available at torch-uncertainty.github.io/quickstart.

TorchUncertainty currently supports classification, probabilistic and pointwise regression, segmentation and pixelwise regression (such as monocular depth estimation).

We also provide the following methods:

To date, the following deep learning uncertainty quantification modes have been implemented. Click 📥 on the methods for tutorials:

• Deep Ensembles, BatchEnsemble, Masksembles, & MIMO
• MC-Dropout
• Packed-Ensembles (see Blog post)
• Variational Bayesian Neural Networks
• Checkpoint Ensembles & Snapshot Ensembles
• Stochastic Weight Averaging & Stochastic Weight Averaging Gaussian
• Deep Evidential Classification & Regression
• Regression with Beta Gaussian NLL Loss
• Test-time adaptation with Zero

The following data augmentation methods have been implemented:

• Mixup, MixupIO, RegMixup, WarpingMixup
• Modernized corruptions to evaluate model performance under distribution shift

To date, the following post-processing methods have been implemented:

• Temperature, Vector, & Matrix scaling
• Conformal Predictions with APS and RAPS
• Monte Carlo Batch Normalization
• Laplace approximation through the Laplace library

It includes the official codes of the following papers:

• Packed-Ensembles for Efficient Uncertainty Estimation - ICLR 2023 - Tutorial
• LP-BNN: Encoding the latent posterior of Bayesian Neural Networks for uncertainty quantification - IEEE TPAMI 2023
• MUAD: Multiple Uncertainties for Autonomous Driving, a benchmark for multiple uncertainty types and tasks - BMVC 2022

Check out all our tutorials at torch-uncertainty.github.io/auto_tutorials.

The following projects use TorchUncertainty:

• Towards Understanding Why Label Smoothing Degrades Selective Classification and How to Fix It - ICLR 2025
• A Symmetry-Aware Exploration of Bayesian Neural Network Posteriors - ICLR 2024

If you are using TorchUncertainty in your project, please let us know, and we will add your project to this list!