Uncertainty quantification with PyTorch
See the code
Installation | Quickstart | Methods | Friends | Contributing | Citation | Documentation | Paper
posteriors?General purpose python library for uncertainty quantification with PyTorch.
transformers, lightning, torchopt, torch.distributions, pyro and more!posteriors is available on PyPI and can be installed via pip:
pip install posteriors
posteriors is functional first and aims to be easy to use and extend. Let's try it out
by training a simple model with variational inference:
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch import nn, utils, func
import torchopt
import posteriors
dataset = MNIST(root="./data", transform=ToTensor())
train_loader = utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
num_data = len(dataset)
classifier = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 10))
params = dict(classifier.named_parameters())
def log_posterior(params, batch):
images, labels = batch
images = images.view(images.size(0), -1)
output = func.functional_call(classifier, params, images)
log_post_val = (
-nn.functional.cross_entropy(output, labels)
+ posteriors.diag_normal_log_prob(params) / num_data
)
return log_post_val, output
transform = posteriors.vi.diag.build(
log_posterior, torchopt.adam(), temperature=1 / num_data
) # Can swap out for any posteriors algorithm
state = transform.init(params)
for batch in train_loader:
state, aux = transform.update(state, batch)
Observe that posteriors recommends specifying log_posterior and temperature such that
log_posterior remains on the same scale for different batch sizes. posteriors
algorithms are designed to be stable as temperature goes to zero.
Further, the output of log_posterior is a tuple containing the evaluation
(single-element Tensor) and an additional argument (TensorTree) containing any
auxiliary information we'd like to retain from the model call, here the model predictions.
If you have no auxiliary information, you can simply return torch.tensor([]) as
the second element. For more info see torch.func.grad
(with has_aux=True) or the documentation.
Check out the tutorials for more detailed usage!
posteriors supports a variety of methods for uncertainty quantification, including:
With full details available in the API documentation.
posteriors is designed to be easily extensible, if you're favorite method is not listed above,
raise an issue and we'll see what we can do!
Interfaces seamlessly with:
torch and in particular torch.func.torch.distributions for distributions and sampling, (note that it's typically required to set validate_args=False to conform with the control flows in torch.func).torchopt.transformers for pre-trained models.lightning for convenient training and logging, see examples/lightning_autoencoder.py.The functional transform interface is strongly inspired by frameworks such as
optax and blackjax.
As well as other UQ libraries fortuna,
laplace, numpyro,
pymc and uncertainty-baselines.
You can report a bug or request a feature by creating a new issue on GitHub.
If you want to contribute code, please check the contributing guide.
If you use posteriors in your research, please cite the library using the following BibTeX entry:
@inproceedings{duffield2025scalable,
title={Scalable Bayesian Learning with posteriors},
author={Samuel Duffield and Kaelan Donatella and Johnathan Chiu and Phoebe Klett and Daniel Simpson},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=fifXzmzeGy}
}
Python
100.0%
Uncertainty quantification with PyTorch
See the code
Installation | Quickstart | Methods | Friends | Contributing | Citation | Documentation | Paper
posteriors?General purpose python library for uncertainty quantification with PyTorch.
transformers, lightning, torchopt, torch.distributions, pyro and more!posteriors is available on PyPI and can be installed via pip:
pip install posteriors
posteriors is functional first and aims to be easy to use and extend. Let's try it out
by training a simple model with variational inference:
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torch import nn, utils, func
import torchopt
import posteriors
dataset = MNIST(root="./data", transform=ToTensor())
train_loader = utils.data.DataLoader(dataset, batch_size=32, shuffle=True)
num_data = len(dataset)
classifier = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 10))
params = dict(classifier.named_parameters())
def log_posterior(params, batch):
images, labels = batch
images = images.view(images.size(0), -1)
output = func.functional_call(classifier, params, images)
log_post_val = (
-nn.functional.cross_entropy(output, labels)
+ posteriors.diag_normal_log_prob(params) / num_data
)
return log_post_val, output
transform = posteriors.vi.diag.build(
log_posterior, torchopt.adam(), temperature=1 / num_data
) # Can swap out for any posteriors algorithm
state = transform.init(params)
for batch in train_loader:
state, aux = transform.update(state, batch)
Observe that posteriors recommends specifying log_posterior and temperature such that
log_posterior remains on the same scale for different batch sizes. posteriors
algorithms are designed to be stable as temperature goes to zero.
Further, the output of log_posterior is a tuple containing the evaluation
(single-element Tensor) and an additional argument (TensorTree) containing any
auxiliary information we'd like to retain from the model call, here the model predictions.
If you have no auxiliary information, you can simply return torch.tensor([]) as
the second element. For more info see torch.func.grad
(with has_aux=True) or the documentation.
Check out the tutorials for more detailed usage!
posteriors supports a variety of methods for uncertainty quantification, including:
With full details available in the API documentation.
posteriors is designed to be easily extensible, if you're favorite method is not listed above,
raise an issue and we'll see what we can do!
Interfaces seamlessly with:
torch and in particular torch.func.torch.distributions for distributions and sampling, (note that it's typically required to set validate_args=False to conform with the control flows in torch.func).torchopt.transformers for pre-trained models.lightning for convenient training and logging, see examples/lightning_autoencoder.py.The functional transform interface is strongly inspired by frameworks such as
optax and blackjax.
As well as other UQ libraries fortuna,
laplace, numpyro,
pymc and uncertainty-baselines.
You can report a bug or request a feature by creating a new issue on GitHub.
If you want to contribute code, please check the contributing guide.
If you use posteriors in your research, please cite the library using the following BibTeX entry:
@inproceedings{duffield2025scalable,
title={Scalable Bayesian Learning with posteriors},
author={Samuel Duffield and Kaelan Donatella and Johnathan Chiu and Phoebe Klett and Daniel Simpson},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=fifXzmzeGy}
}
Python
100.0%