def prior(kernel_size, bias_size=0, dtype=None):
from tensorflow_probability import layers
from tensorflow_probability import distributions as tfd
import numpy as np
import tensorflow as tf
n = kernel_size + bias_size
c = np.log(np.expm1(1.))
return tf.keras.Sequential([
layers.VariableLayer(n, dtype=dtype),
layers.DistributionLambda(lambda t: tfd.Independent(
tfd.Normal(loc=t, scale=1.), reinterpreted_batch_ndims=1
)), #[:n]#1e-5 + tf.math.softplus(c + t[n:])
])
def posterior_mean_field(kernel_size, bias_size=0, dtype=None):
"""Posterior function for variational layer."""
n = kernel_size + bias_size
c = np.log(np.expm1(1e-5))
variable_layer = tfpl.VariableLayer(
2 * n,
dtype=dtype,
initializer=tfpl.BlockwiseInitializer([
tf.keras.initializers.TruncatedNormal(mean=0., stddev=.05, seed=None),
tf.keras.initializers.Constant(np.log(np.expm1(1e-5)))],
sizes=[n, n]))
def distribution_fn(t):
scale = 1e-5 + tf.nn.softplus(c + t[Ellipsis, n:])
return tfd.Independent(tfd.Normal(loc=t[Ellipsis, :n], scale=scale),
reinterpreted_batch_ndims=1)
distribution_layer = tfpl.DistributionLambda(distribution_fn)
return tf.keras.Sequential([variable_layer, distribution_layer])
def prior_trainable(kernel_size, bias_size=0, dtype=None, num_updates=1):
"""Prior function for variational layer."""
n = kernel_size + bias_size
c = np.log(np.expm1(1e-5))
def regularizer(t):
out = tfd.LogNormal(0., 1.).log_prob(1e-5 + tf.nn.softplus(c + t[Ellipsis, -1]))
return -tf.reduce_sum(out) / num_updates
# Include the prior on the scale parameter as a regularizer in the loss.
variable_layer = tfpl.VariableLayer(n, dtype=dtype, regularizer=regularizer)
def distribution_fn(t):
scale = 1e-5 + tf.nn.softplus(c + t[Ellipsis, -1])
return tfd.Independent(tfd.Normal(loc=t[Ellipsis, :n], scale=scale),
reinterpreted_batch_ndims=1)
distribution_layer = tfpl.DistributionLambda(distribution_fn)
return tf.keras.Sequential([variable_layer, distribution_layer])
