Google I/O is a wrap! Catch up on TensorFlow sessions

# 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.

: 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: ... kernel_initializer: ... Default value: `None` (i.e., `tfp.experimental.nn.initializers.glorot_uniform()`). bias_initializer: ... 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`. activation_fn: ... Default value: `None`.
`seed` ... Default value: `None` (i.e., no seed).
`validate_args` ... Default value: `False`. name: ... Default value: `None` (i.e., `'AffineVariationalFlipout'`).

`activation_fn`

`also_track`

`dtype`

`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.
`posterior`

`posterior_value`

`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`

`validate_args` Python `bool` indicating possibly expensive checks are enabled.
`variables` Sequence of variables owned by this module and its submodules.

## Methods

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.

[]
[]