def testKLDivergenceIsZero(self):
"""Tests that KL divergence of identical distributions is zero."""
with self.session() as sess:
mean = tf.random.normal((3, 3, 3))
std = tf.random.normal((3, 3, 3))
kl_divergence = dynamics.KLDivergence()([mean, std, mean, std])
result = sess.run([kl_divergence])[0]
np.testing.assert_array_equal(result, result * 0.0)
def testKLDivergenceAnnealing(self):
inputs = tf.keras.Input(1)
outputs = dynamics.KLDivergence(kl_annealing_steps=4)([
inputs, inputs,
tf.keras.layers.Lambda(lambda x: x + 1.0)(inputs), inputs
])
model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.compile('sgd', 'mse')
obtained_kl = [model.predict(x=[1], steps=1)] # Should be zero before fit.
for _ in range(5):
model.fit(x=[1], y=[1], epochs=1, steps_per_epoch=1)
obtained_kl.append(model.predict(x=[1], steps=1))
obtained_kl = np.array(obtained_kl).ravel()
np.testing.assert_array_almost_equal(
obtained_kl, [0, 0.125, 0.25, 0.375, 0.5, 0.5])
def testNonzeroKLDivergence(self):
"""Test that KL divergence layer provides correct result."""
mu = 2.0
sigma = 2.0
n = 3
# https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Multivariate_normal_distributions
result_np = 0.5 * (sigma ** 2 + mu ** 2 - np.log(sigma ** 2) - 1) * n
with self.session() as sess:
kl_divergence = dynamics.KLDivergence()(
[tf.zeros(n), tf.ones(n), tf.zeros(n) + mu, tf.ones(n) * sigma])
result_tf = sess.run(kl_divergence)
np.testing.assert_almost_equal(result_tf, result_np, decimal=4)