fastbnns.bnn package

Submodules

fastbnns.bnn.base module

Bayesian neural network base module(s) and utilities.

class fastbnns.bnn.base.BNN(nn: torch.nn.Module | laplace.DiagLaplace, convert_in_place: bool = False, *args, **kwargs)

Bases: Module

Bayesian neural network base class.

forward(input: MuVar | Tensor, *args, **kwargs) → Any

Forward pass through BNN.

laplace_init(dataloader: DataLoader, laplace_kwargs: dict = {}, laplace_prior_opt_kwargs: dict = {}) → None

Initialize model parameter variances to Laplace approximated values.

named_parameters_tagged(tag: str) → Iterator

Return named parameters whose name contains tag.

fastbnns.bnn.base.bnn_params_from_laplace(laplace_model: laplace.DiagLaplace) → dict

Create dictionary of parameters for a BNN from a diagonal Laplace approximation.

Parameters:

laplace_model – Diagonal Laplace approximation instance whose parameters will be reorganized for ingestion into a BNN instance.

fastbnns.bnn.inference module

Inference modules for Bayesian neural network layers.

Custom propagators for specific layers (e.g., “Linear” for Bayesian analog of torch.nn.Linear) should share a name with the layer such that getattr(inference, layer.__class__.__name__) will return the custom propagator for that layer if available.

class fastbnns.bnn.inference.AvgPool1d

Bases: AvgPoolNd

Deterministic moment propagation of mean and variance through AvgPool1d layers.

functional()

avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True) -> Tensor

Applies a 1D average pooling over an input signal composed of several input planes.

See AvgPool1d for details and output shape.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iW)\)

• kernel_size – the size of the window. Can be a single number or a tuple (kW,)

• stride – the stride of the window. Can be a single number or a tuple (sW,). Default: kernel_size

• padding – implicit zero paddings on both sides of the input. Can be a single number or a tuple (padW,). Should be at most half of effective kernel size, that is \(((kernelSize - 1) * dilation + 1) / 2\). Default: 0

• ceil_mode – when True, will use ceil instead of floor to compute the output shape. Default: False

• count_include_pad – when True, will include the zero-padding in the averaging calculation. Default: True

Examples:

>>> # pool of square window of size=3, stride=2
>>> input = torch.tensor([[[1, 2, 3, 4, 5, 6, 7]]], dtype=torch.float32)
>>> F.avg_pool1d(input, kernel_size=3, stride=2)
tensor([[[ 2.,  4.,  6.]]])

n_dim = 1

class fastbnns.bnn.inference.AvgPool2d

Bases: AvgPoolNd

Deterministic moment propagation of mean and variance through AvgPool2d layers.

functional()

avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor

Applies 2D average-pooling operation in \(kH \times kW\) regions by step size \(sH \times sW\) steps. The number of output features is equal to the number of input planes.

See AvgPool2d for details and output shape.

Parameters:

• input – input tensor \((\text{minibatch} , \text{in\_channels} , iH , iW)\)

• kernel_size – size of the pooling region. Can be a single number, a single-element tuple or a tuple (kH, kW)

• stride – stride of the pooling operation. Can be a single number, a single-element tuple or a tuple (sH, sW). Default: kernel_size

• padding – implicit zero paddings on both sides of the input. Can be a single number, a single-element tuple or a tuple (padH, padW). Should be at most half of effective kernel size, that is \(((kernelSize - 1) * dilation + 1) / 2\). Default: 0

• ceil_mode – when True, will use ceil instead of floor in the formula to compute the output shape. Default: False

• count_include_pad – when True, will include the zero-padding in the averaging calculation. Default: True

• divisor_override – if specified, it will be used as divisor, otherwise size of the pooling region will be used. Default: None

n_dim = 2

class fastbnns.bnn.inference.AvgPool3d

Bases: AvgPoolNd

Deterministic moment propagation of mean and variance through AvgPool3d layers.

functional()

avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None) -> Tensor

Applies 3D average-pooling operation in \(kT \times kH \times kW\) regions by step size \(sT \times sH \times sW\) steps. The number of output features is equal to \(\lfloor\frac{\text{input planes}}{sT}\rfloor\).

See AvgPool3d for details and output shape.

Parameters:

