zuko.bayesian

Bayesian models.

Classes

BayesianModel

Creates a Bayesian wrapper around a base model.

Descriptions

class zuko.bayesian.BayesianModel(base, init_logvar=-9.0, include_params=('',), exclude_params=())

Creates a Bayesian wrapper around a base model.

Parameters:

• base (nn.Module) – A base model.

• include_params (Sequence[str]) – A list of parameter name prefixes to include. In a prefix, a single star * matches alphanumeric strings ([a-zA-Z0-9_]+) and a double star ** matches dot-separated alphanumeric strings ([a-zA-Z0-9_.]+). By default, all parameters are included.

• exclude_params (Sequence[str]) – A list of parameter name prefixes to exclude.

sample_params()

Returns model parameters sampled from the posterior.

sample_model()

Returns a model sampled from the posterior.

The returned model is a standalone copy of the base model where parameters have been replaced. It can be used exactly in the same ways as the base model.

Warning

This method should not be used for training as it cannot propagate gradients to the Bayesian model’s internal parameters.

reparameterize(local_trick=False)

Reparameterizes the base model from the posterior.

Within this context, the base model is temporarily reparametrized and behaves deterministically, meaning that two calls to a method with the same inputs, leads to the same outputs.

Parameters:

local_trick (bool) – Whether to use the local reparameterization trick for linear layers.

Yields:

The reparametrized base model.

References

Variational Dropout and the Local Reparameterization Trick (Kingma et al., 2015)

https://arxiv.org/abs/1506.02557

kl_divergence(prior_var=1.0)

Returns the KL divergence between the posterior and the prior.