def add_p_z_given_y_reg_loss(loss):
if not config.p_z_given_y_reg:
return loss
n_z = config.p_z_given_y_reg_samples
y = tf.range(config.n_clusters, dtype=tf.int32)
prior = gaussian_mixture_prior(y=y,
z_dim=config.z_dim,
n_clusters=config.n_clusters,
use_concrete=False)
prior0 = Normal(mean=0., logstd=0.)
z = prior.sample(n_z, is_reparameterized=True)
# Note: p(y) = 1/n_clusters, which simplies the
# following deduction.
if config.p_z_given_y_reg == 'kl_p_z_given_y':
# :math:`E_{y}[KL(p(z|y)\|p_0(z))] =
# E_{y,z \sim p(z|y)}[\log p(z|y) - \log p_0(z)]`
reg_loss = tf.reduce_mean(
prior.log_prob(z, group_ndims=1) -
prior0.log_prob(z, group_ndims=1))
elif config.p_z_given_y_reg == 'kl_p_z':
# :math:`KL(E_{y}(p(z|y))\|p_0(z)) =
# E_{y,z \sim p(z|y)}[\log \E_{y}(p(z|y)) - \log p_0(z)]`
log_p_z = log_mean_exp(tfops,
prior.log_prob(z, group_ndims=1),
axis=-1)
log_p0_z = prior0.log_prob(z, group_ndims=1)
reg_loss = tf.reduce_mean(log_p_z) - tf.reduce_mean(log_p0_z)
else:
raise ValueError(
'Unexpected value for config `p_z_given_y_reg`: {}'.format(
config.p_z_given_y_reg))
return loss + config.p_z_given_y_reg_factor * reg_loss
def test_local_log_prob(self):
x_observed = np.arange(24, dtype=np.float32).reshape([2, 3, 4])
net = BayesianNet({'x': x_observed})
normal = Normal(tf.zeros([3, 4]), tf.ones([3, 4]))
x = net.add('x', normal)
y = net.add('y', normal)
# test single query
x_log_prob = net.local_log_prob('x')
self.assertIsInstance(x_log_prob, tf.Tensor)
with self.test_session():
np.testing.assert_allclose(x_log_prob.eval(),
normal.log_prob(x_observed).eval())
# test multiple query
x_log_prob, y_log_prob = net.local_log_probs(iter(['x', 'y']))
self.assertIsInstance(x_log_prob, tf.Tensor)
self.assertIsInstance(y_log_prob, tf.Tensor)
with self.test_session() as sess:
np.testing.assert_allclose(x_log_prob.eval(),
normal.log_prob(x_observed).eval())
x_log_prob_val, x_log_prob_res, y_log_prob_val, y_log_prob_res = \
sess.run([
x_log_prob, normal.log_prob(x.tensor),
y_log_prob, normal.log_prob(y.tensor),
])
np.testing.assert_allclose(x_log_prob_val, x_log_prob_res)
np.testing.assert_allclose(y_log_prob_val, y_log_prob_res)
def test_log_prob_value_and_group_ndims(self):
tf.set_random_seed(123456)
mean = tf.constant([0., 1., 2.], dtype=tf.float64)
normal = Normal(mean=mean, std=tf.constant(1., dtype=tf.float64))
y = tf.random_normal(shape=[2, 5, 3], dtype=tf.float64)
with self.test_session() as sess:
# test value_ndims = 0, group_ndims = 1
flow = QuadraticFlow(2., 5.)
flow.build(tf.zeros([2, 5, 3], dtype=tf.float64))
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 0)
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (2, 5, 3))
self.assertTupleEqual(get_static_shape(log_det), (2, 5, 3))
log_py = tf.reduce_sum(normal.log_prob(x) + log_det, axis=-1)
np.testing.assert_allclose(
*sess.run([distrib.log_prob(y, group_ndims=1), log_py]),
rtol=1e-5
)
# test value_ndims = 1, group_ndims = 0
flow = QuadraticFlow(2., 5., value_ndims=1)
flow.build(tf.zeros([2, 5, 3], dtype=tf.float64))
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 1)
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (2, 5, 3))
self.assertTupleEqual(get_static_shape(log_det), (2, 5))
log_py = normal.log_prob(x, group_ndims=1) + log_det
np.testing.assert_allclose(
*sess.run([distrib.log_prob(y, group_ndims=0), log_py]),
rtol=1e-5
)
# test value_ndims = 1, group_ndims = 2
flow = QuadraticFlow(2., 5., value_ndims=1)
flow.build(tf.zeros([2, 5, 3], dtype=tf.float64))
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 1)
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (2, 5, 3))
self.assertTupleEqual(get_static_shape(log_det), (2, 5))
log_py = tf.reduce_sum(
log_det + tf.reduce_sum(normal.log_prob(x), axis=-1))
np.testing.assert_allclose(
*sess.run([distrib.log_prob(y, group_ndims=2), log_py]),
rtol=1e-5
)
def test_sample_value_and_group_ndims(self):
tf.set_random_seed(123456)
mean = tf.constant([0., 1., 2.], dtype=tf.float64)
normal = Normal(mean=mean, std=tf.constant(1., dtype=tf.float64))
with self.test_session() as sess:
# test value_ndims = 0, group_ndims = 1
flow = QuadraticFlow(2., 5.)
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 0)
y = distrib.sample(n_samples=5, group_ndims=1)
self.assertTupleEqual(get_static_shape(y), (5, 3))
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (5, 3))
self.assertTupleEqual(get_static_shape(log_det), (5, 3))
log_py = tf.reduce_sum(normal.log_prob(x) + log_det, axis=-1)
np.testing.assert_allclose(*sess.run([y.log_prob(), log_py]),
rtol=1e-5)
# test value_ndims = 1, group_ndims = 0
flow = QuadraticFlow(2., 5., value_ndims=1)
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 1)
y = distrib.sample(n_samples=5, group_ndims=0)
self.assertTupleEqual(get_static_shape(y), (5, 3))
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (5, 3))
self.assertTupleEqual(get_static_shape(log_det), (5,))
log_py = log_det + tf.reduce_sum(normal.log_prob(x), axis=-1)
np.testing.assert_allclose(*sess.run([y.log_prob(), log_py]),
rtol=1e-5)
# test value_ndims = 1, group_ndims = 1
flow = QuadraticFlow(2., 5., value_ndims=1)
distrib = FlowDistribution(normal, flow)
self.assertEqual(distrib.value_ndims, 1)
y = distrib.sample(n_samples=5, group_ndims=1)
self.assertTupleEqual(get_static_shape(y), (5, 3))
x, log_det = flow.inverse_transform(y)
self.assertTupleEqual(get_static_shape(x), (5, 3))
self.assertTupleEqual(get_static_shape(log_det), (5,))
log_py = tf.reduce_sum(
log_det + tf.reduce_sum(normal.log_prob(x), axis=-1))
np.testing.assert_allclose(*sess.run([y.log_prob(), log_py]),
rtol=1e-5)
def test_sample(self):
tf.set_random_seed(123456)
mean = tf.constant([0., 1., 2.], dtype=tf.float64)
normal = Normal(mean=mean, std=tf.constant(1., dtype=tf.float64))
flow = QuadraticFlow(2., 5.)
distrib = FlowDistribution(normal, flow)
# test ordinary sample, is_reparameterized = None
y = distrib.sample(n_samples=5)
self.assertTrue(y.is_reparameterized)
grad = tf.gradients(y * 1., mean)[0]
self.assertIsNotNone(grad)
self.assertEqual(get_static_shape(y), (5, 3))
self.assertIsNotNone(y._self_log_prob)
x, log_det = flow.inverse_transform(y)
log_py = normal.log_prob(x) + log_det
with self.test_session() as sess:
np.testing.assert_allclose(*sess.run([log_py, y.log_prob()]),
rtol=1e-5)
# test stop gradient sample, is_reparameterized = False
y = distrib.sample(n_samples=5, is_reparameterized=False)
self.assertFalse(y.is_reparameterized)
grad = tf.gradients(y * 1., mean)[0]
self.assertIsNone(grad)
def test_log_prob(self):
mean = tf.constant([0., 1., 2.], dtype=tf.float64)
normal = Normal(mean=mean, std=tf.constant(1., dtype=tf.float64))
flow = QuadraticFlow(2., 5.)
flow.build(tf.constant(0., dtype=tf.float64))
distrib = FlowDistribution(normal, flow)
y = tf.constant([1., -1., 2.], dtype=tf.float64)
x, log_det = flow.inverse_transform(y)
log_py = normal.log_prob(x) + log_det
py = tf.exp(log_py)
log_prob = distrib.log_prob(y)
self.assertIsInstance(log_prob, FlowDistributionDerivedTensor)
self.assertIsInstance(log_prob.flow_origin, StochasticTensor)
self.assertIs(log_prob.flow_origin.distribution, normal)
prob = distrib.prob(y)
self.assertIsInstance(prob, FlowDistributionDerivedTensor)
self.assertIsInstance(prob.flow_origin, StochasticTensor)
self.assertIs(prob.flow_origin.distribution, normal)
with self.test_session() as sess:
np.testing.assert_allclose(
*sess.run([log_prob.flow_origin, x]), rtol=1e-5)
np.testing.assert_allclose(
*sess.run([log_py, log_prob]), rtol=1e-5)
np.testing.assert_allclose(
*sess.run([py, prob]), rtol=1e-5)
def test_query(self):
x_observed = np.arange(24, dtype=np.float32).reshape([2, 3, 4])
net = BayesianNet({'x': x_observed})
normal = Normal(tf.zeros([3, 4]), tf.ones([3, 4]))
x = net.add('x', normal)
y = net.add('y', normal)
[(x_out, x_log_prob), (y_out, y_log_prob)] = net.query(iter(['x', 'y']))
for o in [x_out, x_log_prob, y_out, y_log_prob]:
self.assertIsInstance(o, tf.Tensor)
self.assertIs(x_out, x.tensor)
self.assertIs(y_out, y.tensor)
with self.test_session() as sess:
np.testing.assert_allclose(
x_log_prob.eval(), normal.log_prob(x_observed).eval())
x_log_prob_val, x_log_prob_res, y_log_prob_val, y_log_prob_res = \
sess.run([
x_log_prob, normal.log_prob(x.tensor),
y_log_prob, normal.log_prob(y.tensor),
])
np.testing.assert_allclose(x_log_prob_val, x_log_prob_res)
np.testing.assert_allclose(y_log_prob_val, y_log_prob_res)
def test_log_prob(self):
mean = tf.constant([0., 1., 2.], dtype=tf.float64)
normal = Normal(mean=mean, std=tf.constant(1., dtype=tf.float64))
flow = QuadraticFlow(2., 5.)
flow.build(tf.constant(0., dtype=tf.float64))
distrib = FlowDistribution(normal, flow)
y = tf.constant([1., -1., 2.], dtype=tf.float64)
x, log_det = flow.inverse_transform(y)
log_py = normal.log_prob(x) + log_det
with self.test_session() as sess:
np.testing.assert_allclose(*sess.run([log_py,
distrib.log_prob(y)]),
rtol=1e-5)
def test_basic(self):
x = Normal(mean=[0., 1.], logstd=0.).sample(10, group_ndims=1)
log_p = x.log_prob()
self.assertListEqual([10, 2], x.get_shape().as_list())
self.assertListEqual([10], log_p.get_shape().as_list())
transformed = x * 2.
transformed_log_p = log_p * .5
d = TransformedDistribution(x,
transformed,
transformed_log_p,
is_reparameterized=True,
is_continuous=False)
self.assertIs(x, d.origin)
self.assertIs(transformed, d.transformed)
self.assertIs(transformed_log_p, d.transformed_log_p)
self.assertEquals(x.dtype, d.dtype)
self.assertTrue(d.is_reparameterized)
self.assertFalse(d.is_continuous)
self.assertIs(transformed_log_p, d.log_prob(transformed,
group_ndims=1))
with self.test_session() as sess:
np.testing.assert_allclose(*sess.run(
[tf.exp(log_p * .5),
d.prob(transformed, group_ndims=1)]))
with_error = lambda: pytest.raises(
ValueError,
match='`given` must be `self.transformed` and '
'`group_ndims` must be `self.origin.group_'
'ndims`.')
with with_error():
_ = d.log_prob(tf.constant(1.), group_ndims=1)
with with_error():
_ = d.log_prob(transformed, group_ndims=0)
with with_error():
_ = d.prob(tf.constant(1.), group_ndims=1)
with with_error():
_ = d.prob(transformed, group_ndims=0)