• input – input tensor \((\text{minibatch} , \text{in\_channels} , iT \times iH , iW)\)

• kernel_size – size of the pooling region. Can be a single number or a tuple (kT, kH, kW)

• stride – stride of the pooling operation. Can be a single number or a tuple (sT, sH, sW). Default: kernel_size

• padding – implicit zero paddings on both sides of the input. Can be a single number or a tuple (padT, padH, padW). Should be at most half of effective kernel size, that is \(((kernelSize - 1) * dilation + 1) / 2\). Default: 0

• ceil_mode – when True, will use ceil instead of floor in the formula to compute the output shape

• count_include_pad – when True, will include the zero-padding in the averaging calculation

• divisor_override – if specified, it will be used as divisor, otherwise size of the pooling region will be used. Default: None

n_dim = 3

class fastbnns.bnn.inference.AvgPoolNd

Bases: MomentPropagator

Deterministic moment propagation of mean and variance through AvgPoolNd layers.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Analytical moment propagation through layer.

class fastbnns.bnn.inference.BasicPropagator

Bases: MomentPropagator

Propagate mean and variance through module.

This propagator can be used with modules that have “simple” forward passes for which propagation rules are already defined by methods in types.MuVar (e.g., forward pass is just input*param1 + param2).

forward(module: Callable, input: Iterable[Tensor]) → Iterable[Tensor]

Propagate moments by relying intrinsically on methods in types.MuVar.

class fastbnns.bnn.inference.Conv1d

Bases: ConvNd

Deterministic moment propagation of mean and variance through a Conv1d layer.

The internal logic is identical to ConvNd so we just create this class for compatibility with module-name-based searches.

functional()

conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor

Applies a 1D convolution over an input signal composed of several input planes.

This operator supports TensorFloat32.

See Conv1d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Note

This operator supports complex data types i.e. complex32, complex64, complex128.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iW)\)

• weight – filters of shape \((\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kW)\)

• bias – optional bias of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a one-element tuple (sW,). Default: 1

• padding –

  implicit paddings on both sides of the input. Can be a string {‘valid’, ‘same’}, single number or a one-element tuple (padW,). Default: 0 padding='valid' is the same as no padding. padding='same' pads the input so the output has the same shape as the input. However, this mode doesn’t support any stride values other than 1.

  Warning

  For padding='same', if the weight is even-length and dilation is odd in any dimension, a full pad() operation may be needed internally. Lowering performance.

• dilation – the spacing between kernel elements. Can be a single number or a one-element tuple (dW,). Default: 1

• groups – split input into groups, \(\text{in\_channels}\) should be divisible by the number of groups. Default: 1

Examples:

>>> inputs = torch.randn(33, 16, 30)
>>> filters = torch.randn(20, 16, 5)
>>> F.conv1d(inputs, filters)

class fastbnns.bnn.inference.Conv2d

Bases: ConvNd

Deterministic moment propagation of mean and variance through a Conv2d layer.

The internal logic is identical to ConvNd so we just create this class for compatibility with module-name-based searches.

functional()

conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor

Applies a 2D convolution over an input image composed of several input planes.

This operator supports TensorFloat32.

See Conv2d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Note

This operator supports complex data types i.e. complex32, complex64, complex128.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iH , iW)\)

• weight – filters of shape \((\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kH , kW)\)

• bias – optional bias tensor of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sH, sW). Default: 1

• padding –

  implicit paddings on both sides of the input. Can be a string {‘valid’, ‘same’}, single number or a tuple (padH, padW). Default: 0 padding='valid' is the same as no padding. padding='same' pads the input so the output has the same shape as the input. However, this mode doesn’t support any stride values other than 1.

  Warning

  For padding='same', if the weight is even-length and dilation is odd in any dimension, a full pad() operation may be needed internally. Lowering performance.

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dH, dW). Default: 1

• groups – split input into groups, both \(\text{in\_channels}\) and \(\text{out\_channels}\) should be divisible by the number of groups. Default: 1

Examples:

>>> # With square kernels and equal stride
>>> filters = torch.randn(8, 4, 3, 3)
>>> inputs = torch.randn(1, 4, 5, 5)
>>> F.conv2d(inputs, filters, padding=1)

