TyXe: Pyro-based Bayesian neural nets for Pytorch

by   Hippolyt Ritter, et al.

We introduce TyXe, a Bayesian neural network library built on top of Pytorch and Pyro. Our leading design principle is to cleanly separate architecture, prior, inference and likelihood specification, allowing for a flexible workflow where users can quickly iterate over combinations of these components. In contrast to existing packages TyXe does not implement any layer classes, and instead relies on architectures defined in generic Pytorch code. TyXe then provides modular choices for canonical priors, variational guides, inference techniques, and layer selections for a Bayesian treatment of the specified architecture. Sampling tricks for variance reduction, such as local reparameterization or flipout, are implemented as effect handlers, which can be applied independently of other specifications. We showcase the ease of use of TyXe to explore Bayesian versions of popular models from various libraries: toy regression with a pure Pytorch neural network; large-scale image classification with torchvision ResNets; graph neural networks based on DGL; and Neural Radiance Fields built on top of Pytorch3D. Finally, we provide convenient abstractions for variational continual learning. In all cases the change from a deterministic to a Bayesian neural network comes with minimal modifications to existing code, offering a broad range of researchers and practitioners alike practical access to uncertainty estimation techniques. The library is available at https://github.com/TyXe-BDL/TyXe.


SequeL: A Continual Learning Library in PyTorch and JAX

Continual Learning is an important and challenging problem in machine le...

Avalanche: A PyTorch Library for Deep Continual Learning

Continual learning is the problem of learning from a nonstationary strea...

Indian Buffet Neural Networks for Continual Learning

We place an Indian Buffet Process (IBP) prior over the neural structure ...

Bayesian Deep Net GLM and GLMM

Deep feedforward neural networks (DFNNs) are a powerful tool for functio...

Epistemic Neural Networks

We introduce the epistemic neural network (ENN) as an interface for unce...

torchgfn: A PyTorch GFlowNet library

The increasing popularity of generative flow networks (GFlowNets or GFNs...

Please sign up or login with your details

Forgot password? Click here to reset