torch_blue.vi

This module provides basic layers and loss functions for BNN-training with variational inference.

Submodules

• torch_blue.vi.distributions
• torch_blue.vi.utils

Classes

AnalyticalKullbackLeiblerLoss	Analytical Kullback-Leibler loss function.
KullbackLeiblerModule	Base class for modules calculating the Kullback-Leibler divergence from distribution parameters.
NonBayesianDivergence	Placeholder Kullback-Leibler divergence for non-Bayesian models.
NormalNormalDivergence	Kullback-Leibler divergence between two normal distributions.
UniformNormalDivergence	Kullback-Leibler divergence between a uniform and normal distribution.
VIModule	Base class for Modules using Variational Inference.
VIConv1d	Applies a 1D convolution over an input signal composed of several input planes.
VIConv2d	Applies a 2D convolution over an input signal composed of several input planes.
VIConv3d	Applies a 3D convolution over an input signal composed of several input planes.
KullbackLeiblerLoss	Kullback-Leibler (KL) divergence loss.
VILinear	Applies an affine linear transformation to the incoming data: \(y = xA^T + b\).
VIResidualConnection	A version of VISequential that supports residual connections.
VISequential	Sequential container for VIModule.
VIMultiheadAttention	Allows the model to jointly attend to information from different representation subspaces.
VITransformer	A Bayesian Transformer model.
VITransformerDecoder	TransformerDecoder is a stack of N decoder layers.
VITransformerDecoderLayer	TransformerDecoderLayer is made up of self-attn, multi-head-attn and feedforward network.
VITransformerEncoder	TransformerEncoder is a stack of N encoder layers.
VITransformerEncoderLayer	TransformerEncoderLayer is made up of self-attn and feedforward network.
VIkwargs	Common keyword arguments for most VIModules.
VIReturn	A subclass of torch.Tensor that also stores log probabilities.

Functions

convert_to_vimodule(, prior, rescale_prior, ...)

Convert a PyTorch module to a VIModule.

Package Contents

torch_blue.vi.convert_to_vimodule(module: torch.nn.Module, keep_weights: bool = False, variational_distribution: torch_blue.vi.utils.common_types._dist_any_t = MeanFieldNormal(), prior: torch_blue.vi.utils.common_types._dist_any_t = MeanFieldNormal(), rescale_prior: bool = False, kaiming_initialization: bool = True, prior_initialization: bool = False, return_log_probs: bool = True) → None

Convert a PyTorch module to a VIModule.

This method automatically converts a PyTorch module to a VIModule. This also works for any model that is compatible with torch.vmap() (documentation here). Usually, this will be the case if you do not use the +=, -=, *=, and /= operators (their long form, i.e. a = a + b instead of a += b does not cause issues).

To configure the model the usual VIkwargs can be used, except device and dtype which are copied independently for each weight matrix from the original model. This allows to convert a distributed model and maintain the distributed structure. This is compatible with torch.ddp, but DDP needs to be applied after the conversion.

For Bayesian pretraining or custom initialization schemes keep_weights can be set to True, which will maintain the original weights as value for the primary parameter of the weight distribution. While this can vary the primary parameter is typically the distribution mean.

The model is converted inplace so to continue using the original model a copy needs to be made before conversion.

A good method to verify correct model conversion is to set variational distribution and prior to NonBayesian and keep_weights to True. This should make the converted model produce the same results as the original. (The output will have an additional sampling dimension and multiple copies of the same result unless you pass samples=1 to the forward call)

Standard PyTorch layers will automatically be converted to the optimized implementation of this library. Therefore, you should try to use class names that already exist in PyTorch (like “Transformer”). If you cannot do this, you can use ban_torch_convert() to add classes to a banlist (or later remove them from it) that will make them be converted normally.

Advanced note: Auto-conversion will ignore many standard modules from PyTorch since they do not have weights (making conversion irrelevant) or should not be converted (this mostly applies to norm layers).

While not recommended you can use convert_norms() to enable or disable conversion of all PyTorch norms.

Furthermore, you can add (or remove) any class to the conversion ban list with ban_convert(). By default, it will add the class to the global banlist so all instances of this class are left non-Bayesian. There are several more specific behaviors that can be disabled for specific classes with this method using the keyword argument ban_mode. Note, that it only affects the specific class not any subclasses, so if you wish to stop conversion of only one specific layer you can implement it as a subclass with different name but all methods unchanged.

Auto-conversion will try to reuse auto-converted classes. This means if you implemented a custom layer type and used it multiple times all instances will still be instances of the same (converted) class. To disable this behavior you can use the keyword argument ban_mode=reuse.

If there is a torch_blue implementation of a PyTorch layer auto-convert will use that instead of an automatically created class. Note, that this will cause problems if you name custom classes the same as PyTorch classes. To disable this behavior you can ban the replacement of the custom class with ban_mode=replace.

Finally, you can stop conversion of a module and all of its submodules with ban_mode=submodule.

Parameters:

• module (nn.Module) – The Pytorch module to convert.

• keep_weights (bool) – If True keep the original weights as value for the primary parameter.

• VIkwargs – Several standard keyword arguments. See VIkwargs for details. device and dtype cannot be used since they are copied from the original weights.