class fastbnns.bnn.inference.Conv3d

Bases: ConvNd

Deterministic moment propagation of mean and variance through a Conv3d layer.

The internal logic is identical to ConvNd so we just create this class for compatibility with module-name-based searches.

functional()

conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor

Applies a 3D convolution over an input image composed of several input planes.

This operator supports TensorFloat32.

See Conv3d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Note

This operator supports complex data types i.e. complex32, complex64, complex128.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iT , iH , iW)\)

• weight – filters of shape \((\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kT , kH , kW)\)

• bias – optional bias tensor of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sT, sH, sW). Default: 1

• padding –

  implicit paddings on both sides of the input. Can be a string {‘valid’, ‘same’}, single number or a tuple (padT, padH, padW). Default: 0 padding='valid' is the same as no padding. padding='same' pads the input so the output has the same shape as the input. However, this mode doesn’t support any stride values other than 1.

  Warning

  For padding='same', if the weight is even-length and dilation is odd in any dimension, a full pad() operation may be needed internally. Lowering performance.

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dT, dH, dW). Default: 1

• groups – split input into groups, \(\text{in\_channels}\) should be divisible by the number of groups. Default: 1

Examples:

>>> filters = torch.randn(33, 16, 3, 3, 3)
>>> inputs = torch.randn(20, 16, 50, 10, 20)
>>> F.conv3d(inputs, filters)

class fastbnns.bnn.inference.ConvNd

Bases: MomentPropagator

Deterministic moment propagation of mean and variance through a ConvNd layer.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Analytical moment propagation through layer.

class fastbnns.bnn.inference.ConvTranspose1d

Bases: ConvTransposeNd

Deterministic moment propagation of mean and variance through a ConvTranspose1d layer.

functional()

conv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor

Applies a 1D transposed convolution operator over an input signal composed of several input planes, sometimes also called “deconvolution”.

This operator supports TensorFloat32.

See ConvTranspose1d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iW)\)

• weight – filters of shape \((\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kW)\)

• bias – optional bias of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sW,). Default: 1

• padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padW,). Default: 0

• output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padW). Default: 0

• groups – split input into groups, \(\text{in\_channels}\) should be divisible by the number of groups. Default: 1

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dW,). Default: 1

Examples:

>>> inputs = torch.randn(20, 16, 50)
>>> weights = torch.randn(16, 33, 5)
>>> F.conv_transpose1d(inputs, weights)

num_spatial_dims = 1

class fastbnns.bnn.inference.ConvTranspose2d

Bases: ConvTransposeNd

Deterministic moment propagation of mean and variance through a ConvTranspose2d layer.

functional()

conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor

Applies a 2D transposed convolution operator over an input image composed of several input planes, sometimes also called “deconvolution”.

This operator supports TensorFloat32.

See ConvTranspose2d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iH , iW)\)

• weight – filters of shape \((\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kH , kW)\)

• bias – optional bias of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sH, sW). Default: 1

• padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padH, padW). Default: 0

• output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padH, out_padW). Default: 0

• groups – split input into groups, \(\text{in\_channels}\) should be divisible by the number of groups. Default: 1

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dH, dW). Default: 1

Examples:

>>> # With square kernels and equal stride
>>> inputs = torch.randn(1, 4, 5, 5)
>>> weights = torch.randn(4, 8, 3, 3)
>>> F.conv_transpose2d(inputs, weights, padding=1)

num_spatial_dims = 2

class fastbnns.bnn.inference.ConvTranspose3d

Bases: ConvTransposeNd

Deterministic moment propagation of mean and variance through a ConvTranspose3d layer.

functional()

conv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1) -> Tensor

Applies a 3D transposed convolution operator over an input image composed of several input planes, sometimes also called “deconvolution”

This operator supports TensorFloat32.

See ConvTranspose3d for details and output shape.

Note

In some circumstances when given tensors on a CUDA device and using CuDNN, this operator may select a nondeterministic algorithm to increase performance. If this is undesirable, you can try to make the operation deterministic (potentially at a performance cost) by setting torch.backends.cudnn.deterministic = True. See /notes/randomness for more information.

Parameters:

• input – input tensor of shape \((\text{minibatch} , \text{in\_channels} , iT , iH , iW)\)

