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_property(self):
normal = Normal(mean=tf.constant([0., 1., 2.], dtype=tf.float64),
std=tf.constant(1., dtype=tf.float64))
flow = QuadraticFlow(2., 5., value_ndims=1)
distrib = FlowDistribution(normal, flow)
self.assertIs(distrib.flow, flow)
self.assertIs(distrib.base_distribution, normal)
self.assertEqual(distrib.dtype, tf.float64)
self.assertTrue(distrib.is_continuous)
self.assertTrue(distrib.is_reparameterized)
self.assertEqual(distrib.value_ndims, 1)
# self.assertEqual(distrib.get_value_shape(), normal.get_value_shape())
# self.assertEqual(distrib.get_batch_shape(), normal.get_batch_shape())
# with self.test_session() as sess:
# np.testing.assert_equal(
# *sess.run([distrib.value_shape, normal.value_shape]))
# np.testing.assert_equal(
# *sess.run([distrib.batch_shape, normal.batch_shape]))
# test is_reparameterized = False
normal = Normal(mean=[0., 1., 2.], std=1., is_reparameterized=False)
distrib = FlowDistribution(normal, flow)
self.assertFalse(distrib.is_reparameterized)
# test y_value_ndims = 2
distrib = FlowDistribution(normal, ReshapeFlow(1, [-1, 1]))
self.assertEqual(distrib.value_ndims, 2)
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_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_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_add(self):
x_observed = np.arange(24, dtype=np.float32).reshape([2, 3, 4])
net = BayesianNet({'x': x_observed})
self.assertNotIn('x', net)
self.assertNotIn('y', net)
# add an observed node
x = net.add('x',
Normal(tf.zeros([3, 4]), tf.ones([3, 4])),
n_samples=2,
group_ndims=1,
is_reparameterized=False)
self.assertIs(net.get('x'), x)
self.assertIs(net['x'], x)
self.assertIn('x', net)
self.assertListEqual(list(net), ['x'])
self.assertIsInstance(x, StochasticTensor)
self.assertEqual(x.n_samples, 2)
self.assertEqual(x.group_ndims, 1)
self.assertEqual(x.is_reparameterized, False)
with self.test_session():
np.testing.assert_allclose(x.eval(), x_observed)
np.testing.assert_equal(tf.shape(x).eval(), [2, 3, 4])
# add an unobserved node
y = net.add('y',
Normal(tf.zeros([3, 4]), tf.ones([3, 4])),
n_samples=2,
group_ndims=1,
is_reparameterized=False)
self.assertIs(net.get('y'), y)
self.assertIs(net['y'], y)
self.assertIn('y', net)
self.assertListEqual(list(net), ['x', 'y'])
self.assertIsInstance(y, StochasticTensor)
self.assertEqual(y.n_samples, 2)
self.assertEqual(y.group_ndims, 1)
self.assertEqual(y.is_reparameterized, False)
with self.test_session():
np.testing.assert_equal(tf.shape(y).eval(), [2, 3, 4])
# error adding non-string name
with pytest.raises(TypeError, match='`name` must be a str'):
_ = net.add(1, Normal(0., 1.))
# error adding the same variable
with pytest.raises(
KeyError,
match='StochasticTensor with name \'x\' already exists '
'in the BayesianNet. Names must be unique.'):
_ = net.add('x', Normal(2., 3.))
# default is_reparameterized of Normal
z = net.add('z', Normal(0., 1.))
self.assertTrue(z.is_reparameterized)
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_variational_chain(self):
q_net = BayesianNet({'x': [1.]})
q_net.add('z', Normal(q_net.observed['x'], 1.))
q_net.add('y', Normal(q_net.observed['x'] * 2, 2.))
def model_builder(observed):
model = BayesianNet(observed)
z = model.add('z', Normal([0.], [1.]))
y = model.add('y', Normal([0.], [2.]))
x = model.add('x', Normal(z + y, [1.]))
return model
model_builder = Mock(wraps=model_builder)
# test chain with default parameters
chain = q_net.variational_chain(model_builder)
self.assertEqual(
model_builder.call_args,
(({'y': q_net['y'].tensor, 'z': q_net['z'].tensor},),)
)
self.assertEqual(chain.latent_names, ('z', 'y'))
self.assertIsNone(chain.latent_axis)
# test chain with latent_names
chain = q_net.variational_chain(model_builder, latent_names=['y'])
self.assertEqual(
model_builder.call_args,
(({'y': q_net['y'].tensor},),)
)
self.assertEqual(chain.latent_names, ('y',))
# test chain with latent_axis
chain = q_net.variational_chain(model_builder, latent_axis=-1)
self.assertEqual(chain.latent_axis, -1)
# test chain with observed
chain = q_net.variational_chain(model_builder, observed=q_net.observed)
self.assertEqual(
model_builder.call_args,
(({'x': q_net.observed['x'], 'y': q_net['y'].tensor,
'z': q_net['z'].tensor},),)
)
self.assertEqual(chain.latent_names, ('z', 'y'))
# test model_builder with log_joint
def model_builder_1(observed):
return model_builder(observed), fake_log_joint
fake_log_joint = tf.constant(0.)
chain = q_net.variational_chain(model_builder_1)
with self.test_session():
np.testing.assert_equal(
chain.log_joint.eval(), fake_log_joint.eval())
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_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_errors(self):
# errors in constructor
normal = Normal(mean=tf.zeros([3]), std=1.)
with pytest.raises(TypeError, match='`flow` is not an instance of '
'`BaseFlow`: 123'):
_ = FlowDistribution(normal, 123)
flow = QuadraticFlow(2., 5., value_ndims=1)
with pytest.raises(ValueError,
match='cannot be transformed by a flow, because '
'it is not continuous'):
_ = FlowDistribution(Categorical(logits=[0., 1., 2.]), flow)
with pytest.raises(ValueError,
match='cannot be transformed by a flow, because '
'its data type is not float'):
_ = FlowDistribution(Mock(normal, dtype=tf.int32), flow)
with pytest.raises(ValueError,
match='cannot be transformed by flow .*, because '
'distribution.value_ndims is larger than '
'flow.x_value_ndims'):
_ = FlowDistribution(Mock(normal, value_ndims=2), flow)
# errors in sample
distrib = FlowDistribution(normal, flow)
with pytest.raises(RuntimeError,
match='`FlowDistribution` requires `compute_prob` '
'not to be False'):
_ = distrib.sample(compute_density=False)
def p_net(observed=None, n_z=None, is_training=True):
logging.info('p_net builder: %r', locals())
net = BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_x, s1, s2 = flatten(z, 2)
h_x = dense(h_x, 500)
h_x = dense(h_x, 500)
# sample x ~ p(x|z)
x_logits = unflatten(dense(h_x, config.x_dim, name='x_logits'), s1, s2)
x = net.add('x', Bernoulli(logits=x_logits), group_ndims=1)
return net
def gaussian_mixture_prior(config, y, z_dim, n_clusters):
# derive the learnt z_mean
prior_mean = tf.get_variable('z_prior_mean',
dtype=tf.float32,
shape=[n_clusters, z_dim],
initializer=tf.random_normal_initializer())
z_mean = tf.nn.embedding_lookup(prior_mean, y)
# derive the learnt z_std
z_logstd = z_std = None
if config.p_z_given_y_std == 'one':
z_logstd = tf.zeros_like(z_mean)
else:
prior_std_or_logstd = tf.get_variable(
'z_prior_std_or_logstd',
dtype=tf.float32,
shape=[n_clusters, z_dim],
initializer=tf.zeros_initializer())
z_std_or_logstd = tf.nn.embedding_lookup(prior_std_or_logstd, y)
if config.p_z_given_y_std == 'one_plus_softplus_std':
z_std = 1. + tf.nn.softplus(z_std_or_logstd)
elif config.p_z_given_y_std == 'softplus_logstd':
z_logstd = tf.nn.softplus(z_std_or_logstd)
elif config.p_z_given_y_std == 'unbound_logstd':
z_logstd = z_std_or_logstd
else:
raise ValueError(
'Unexpected value for config `p_z_given_y_std`: {}'.format(
config.p_z_given_y_std))
return Normal(mean=z_mean, std=z_std, logstd=z_logstd)
def test_add_with_flow(self):
normal = Normal(mean=tf.constant([0., 1., 2.]), std=1.)
flow = QuadraticFlow(2., 5.)
# test add with sample
net = BayesianNet()
x = net.add('x', normal, flow=flow)
self.assertIsInstance(x.distribution, FlowDistribution)
self.assertIs(x.distribution.flow, flow)
# ensure non-invertible flow cannot be added with observed var
class _Flow(BaseFlow):
@property
def explicitly_invertible(self):
return False
net = BayesianNet({'x': tf.zeros([5, 3])})
with pytest.raises(TypeError,
match='The observed variable \'x\' expects `flow` '
'to be explicitly invertible, but it is not'):
_ = net.add('x', normal, flow=_Flow(x_value_ndims=0))
# test add observed with flow
x = net.add('x', normal, flow=flow)
self.assertIsInstance(x.distribution, FlowDistribution)
self.assertIs(x.distribution.flow, flow)
def test_validate_distribution_factory(self):
# test distribution class
factory = validate_distribution_factory(Normal, 'xyz')
self.assertIsInstance(factory, DistributionFactory)
self.assertIs(factory.distribution_class, Normal)
self.assertEqual(factory.default_args, {})
# test zhusuan distribution class
factory = validate_distribution_factory(zd.Normal, 'xyz')
self.assertIsInstance(factory, DistributionFactory)
self.assertIs(factory.distribution_class, zd.Normal)
self.assertEqual(factory.default_args, {})
# test distribution factory
factory_0 = Normal.factory(mean=0.)
factory = validate_distribution_factory(factory_0, 'xyz')
self.assertIs(factory, factory_0)
# test error
with pytest.raises(TypeError,
match='xyz must be a subclass of `Distribution` or '
'`zhusuan.distributions.Distribution`, or '
'an instance of `DistributionFactory`'):
_ = validate_distribution_factory(object, 'xyz')
with pytest.raises(TypeError,
match='xyz must be a subclass of `Distribution` or '
'`zhusuan.distributions.Distribution`, or '
'an instance of `DistributionFactory`'):
_ = validate_distribution_factory(object(), 'xyz')
def test_validate_names(self):
net = BayesianNet({'x': [2., 3., 4.]})
x = net.add('x', Normal(0., 1.))
y = net.add('y', Normal(0., 1.))
for meth in ['output', 'local_log_prob']:
with pytest.raises(TypeError, match='`names` is not a list of str'):
_ = getattr(net, meth)(1)
with pytest.raises(KeyError, match='StochasticTensor with name '
'\'z\' does not exist'):
_ = getattr(net, meth)('z')
for meth in ['outputs', 'local_log_probs', 'query']:
with pytest.raises(TypeError, match='`names` is not a list of str'):
_ = getattr(net, meth)([1, 2])
with pytest.raises(KeyError, match='StochasticTensor with name '
'\'z\' does not exist'):
_ = getattr(net, meth)(['x', 'y', 'z'])
def test_add_transform(self):
class PatchedNormal(Normal):
def sample(self,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
name=None):
return StochasticTensor(
self,
x_samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
net = BayesianNet({'w': tf.constant(0.)})
x_samples = tf.reshape(tf.range(24, dtype=tf.float32), [2, 3, 4])
normal = PatchedNormal(tf.zeros([3, 4]), tf.ones([3, 4]))
# test success call
x = net.add('x',
normal,
n_samples=2,
group_ndims=1,
transform=lambda x, log_p: (x * 2., log_p * .5))
self.assertIsInstance(x.distribution, TransformedDistribution)
self.assertEquals(1, x.group_ndims)
self.assertEquals(2, x.n_samples)
self.assertTrue(x.is_reparameterized)
with self.test_session() as sess:
np.testing.assert_allclose(*sess.run([x_samples * 2., x]))
np.testing.assert_allclose(*sess.run(
[normal.log_prob(x_samples, group_ndims=1) * .5,
x.log_prob()]))
# test errors
I = lambda x, log_p: (x, log_p)
with pytest.raises(TypeError,
match='Cannot add `TransformedDistribution`'):
_ = net.add('y', x.distribution)
with pytest.raises(ValueError,
match='`transform` can only be applied on '
'continuous, re-parameterized variables'):
_ = net.add('y', Categorical([0.], dtype=tf.int32), transform=I)
with pytest.raises(ValueError,
match='`transform` can only be applied on '
'continuous, re-parameterized variables'):
_ = net.add('y',
ExpConcrete(.5, [0.], is_reparameterized=False),
transform=I)
with pytest.raises(ValueError,
match='`observed` variable cannot be transformed.'):
_ = net.add('w', Normal(mean=0., std=0.), transform=I)
with pytest.raises(ValueError,
match='The transformed samples must be continuous'):
T = lambda x, log_p: (tf.cast(x, dtype=tf.int32), log_p)
_ = net.add('y', normal, transform=T)
def q_net(x, observed=None, n_samples=None, tau=None, is_training=True):
use_concrete = config.use_concrete_distribution and tau is not None
logging.info('q_net builder: %r', locals())
net = BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_x = tf.to_float(x)
h_x = dense(h_x, 500)
h_x = dense(h_x, 500)
# sample y ~ q(y|x)
y_logits = dense(h_x, config.n_clusters, name='y_logits')
if use_concrete:
y = net.add('y',
ExpConcrete(tau, y_logits),
is_reparameterized=True,
n_samples=n_samples)
y_one_hot = tf.exp(y)
else:
y = net.add('y', Categorical(y_logits), n_samples=n_samples)
y_one_hot = tf.one_hot(y, config.n_clusters, dtype=tf.float32)
# sample z ~ q(z|y,x)
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
if config.mean_field_assumption_for_q:
# by mean-field-assumption we let q(z|y,x) = q(z|x)
h_z, s1, s2 = flatten(h_x, 2)
z_n_samples = n_samples
else:
if n_samples is not None:
h_z = tf.concat([
tf.tile(tf.reshape(h_x, [1, -1, 500]),
tf.stack([n_samples, 1, 1])), y_one_hot
],
axis=-1)
else:
h_z = tf.concat([h_x, y_one_hot], axis=-1)
h_z, s1, s2 = flatten(h_z, 2)
h_z = dense(h_z, 500)
z_n_samples = None
z_mean = dense(h_z, config.z_dim, name='z_mean')
z_logstd = dense(h_z, config.z_dim, name='z_logstd')
z = net.add('z',
Normal(mean=unflatten(z_mean, s1, s2),
logstd=unflatten(z_logstd, s1, s2),
is_reparameterized=use_concrete),
n_samples=z_n_samples,
group_ndims=1)
return net
def vae(self, is_reparameterized=None, z_group_ndims=1, x_group_ndims=1):
return VAE(p_z=Normal(mean=tf.zeros([BATCH_SIZE, Z_DIMS]),
std=tf.ones([BATCH_SIZE, Z_DIMS])),
p_x_given_z=Normal,
q_z_given_x=Normal,
h_for_p_x=self.h_for_p_x,
h_for_q_z=self.h_for_q_z,
is_reparameterized=is_reparameterized,
z_group_ndims=z_group_ndims,
x_group_ndims=x_group_ndims)
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)
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 __init__(self,
hidden_net_p_x_z,
hidden_net_q_z_x,
x_dims,
z_dims,
std_epsilon=1e-4,
name=None,
scope=None):
if not is_integer(x_dims) or x_dims <= 0:
raise ValueError('`x_dims`必须为正整数')
if not is_integer(z_dims) or z_dims <= 0:
raise ValueError('`z_dims`必须为正整数')
super(Donut, self).__init__(name=name, scope=scope)
with reopen_variable_scope(self.variable_scope):
# 基于VAE构造
self._vae = VAE(
# p(z)：均值和标准差都为z维数量大小的全零数组的一元正态分布
p_z=Normal(mean=tf.zeros([z_dims]), std=tf.ones([z_dims])),
# p(x|h(z))：一元正态分布
p_x_given_z=Normal,
# q(z|h(x))：一元正态分布
q_z_given_x=Normal,
# p(x|h(z))的隐藏网络：mean、std，由p(x|z)隐藏网络输入获得
h_for_p_x=Lambda(partial(wrap_params_net,
h_for_dist=hidden_net_p_x_z,
mean_layer=partial(tf.layers.dense,
units=x_dims,
name='x_mean'),
std_layer=partial(softplus_std,
units=x_dims,
epsilon=std_epsilon,
name='x_std')),
name='p_x_given_z'),
# q(z|h(x))的隐藏网络：mean、std，由q(z|x)隐藏网络输入获得
h_for_q_z=Lambda(partial(wrap_params_net,
h_for_dist=hidden_net_q_z_x,
mean_layer=partial(tf.layers.dense,
units=z_dims,
name='z_mean'),
std_layer=partial(softplus_std,
units=z_dims,
epsilon=std_epsilon,
name='z_std')),
name='q_z_given_x'))
self._x_dims = x_dims
self._z_dims = z_dims
def __init__(self,
h_for_p_x,
h_for_q_z,
x_dims,
z_dims,
std_epsilon=1e-4,
name=None,
scope=None):
if not is_integer(x_dims) or x_dims <= 0:
raise ValueError('`x_dims` must be a positive integer')
if not is_integer(z_dims) or z_dims <= 0:
raise ValueError('`z_dims` must be a positive integer')
super(Donut, self).__init__(name=name, scope=scope)
with reopen_variable_scope(self.variable_scope):
self._vae = VAE(
p_z=Normal(mean=tf.zeros([z_dims]), std=tf.ones([z_dims])),
p_x_given_z=Normal,
q_z_given_x=Normal,
h_for_p_x=Sequential([
h_for_p_x,
DictMapper(
{
'mean':
K.layers.Dense(x_dims),
'std':
lambda x: (std_epsilon + K.layers.Dense(
x_dims, activation=tf.nn.softplus)(x))
},
name='p_x_given_z')
]),
h_for_q_z=Sequential([
h_for_q_z,
DictMapper(
{
'mean':
K.layers.Dense(z_dims),
'std':
lambda z: (std_epsilon + K.layers.Dense(
z_dims, activation=tf.nn.softplus)(z))
},
name='q_z_given_x')
]),
)
self._x_dims = x_dims
self._z_dims = z_dims
def __init__(self, h_for_p_x, h_for_q_z, x_dims, z_dims, std_epsilon=1e-4,
name=None, scope=None) -> object:
if not is_integer(x_dims) or x_dims <= 0:
raise ValueError('`x_dims` must be a positive integer')
if not is_integer(z_dims) or z_dims <= 0:
raise ValueError('`z_dims` must be a positive integer')
super(Donut, self).__init__(name=name, scope=scope)
with reopen_variable_scope(self.variable_scope):
self._vae = VAE(
p_z=Normal(mean=tf.zeros([z_dims]), std=tf.ones([z_dims])),
p_x_given_z=Normal,
q_z_given_x=Normal,
h_for_p_x=Lambda(
partial(
wrap_params_net,
h_for_dist=h_for_p_x,
mean_layer=partial(
tf.layers.dense, units=x_dims, name='x_mean'
),
std_layer=partial(
softplus_std, units=x_dims, epsilon=std_epsilon,
name='x_std'
)
),
name='p_x_given_z'
),
h_for_q_z=Lambda(
partial(
wrap_params_net,
h_for_dist=h_for_q_z,
mean_layer=partial(
tf.layers.dense, units=z_dims, name='z_mean'
),
std_layer=partial(
softplus_std, units=z_dims, epsilon=std_epsilon,
name='z_std'
)
),
name='q_z_given_x'
)
)
self._x_dims = x_dims
self._z_dims = z_dims
def q_net(config,
x,
observed=None,
n_z=None,
is_training=True,
channels_last=False):
net = BayesianNet(observed=observed)
# compute the hidden features
normalizer_fn = None if not config.batch_norm else functools.partial(
batch_norm_2d,
channels_last=channels_last,
training=is_training,
)
dropout_fn = None if not config.dropout else functools.partial(
tf.layers.dropout, training=is_training)
with arg_scope([resnet_block],
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
dropout_fn=dropout_fn,
kernel_regularizer=l2_regularizer(config.l2_reg),
channels_last=channels_last):
h_x = tf.to_float(x)
h_x = tf.reshape(h_x,
[-1, 28, 28, 1] if channels_last else [-1, 1, 28, 28])
h_x = resnet_block(h_x, 16) # output: (16, 28, 28)
h_x = resnet_block(h_x, 32, strides=2) # output: (32, 14, 14)
h_x = resnet_block(h_x, 32) # output: (32, 14, 14)
h_x = resnet_block(h_x, 64, strides=2) # output: (64, 7, 7)
h_x = resnet_block(h_x, 64) # output: (64, 7, 7)
h_x = reshape_conv2d_to_flat(h_x)
# sample z ~ q(z|x)
z_mean = dense(h_x, config.z_dim, name='z_mean')
z_logstd = dense(h_x, config.z_dim, name='z_logstd')
z = net.add('z',
Normal(mean=z_mean, logstd=z_logstd),
n_samples=n_z,
group_ndims=1)
return net
def q_net(config, x, observed=None, n_z=None, is_training=True):
net = BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_x = tf.to_float(x)
h_x = dense(h_x, 500)
h_x = dense(h_x, 500)
# sample z ~ q(z|x)
z_mean = dense(h_x, config.z_dim, name='z_mean')
z_logstd = dense(h_x, config.z_dim, name='z_logstd')
z = net.add('z',
Normal(mean=z_mean, logstd=z_logstd),
n_samples=n_z,
group_ndims=1)
return net
def test_outputs(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_out = net.output('x')
self.assertIs(x_out, x.tensor)
self.assertIsInstance(x_out, tf.Tensor)
with self.test_session():
np.testing.assert_equal(x_out.eval(), x_observed)
# test multiple query
x_out, y_out = net.outputs(iter(['x', 'y']))
self.assertIs(x_out, x.tensor)
self.assertIs(y_out, y.tensor)
self.assertIsInstance(x_out, tf.Tensor)
self.assertIsInstance(y_out, tf.Tensor)
with self.test_session():
np.testing.assert_equal(x_out.eval(), x_observed)
def test_errors(self):
normal = Normal(mean=0., std=1.)
with pytest.raises(TypeError,
match='`flow` is not an instance of '
'`BaseFlow`: 123'):
_ = FlowDistribution(normal, 123)
flow = QuadraticFlow(2., 5.)
with pytest.raises(ValueError,
match='cannot be transformed by a flow, because '
'it is not continuous'):
_ = FlowDistribution(Categorical(logits=[0., 1., 2.]), flow)
with pytest.raises(ValueError,
match='cannot be transformed by a flow, because '
'its data type is not float'):
_ = FlowDistribution(Mock(normal, dtype=tf.int32), flow)
distrib = FlowDistribution(normal, flow)
with pytest.raises(RuntimeError,
match='`FlowDistribution` requires `compute_prob` '
'not to be False'):
_ = distrib.sample(compute_density=False)
def p_net(config,
observed=None,
n_z=None,
is_training=True,
channels_last=False):
net = BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim])),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([deconv_resnet_block],
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg),
channels_last=channels_last):
h_z, s1, s2 = flatten(z, 2)
h_z = tf.reshape(dense(h_z, 64 * 7 * 7),
[-1, 7, 7, 64] if channels_last else [-1, 64, 7, 7])
h_z = deconv_resnet_block(h_z, 64) # output: (64, 7, 7)
h_z = deconv_resnet_block(h_z, 32, strides=2) # output: (32, 14, 14)
h_z = deconv_resnet_block(h_z, 32) # output: (32, 14, 14)
h_z = deconv_resnet_block(h_z, 16, strides=2) # output: (16, 28, 28)
h_z = conv2d(h_z,
1, (1, 1),
padding='same',
name='feature_map_to_pixel',
channels_last=channels_last) # output: (1, 28, 28)
h_z = tf.reshape(h_z, [-1, config.x_dim])
# sample x ~ p(x|z)
x_logits = unflatten(h_z, s1, s2)
x = net.add('x', Bernoulli(logits=x_logits), group_ndims=1)
return net
def q_net(x, observed=None, n_z=None, is_training=True):
logging.info('q_net builder: %r', locals())
net = BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_z = tf.to_float(x)
h_z = dense(h_z, 500)
h_z = dense(h_z, 500)
# sample z ~ q(z|x)
z_mean = tf.layers.dense(h_z, config.z_dim, name='z_mean')
z_logstd = tf.layers.dense(h_z, config.z_dim, name='z_logstd')
z = net.add('z',
Normal(mean=z_mean, logstd=z_logstd),
n_samples=n_z,
group_ndims=1,
transform=posterior_flow)
return net