Have a question? Connect with the community at the TensorFlow Forum

# tfp.experimental.nn.AffineVariationalFlipout

Densely-connected layer class with Flipout estimator.

Inherits From: `VariationalLayer`, `Layer`

This layer implements the Bayesian variational inference analogue to a dense layer by assuming the `kernel` and/or the `bias` are drawn from distributions. By default, the layer implements a stochastic forward pass via sampling from the kernel and bias posteriors,

``````kernel, bias ~ posterior
outputs = tf.linalg.matmul(inputs, kernel) + bias
``````

It uses the Flipout estimator [(Wen et al., 2018)], which performs a Monte Carlo approximation of the distribution integrating over the `kernel` and `bias`. Flipout uses roughly twice as many floating point operations as the reparameterization estimator but has the advantage of significantly lower variance.

The arguments permit separate specification of the surrogate posterior (`q(W|x)`), prior (`p(W)`), and divergence for both the `kernel` and `bias` distributions.

Upon being built, this layer adds losses (accessible via the `losses` property) representing the divergences of `kernel` and/or `bias` surrogate posteriors and their respective priors. When doing minibatch stochastic optimization, make sure to scale this loss such that it is applied just once per epoch (e.g. if `kl` is the sum of `losses` for each element of the batch, you should pass `kl / num_examples_per_epoch` to your optimizer).

#### Examples

We illustrate a Bayesian neural network with variational inference, assuming a dataset of images and length-10 one-hot `targets`.

``````# Using the following substitution, see:
tfn = tfp.experimental.nn
help(tfn.AffineVariationalReparameterization)
BayesAffine = tfn.AffineVariationalFlipout
``````

This example uses reparameterization gradients to minimize the Kullback-Leibler divergence up to a constant, also known as the negative Evidence Lower Bound. It consists of the sum of two terms: the expected negative log-likelihood, which we approximate via Monte Carlo; and the KL divergence, which is added via regularizer terms which are arguments to the layer.

#### References

: Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse. Flipout: Efficient Pseudo-Independent Weight Perturbations on Mini-Batches. In International Conference on Learning Representations, 2018. https://arxiv.org/abs/1803.04386

`input_size` ...
`output_size` ...
`init_kernel_fn` ... Default value: `None` (i.e., `tfp.experimental.nn.initializers.glorot_uniform()`).
`init_bias_fn` ... Default value: `None` (i.e., `tf.initializers.zeros()`).
`make_posterior_fn` ... Default value: `tfp.experimental.nn.util.make_kernel_bias_posterior_mvn_diag`.
`make_prior_fn` ... Default value: `tfp.experimental.nn.util.make_kernel_bias_prior_spike_and_slab`.
`posterior_value_fn` ... Default valye: `tfd.Distribution.sample`
`unpack_weights_fn` Default value: `unpack_kernel_and_bias`
`dtype` ... Default value: `tf.float32`.
`penalty_weight` ... Default value: `None` (i.e., weight is `1`).
`posterior_penalty_fn` ... Default value: `kl_divergence_monte_carlo`.
`activation_fn` ... Default value: `None`.
`seed` ... Default value: `None` (i.e., no seed).
`name` ... Default value: `None` (i.e., `'AffineVariationalFlipout'`).

`activation_fn`

`also_track`

`dtype`

`extra_loss`

`extra_result`

`name` Returns the name of this module as passed or determined in the ctor.

`name_scope` Returns a `tf.name_scope` instance for this class.
`non_trainable_variables` Sequence of non-trainable variables owned by this module and its submodules.
`penalty_weight`

`posterior`

`posterior_penalty_fn`

`posterior_value_fn`

`prior`

`submodules` Sequence of all sub-modules.

Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on).

````a = tf.Module()`
`b = tf.Module()`
`c = tf.Module()`
`a.b = b`
`b.c = c`
`list(a.submodules) == [b, c]`
`True`
`list(b.submodules) == [c]`
`True`
`list(c.submodules) == []`
`True`
```

`trainable_variables` Sequence of trainable variables owned by this module and its submodules.

`unpack_weights_fn`

`variables` Sequence of variables owned by this module and its submodules.

## Methods

View source

View source

View source

View source

### `with_name_scope`

Decorator to automatically enter the module name scope.

````class MyModule(tf.Module):`
`  @tf.Module.with_name_scope`
`  def __call__(self, x):`
`    if not hasattr(self, 'w'):`
`      self.w = tf.Variable(tf.random.normal([x.shape, 3]))`
`    return tf.matmul(x, self.w)`
```

Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose names included the module name:

````mod = MyModule()`
`mod(tf.ones([1, 2]))`
`<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>`
`mod.w`
`<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,`
`numpy=..., dtype=float32)>`
```

Args
`method` The method to wrap.

Returns
The original method wrapped such that it enters the module's name scope.

### `__call__`

View source

Call self as a function.

[{ "type": "thumb-down", "id": "missingTheInformationINeed", "label":"Missing the information I need" },{ "type": "thumb-down", "id": "tooComplicatedTooManySteps", "label":"Too complicated / too many steps" },{ "type": "thumb-down", "id": "outOfDate", "label":"Out of date" },{ "type": "thumb-down", "id": "samplesCodeIssue", "label":"Samples / code issue" },{ "type": "thumb-down", "id": "otherDown", "label":"Other" }]
[{ "type": "thumb-up", "id": "easyToUnderstand", "label":"Easy to understand" },{ "type": "thumb-up", "id": "solvedMyProblem", "label":"Solved my problem" },{ "type": "thumb-up", "id": "otherUp", "label":"Other" }]