• weight – filters of shape \((\text{in\_channels} , \frac{\text{out\_channels}}{\text{groups}} , kT , kH , kW)\)

• bias – optional bias of shape \((\text{out\_channels})\). Default: None

• stride – the stride of the convolving kernel. Can be a single number or a tuple (sT, sH, sW). Default: 1

• padding – dilation * (kernel_size - 1) - padding zero-padding will be added to both sides of each dimension in the input. Can be a single number or a tuple (padT, padH, padW). Default: 0

• output_padding – additional size added to one side of each dimension in the output shape. Can be a single number or a tuple (out_padT, out_padH, out_padW). Default: 0

• groups – split input into groups, \(\text{in\_channels}\) should be divisible by the number of groups. Default: 1

• dilation – the spacing between kernel elements. Can be a single number or a tuple (dT, dH, dW). Default: 1

Examples:

>>> inputs = torch.randn(20, 16, 50, 10, 20)
>>> weights = torch.randn(16, 33, 3, 3, 3)
>>> F.conv_transpose3d(inputs, weights)

num_spatial_dims = 3

class fastbnns.bnn.inference.ConvTransposeNd

Bases: MomentPropagator

Deterministic moment propagation of mean and variance through a ConvTransposeNd layer.

forward(module: BayesianModule, input: types.MuVar, output_size: List[int] | None = None) → types.MuVar

Analytical moment propagation through layer.

class fastbnns.bnn.inference.JointUnscentedTransform(sigma_scale: float | Tensor | None = None, sigma_weights: list[float] | tuple[float, ...] | Tensor | None = None, outer_product: bool = False)

Bases: MomentPropagator

Unscented transform propagation of mean and variance through Bayesian module.

This propagator uses the unscented transform to propagate mean and variance through a Bayesian layer, jointly applying an unscented transform that accounts for the input distribution and the parameter distributions of the layer.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Forward method for MomentPropagator modules.

Parameters:

• module – Instance of the layer through which we will propagate moments.

• input – Input passed to layer, which will typically be a torch.Tensor or types.MuVar.

class fastbnns.bnn.inference.LeakyReLUa

Bases: MomentPropagator

Deterministic moment propagation of a normal random variable through a leaky-ReLU.

NOTE: the suffix “a” is added to prevent fastbnns.bnn.wrappers.covert_to_bnn_() from automatically selecting this propagator for LeakyReLU. This is currently desirable since using UnscentedTransform is faster (though less accurate) than analytical propagation.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Analytical moment propagation through layer.

class fastbnns.bnn.inference.Linear

Bases: MomentPropagator

Deterministic moment propagation of mean and variance through a Linear layer.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Analytical moment propagation through layer.

functional()

linear(input, weight, bias=None) -> Tensor

Applies a linear transformation to the incoming data: \(y = xA^T + b\).

This operation supports 2-D weight with sparse layout

Warning

Sparse support is a beta feature and some layout(s)/dtype/device combinations may not be supported, or may not have autograd support. If you notice missing functionality please open a feature request.

This operator supports TensorFloat32.

Shape:

• Input: \((*, in\_features)\) where * means any number of additional dimensions, including none

• Weight: \((out\_features, in\_features)\) or \((in\_features)\)

• Bias: \((out\_features)\) or \(()\)

• Output: \((*, out\_features)\) or \((*)\), based on the shape of the weight

class fastbnns.bnn.inference.MomentPropagator

Bases: Module

Base class for layer propagators.

forward(module: Callable, input: Iterable[Tensor]) → Iterable[Tensor]

Forward method for MomentPropagator modules.

Parameters:

• module – Instance of the layer through which we will propagate moments.

• input – Input passed to layer, which will typically be a torch.Tensor or types.MuVar.

class fastbnns.bnn.inference.MonteCarlo(n_samples: int = 10, input_sampler: Distribution = <class 'torch.distributions.normal.Normal'>)

Bases: MomentPropagator

Monte Carlo propagation of mean and variance through module.

This propagator is designed to perform Monte Carlo sampling of BOTH the input and the layer being propagated through. This would generally only be used for Bayesian/stochastic layers through which we cannot propagate moments through analytically.

forward(module: Callable, input: MuVar, return_samples: bool = False) → MuVar

