def _fn(kernel_size, bias_size, dtype=None):
smallconst = np.log(np.expm1(1.))
n_weights_block = kernel_size//C
n_bias_block = bias_size//C
n_weight_mean_params = n_weights_block
n_weight_cov_params = tfp.layers.MultivariateNormalTriL.params_size(n_weights_block) - n_weights_block
n_params_total = C*(n_weight_mean_params + n_weight_cov_params) + bias_size
#print("{} params in total".format(n_params_total))
block_param_indices = tf.split(np.arange(n_params_total - bias_size), C)
split_array = [n_weight_mean_params, n_weight_cov_params]
split_param_idxs = [tf.split(x, split_array, axis=0) for x in block_param_indices]
model = tf.keras.Sequential([
tfpl.VariableLayer(n_params_total, dtype=dtype),
tfpl.DistributionLambda(lambda t: tfd.Blockwise(
[
tfd.MultivariateNormalTriL(
loc=tf.gather(t,split_param_idxs[c][0], axis=-1),
scale_tril=tfp.math.fill_triangular(
1e-5 + tf.nn.softplus(smallconst + tf.gather(t,split_param_idxs[c][1], axis=-1)))
) for c in range(C)
] +
[ tfd.VectorDeterministic(loc=t[...,-bias_size:]) ]
)
)
])
return model
def test_doc_string(self):
# Generate data -- as in Figure 1 in [Papamakarios et al. (2017)][1]).
n = 2000
x2 = np.random.randn(n) * 2.
x1 = np.random.randn(n) + (x2 * x2 / 4.)
data = np.stack([x1, x2], axis=-1)
# Density estimation with MADE.
model = tfk.Sequential([
# NOTE: This model takes no input and outputs a Distribution. (We use
# the batch_size and type of the input, but there are no actual input
# values because the last dimension of the shape is 0.)
#
# For conditional density estimation, the model would take the
# conditioning values as input.)
tfkl.InputLayer(input_shape=(0, ), dtype=tf.float32),
# Given the empty input, return a standard normal distribution with
# matching batch_shape and event_shape of [2].
# pylint: disable=g-long-lambda
tfpl.DistributionLambda(lambda t: tfd.MultivariateNormalDiag(
loc=tf.zeros(tf.concat([tf.shape(input=t)[:-1], [2]], axis=0)),
scale_diag=[1., 1.])),
# Transform the standard normal distribution with event_shape of [2] to
# the target distribution with event_shape of [2].
tfpl.AutoregressiveTransform(
tfb.AutoregressiveLayer(params=2,
hidden_units=[10],
activation='relu')),
])
model.compile(optimizer=tf.optimizers.Adam(),
loss=lambda y, rv_y: -rv_y.log_prob(y))
model.fit(
x=np.zeros((n, 0)),
y=data,
batch_size=25,
epochs=1,
steps_per_epoch=1, # Usually n // 25,
verbose=True)
distribution = model(np.zeros((0, )))
self.assertEqual((4, 2), self.evaluate(distribution.sample(4)).shape)
self.assertEqual(
(5, 3),
self.evaluate(
distribution.log_prob(np.zeros((5, 3, 2),
dtype=np.float32))).shape)
def _fn(kernel_size, bias_size, dtype=None):
smallconst = np.log(np.expm1(1.))
n = kernel_size + bias_size
c = np.log(np.expm1(1.))
model = tf.keras.Sequential([
tfp.layers.VariableLayer(2*kernel_size + bias_size, dtype=dtype),
tfpl.DistributionLambda(lambda t: tfd.Blockwise(
[
tfd.MultivariateNormalDiag(
loc=t[...,:kernel_size],
scale_diag=1e-5 + tf.nn.softplus(smallconst + t[...,kernel_size:-bias_size])),
tfd.VectorDeterministic(loc=t[...,-bias_size:])
]
)
)
])
return model