Package: tyxe
TyXe: Pyro-based BNNs for Pytorch usersTyXe aims to simplify the process of turningPytorchneural networks into Bayesian neural networks by
leveraging the model definition and inference capabilities ofPyro.
Our core design principle is to cleanly separate the construction of neural architecture, prior, inference distribution
and likelihood, enabling a flexible workflow where each component can be exchanged independently.
Defining a BNN in TyXe takes as little as 5 lines of code:net = nn.Sequential(nn.Linear(1, 50), nn.Tanh(), nn.Linear(50, 1))
prior = tyxe.priors.IIDPrior(dist.Normal(0, 1))
likelihood = tyxe.likelihoods.HomoskedasticGaussian(scale=0.1)
inference = tyxe.guides.AutoNormal
bnn = tyxe.VariationalBNN(net, prior, likelihood, inference)In the following, we assume that you (roughly) know what a BNN is mathematically.Motivating exampleStandard neural networks give us a single function that fits the data, but many different ones are typically plausible.
With only a single fit, we don't know for what inputs the model is 'certain' (because there is training data nearby) and
where it is uncertain.Maximum likelihood fitPosterior samplesImplementing the former can be achieved easily in a few lines of Pytorch code, but training a BNN that gives a
distribution over different fits is typically more complicated and is specifically what we aim to simplify.TrainingConstructing a BNN object has been shown in the example above.
For fitting the posterior approximation, we provide a high-level.fitmethod similar to libraries such as scikit-learn
or keras:optim = pyro.optim.Adam({"lr": 1e-3})
bnn.fit(data_loader, optim, num_epochs)Prediction & evaluationFurther we provide.predictand.evaluationmethods, which make predictions based on multiple samples from the approximate posterior, average them based on the observation model, and return log likelihoods and an error measure:predictions = bnn.predict(x_test, num_samples)
error, log_likelihood = bnn.evaluate(x_test, y_test, num_samples)Local reparameterizationWe implementlocal reparameterizationfor factorized Gaussians as a poutine, which reduces gradient noise during training.
This means it can be enabled or disabled at both during training and prediction with a context manager:with tyxe.poutine.local_reparameterization():
    bnn.fit(data_loader, optim, num_epochs)
    bnn.predict(x_test, num_predictions)At the moment, this poutine does not work with theAutoNormalandAutoDiagonalNormalguides in pyro, since those draw the weights from a Delta distribution, so you need to usetyxe.guides.ParameterwiseDiagonalNormalas your guide.MCMCWe provide a unified interface to pyro's MCMC implementations, simply use thetyxe.MCMC_BNNclass instead and provide a kernel instead of the guide:kernel = pyro.infer.mcmcm.NUTS
bnn = tyxe.MCMC_BNN(net, prior, likelihood, kernel)Any parameters that pyro'sMCMCclass accepts can be passed through the keyword arguments of the.fitmethod.Continual learningDue to our design that cleanly separates the prior from guide, architecture and likelihood, it is easy to update it in a continual setting.
For example, you can construct atyxe.priors.DictPriorby extracting the distributions over all weights and biases from aParameterwiseDiagonalNormalinstance using theget_detached_distributionsmethod and pass it tobnn.update_priorto implementVariational Continual Learningin a few lines of code.
Seeexamples/vcl.pyfor a basic example on split-MNIST and split-CIFAR.Network architecturesWe don't implement any layer classes.
You construct your network in Pytorch and then turn it into a BNN, which makes it easy to apply the same prior and inference strategies to different neural networks.InferenceFor inference, we mainly provide an equivalent to pyro'sAutoDiagonalNormalthat is compatible with local reparameterization intyxe.guides.
This module also contains a few helper functions for initialization of Gaussian mean parameters, e.g. to the values of a pre-trained network.
It should be possible to use any of pyro's autoguides for variational inference.
Seeexamples/resnet.pyfor a few options as well as initializing to pre-trained weights.PriorsThe priors can be found intyxe.priors.
We currently only support placing priors on the parameters.
Through the expose and hide arguments in the init method you can specify layers, types of layers and specific parameters over which you want to place a prior.
This helps, for example in learning the parameters of BatchNorm layers deterministically.Likelihoodstyxe.observation_modelscontains classes that wrap the most commontorch.distributionsfor specifying noise models of data toInstallationWe recommend installing TyXe using conda with the providedenvironment.yml, which also installs all the dependencies for the examples except for Pytorch3d, which needs to be added manually.
The environment assumes that you are using CUDA11.0, if this is not the case, simply change thecudatoolkitanddgl-cudaversions before running:conda env create -f environment.yml
conda activate tyxe
pip install -e .CitationIf you use TyXe, please consider citing:@inproceedings{ritter2022tyxe,
  title={Ty{X}e: Pyro-based {B}ayesian neural nets for {P}ytorch},
  author={Ritter, Hippolyt and Karaletsos, Theofanis},
  booktitle={Proceedings of Machine Learning and Systems},
  year={2022}
}