Propagate moments by averaging over n_samples forward passes of module.

class fastbnns.bnn.inference.ReLUa

Bases: MomentPropagator

Deterministic moment propagation of a normal random variable through a ReLU.

NOTE: the suffix “a” is added to prevent fastbnns.bnn.wrappers.covert_to_bnn_() from automatically selecting this propagator for ReLU. This is currently desirable since using UnscentedTransform is faster (though less accurate) than analytical propagation.

forward(module: BayesianModule, input: types.MuVar) → types.MuVar

Analytical moment propagation through layer.

class fastbnns.bnn.inference.UnscentedTransform(sigma_scale: Tensor | None = None, sigma_weights: Tensor | None = None, n_module_samples: int = 1)

Bases: MomentPropagator

Unscented transform propagation of mean and variance through module.

This propagator uses the unscented transform to propagate mean and variance through a deterministic layer.

forward(module: Callable, input: MuVar, return_samples: bool = False) → MuVar

Propagate moments using the unscented transform.

fastbnns.bnn.losses module

Losses and helpers useful for Bayesian neural network training/evaluation.

class fastbnns.bnn.losses.BNNLoss(size_average=None, reduce=None, reduction: str = 'mean')

Bases: ABC, _Loss

Abstract class for ELBO-like losses used to train Bayesian Neural Networks.

abstract property beta: FloatTensor

Scale factor for KL divergence loss term.

abstract property kl_divergence: _Loss

_Loss to compute the KL divergence of a model.

abstract property neg_log_likelihood: _Loss

_Loss to compute the negative log-likelihood term in the ELBO.

class fastbnns.bnn.losses.ELBO(neg_log_likelihood: _Loss | None = None, kl_divergence: _Loss = KLDivergence(), beta: float = 1.0, reduction: str = 'sum')

Bases: BNNLoss

Evidence lower bound with scaled KL.

property beta: bool

Return property beta.

This property is written as an @property method for compatibility with the abstract parent class.

forward(model: Module | None = None, **kwargs) → Tensor

Compute the ELBO loss.

Parameters:

• model – torch.nn.Module that may have some layers.BayesianLayers as sub-modules, for which we’ll compute the KL divergence w.r.t their prior. Passing None is treated as no model, i.e., KL = 0.0

• kwargs – Keyword arguments to pass to self.log_likelihood(**kwargs)

property kl_divergence: bool

Return property kl_divergence.

This property is written as an @property method for compatibility with the abstract parent class.

property neg_log_likelihood: bool

Return property neg_log_likelihood.

This property is written as an @property method for compatibility with the abstract parent class.

class fastbnns.bnn.losses.KLDivergence(prior: dict | Distribution | None = None)

Bases: _Loss

KL divergence loss for Bayesian neural networks.

forward(model: Module) → Tensor

Compute KL divergence for Bayesian sub-modules of model.

Parameters:

model – torch.nn.Module that may have some Bayesian layers as sub-modules, for which we’ll compute the KL divergence w.r.t their prior.

fastbnns.bnn.losses.kl_divergence_sampled(dist0: Distribution, dist1: Distribution, n_samples: int = 1)

KL divergence KL(dist0||dist1) approximated by Monte Carlo sampling.

fastbnns.bnn.priors module

Definitions of prior distributions over neural network parameters.

class fastbnns.bnn.priors.Distribution(distribution: Distribution | None = None, *args: Any, **kwargs: Any)

Bases: Module

Distribution wrapper to facilitate device transfers.

property distribution: Distribution

Prepare an instance of the distribution.

log_prob(x: Tensor) → Tensor

Compute the log PDF of the prior at points x.

sample(sample_shape: Iterable = ()) → Tensor

Generate samples from the prior of size sample_shape.

class fastbnns.bnn.priors.SpikeSlab(loc: Tensor = tensor([0., 0.]), scale: Tensor = tensor([0.1000, 1.0000]), probs: Tensor = tensor([0.5000, 0.5000]))

Bases: Distribution

Spike-slab Gaussian Mixture Model prior.

property distribution: Distribution

Prepare an instance of the distribution.

fastbnns.bnn.types module

Custom types and associated functionality.

