TyXe: Pyro-based Bayesian neural nets for Pytorch

Abstract

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.

arXiv commentPreviously presented at PROBPROG 2020

References in corpus (20)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• TensorFlow: A system for large-scale machine learning
• The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian Monte Carlo
• Weight Uncertainty in Neural Networks
• Continual Learning Through Synaptic Intelligence
• Deep Graph Library: A Graph-Centric, Highly-Performant Package for Graph Neural Networks
• Probabilistic Backpropagation for Scalable Learning of Bayesian Neural Networks
• Variational Dropout and the Local Reparameterization Trick
• Multiplicative Normalizing Flows for Variational Bayesian Neural Networks
• Structured and Efficient Variational Deep Learning with Matrix Gaussian Posteriors
• Automated Variational Inference in Probabilistic Programming
• Improving and Understanding Variational Continual Learning
• 'In-Between' Uncertainty in Bayesian Neural Networks
• Overpruning in Variational Bayesian Neural Networks
• The k-tied Normal Distribution: A Compact Parameterization of Gaussian Mean Field Posteriors in Bayesian Neural Networks
• BNNpriors: A library for Bayesian neural network inference with different prior distributions
• Uncertainty Quantification 360: A Holistic Toolkit for Quantifying and Communicating the Uncertainty of AI
• Probabilistic Meta-Representations Of Neural Networks
• Bayesian Graph Neural Networks for Molecular Property Prediction
• Sparse Uncertainty Representation in Deep Learning with Inducing Weights