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