class fastbnns.bnn.types.MuVar(mu: Tensor | list[Tensor, Tensor] | tuple[Tensor, Tensor] | MuVar, var: Tensor | None = None)

Bases: object

Custom object holding mean and variance of some distribution.

WARNING: Some functionality, like __pow__(), assumes the normal distribution!

add(input: int | float | Tensor | MuVar) → MuVar

add_(input: int | float | Tensor | MuVar, *, alpha: int = 1) → MuVar

Custom inplace add for MuVar types assuming self and input are independent.

addcmul(tensor1, tensor2, *, value=1) → torch.MuVar

Custom implementation of torch.addcmul.

apply(func: Callable, *args, **kwargs) → MuVar

Generic apply() for functions that act separately on mu and var.

apply_method(name: str, *args, **kwargs) → MuVar

Generic apply_method() for methods that act separately on mu and var.

div(input: int | float | Tensor) → MuVar

mean(*args, **kwargs) → MuVar

Custom replacement of torch.mean() for MuVar type.

mul(input: int | float | Tensor | MuVar) → MuVar

numel() → int

Custom numel() to avoid complicated logic in __getattr__ above.

pow(input: int | float | Tensor) → MuVar

sub(input: int | float | Tensor | MuVar) → MuVar

sub_(input: int | float | Tensor | MuVar, *, alpha: int = 1) → MuVar

Custom inplace sub for MuVar types assuming self and input are independent.

fastbnns.bnn.types.implements(*functions: Callable[[...], Any])

Register a custom MuVar implementation for one or more torch functions.

fastbnns.bnn.types.muvar_avg_pool(func, input: Tensor, kernel_size: int | Size | list[int] | tuple[int, ...], stride: int | Size | list[int] | tuple[int, ...] | None = None, padding: int | Size | list[int] | tuple[int, ...] = 0, ceil_mode: bool = False, count_include_pad: bool = True, *args, **kwargs)

fastbnns.bnn.types.muvar_dropout(func, input: Tensor, *args, **kwargs)

fastbnns.bnn.wrappers module

Collections of Bayesian neural network layers.

class fastbnns.bnn.wrappers.BayesianModule(module: Module, samplers: dict | None = None, samplers_init: dict | None = None, resample_mean: bool = True, priors: dict | None = None, moment_propagator: MomentPropagator | None = None, learn_var: bool = True, *args, **kwargs)

Bases: BayesianModuleBase

Base class for BayesianModule modules to make PyTorch modules BNN compatible.

compute_kl_divergence(priors: dict | Distribution | None = None, n_samples: int = 1) → Tensor

Compute the KL divergence between self.prior and module parameters.

Parameters:

• priors – Prior distribution over parameters. This can be a single distribution for all parameters or a dictionary whose keys correspond to named parameters of this layer. By default, None will use self.priors. This argument is used to allow external passing of priors not defined at initialization of this layer.

• n_samples – Number of Monte Carlo samples used to estimate KL divergence for distributions not compatible with torch.nn.distributions.kl_divergence()

forward(input: MuVar | Tensor, *args, **kwargs) → MuVar | Tensor

Forward pass through layer.

get_named_sampler(name: str) → Distribution

Initialize and return the requested module parameter sampler.

property learn_var: bool

Return property learn_var.

This property is written as an @property method for compatibility with the abstract parent class.

property module: Module

Prepare a callable that acts like input module with random parameters.

property module_map: Module

Return module instance with parameters set to learned means.

property module_params: dict

Return a sample of this module’s parameters.

property moment_propagator: bool

Return property moment_propagator.

This property is written as an @property method for compatibility with the abstract parent class.

reset_parameters() → None

Resample layer parameters from initial distributions.

property samplers: dict

Return initialized variational distribution samplers.

class fastbnns.bnn.wrappers.BayesianModuleBase(*args: Any, **kwargs: Any)

Bases: ABC, Module

Abstract base class for Bayesian modules.

abstractmethod compute_kl_divergence(*args, **kwargs) → Tensor

Method that computes the KL divergence of this module w.r.t. some prior.

abstractmethod forward(input: MuVar | Tensor, *args, **kwargs) → MuVar | Tensor

Method that computes a forward pass through this module.

abstract property learn_var: bool

Flag indicating variance of module parameters should be learnable.

abstract property module: Module

Return a sampled instance of this module.

abstract property module_map: Module

Return a non-Bayesian instance of self with parameters set to learned means.

abstract property module_params: ParameterDict

Dictionary organizing parameters of this module.

abstract property moment_propagator: MomentPropagator

Function used to propagate mean and variance through this module.

scale_tform()

softplus(input, beta=1, threshold=20) -> Tensor

Applies element-wise, the function \(\text{Softplus}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x))\).

For numerical stability the implementation reverts to the linear function when \(input \times \beta > threshold\).

See Softplus for more details.

scale_tform_inv() → Tensor

class fastbnns.bnn.wrappers.BroadcastModule(module: Module, *args, **kwargs)

Bases: Module

BroadcastModule for compatibility with other Bayesian layers.

forward(input: MuVar | Tensor, *args, **kwargs) → MuVar | Tensor

Forward pass through layer.

fastbnns.bnn.wrappers.convert_to_bnn_(model: Module, layer_wrappers: dict = {}, layer_wrappers_tag: dict = {}, layer_wrappers_type: dict = {}, wrapper_kwargs: dict = {}, wrapper_kwargs_tag: dict = {}, wrapper_kwargs_type: dict = {}, wrapper_kwargs_global: dict = {}) → None

Convert layers of model to Bayesian counterparts.

Parameters:

• model – Model to be converted to Bayesian counterpart.

• layer_wrappers – Dictionary of manually-specified wrappers for specific module layers. The keys are names of leaf modules (e.g., “module1.layer1”) and the values are the corresponding wrapper class present in this module (e.g., “BroadcastModule”).

• layer_wrappers_tag – Extends functionality of layer_wrappers where keys don’t have to be exact layer names but instead can be tags, e.g., {“encoder”: “BroadcastModule”} specifies that any submodule in [name for name, _ in model.named_modules()] satisfying “encoder” in name will be wrapped with a BroadcastModule.

• layer_wrappers_type – Extends functionality of layer_wrappers where keys are now type names, e.g., {“BatchNorm2d”: “BroadcastModule”} specifies that any submodule in [m for m in model.modules()] satisfying “BatchNorm2d” == type(m).__name__” will be wrapped with a BroadcastModule.

• wrapper_kwargs – Additional keyword arguments passed to initialization of named Bayesian layers. For example, if model has a module named “module1”, we’ll convert “module1” as Converter(module1, **wrapper_kwargs[“module1”]) where Converter is a module converter.

• wrapper_kwargs_tag – Extends functionality of wrapper_kwargs where keys don’t have to be exact layer names but instead can be tags, e.g., {“encoder”: encoder_kwargs} specifies that any submodule in [name for name, _ in model.named_modules()] satisfying “encoder” in name will use encoder_kwargs for their wrapper.

• wrapper_kwargs_type – Extends functionality of wrapper_kwargs where keys are now type names, e.g., {“BatchNorm2d”: batchnorm_kwargs} specifies that any submodule in [m for m in model.modules()] satisfying “BatchNorm2d” == type(m).__name__” will use batchnorm_kwargs for their wrapper.

• wrapper_kwargs_global – Keyword arguments that we’ll merge with values of wrapper_kwargs, wrapper_kwargs_tag, and wrapper_kwargs_type as appropriate for each module.

fastbnns.bnn.wrappers.convert_to_nn(bnn: Module) → None

Inverse of convert_to_bnn_ to convert a BNN back to a standard NN

Parameters:

model – Bayesian NN to be converted back to a standard NN.

fastbnns.bnn.wrappers.isolate_leaf_module_names(module_names: list[str]) → list[str]

Prepare a list of leaf modules of model.

This function filters module_names to eliminate the names of parent modules. For example, if we have a model: torch.nn.Module with named modules m = [“”, “module1”, “module2”, “module1.submodule”, “module2.submodule”], isolate_leaf_module_names(m) == [“module1.submodule”, “module2.submodule”]

fastbnns.bnn.wrappers.select_default_propagator(module: Module, is_bayesian: bool = True) → MomentPropagator

Select a compatible moment propagator for module.

Parameters:

• module – Module that we’ll choose a propagator for.

• is_bayesian – Flag indicating parameters of module will be treated as distributions.

Module contents