def p_net(observed=None, n_z=None, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = functools.partial(spt.layers.act_norm,
initializing=is_initializing)
# sample z ~ p(z)
z = net.add('z',
spt.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([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
weight_norm=True,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_z = z
h_z = spt.layers.dense(h_z, 500)
h_z = spt.layers.dense(h_z, 500)
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_z, config.x_dim, name='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_y=None, n_z=None, n_samples=None):
if n_samples is not None:
warnings.warn('`n_samples` is deprecated, use `n_y` instead.')
n_y = n_samples
net = spt.BayesianNet(observed=observed)
# sample y
y = net.add('y',
spt.Categorical(tf.zeros([1, config.n_clusters])),
n_samples=n_y)
# sample z ~ p(z|y)
z = net.add('z',
gaussian_mixture_prior(y, config.z_dim, config.n_clusters),
group_ndims=1,
n_samples=n_z,
is_reparameterized=False)
# compute the hidden features for x
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_x = z
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_x, config.x_dim, name='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z', spt.Bernoulli(tf.zeros([1, config.z_dim])),
group_ndims=1, n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
h_z = tf.to_float(z)
h_z = spt.layers.dense(h_z, 500)
h_z = spt.layers.dense(h_z, 500)
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_z, config.x_dim, name='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_z=None, is_training=False, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = None if not config.act_norm else functools.partial(
spt.layers.act_norm,
axis=-1 if config.channels_last else -3,
initializing=is_initializing,
value_ndims=3,
)
# sample z ~ p(z)
z = net.add('z',
spt.Normal(mean=tf.zeros([1, config.z_dim]),
std=tf.ones([1, config.z_dim]) *
config.truncated_sigma),
group_ndims=1,
n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.resnet_deconv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=config.channels_last):
h_z = spt.layers.dense(z, 64 * 7 * 7)
h_z = spt.ops.reshape_tail(h_z,
ndims=1,
shape=(7, 7,
64) if config.channels_last else
(64, 7, 7))
h_z = spt.layers.resnet_deconv2d_block(h_z, 64) # output: (64, 7, 7)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 32, strides=2) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(h_z, 32) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 16, strides=2) # output: (16, 28, 28)
# sample x ~ p(x|z)
x_logits = spt.layers.conv2d(
h_z,
1, (1, 1),
padding='same',
name='feature_map_to_pixel',
channels_last=config.channels_last) # output: (1, 28, 28)
x = net.add('x',
spt.Bernoulli(logits=x_logits, dtype=tf.float32),
group_ndims=3)
return net
def q_net(x, observed=None, n_z=None):
net = spt.BayesianNet(observed=observed)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg)):
h_x = tf.to_float(x)
h_x = spt.layers.dense(h_x, 500)
h_x = spt.layers.dense(h_x, 500)
# sample z ~ q(z|x)
z_logits = spt.layers.dense(h_x, config.z_dim, name='z_logits')
z = net.add('z', spt.Bernoulli(logits=z_logits), n_samples=n_z,
group_ndims=1)
return net
def p_net(observed=None, n_z=None, is_initializing=False):
net = spt.BayesianNet(observed=observed)
normalizer_fn = functools.partial(spt.layers.act_norm,
initializing=is_initializing)
# sample z ~ p(z)
def make_component(i):
normal = spt.Normal(mean=tf.get_variable('mean_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True),
logstd=tf.maximum(
tf.get_variable('logstd_{}'.format(i),
shape=[1, config.z_dim],
dtype=tf.float32,
trainable=True),
config.z_logstd_min))
return normal.expand_value_ndims(1)
components = [
make_component(i) for i in range(config.n_mixture_components)
]
mixture = spt.Mixture(categorical=spt.Categorical(
logits=tf.zeros([1, config.n_mixture_components])),
components=components,
is_reparameterized=True)
z = net.add('z', mixture, n_samples=n_z)
# compute the hidden features
with arg_scope([spt.layers.dense],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=normalizer_fn,
weight_norm=True,
kernel_regularizer=spt.layers.l2_regularizer(
config.l2_reg)):
h_z = z
h_z = spt.layers.dense(h_z, 500)
h_z = spt.layers.dense(h_z, 500)
# sample x ~ p(x|z)
x_logits = spt.layers.dense(h_z, config.x_dim, name='x_logits')
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_z=None, is_training=True, channels_last=True):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
spt.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([spt.layers.resnet_deconv2d_block],
kernel_size=config.kernel_size,
shortcut_kernel_size=config.shortcut_kernel_size,
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=spt.layers.l2_regularizer(config.l2_reg),
channels_last=channels_last):
h_z = spt.layers.dense(z, 64 * 7 * 7)
h_z = spt.utils.reshape_tail(
h_z, ndims=1, shape=[7, 7, 64] if channels_last else [64, 7, 7])
h_z = spt.layers.resnet_deconv2d_block(h_z, 64) # output: (64, 7, 7)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 32, strides=2) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(h_z, 32) # output: (32, 14, 14)
h_z = spt.layers.resnet_deconv2d_block(
h_z, 16, strides=2) # output: (16, 28, 28)
# sample x ~ p(x|z)
h_z = spt.layers.conv2d(h_z,
1, (1, 1),
padding='same',
name='feature_map_to_pixel',
channels_last=channels_last) # output: (1, 28, 28)
x_logits = spt.utils.reshape_tail(h_z, 3, [config.x_dim])
x = net.add('x', spt.Bernoulli(logits=x_logits), group_ndims=1)
return net
def p_net(observed=None, n_z=None, beta=1.0, mcmc_iterator=0, log_Z=0.0, initial_z=None,
mcmc_alpha=config.smallest_step):
net = spt.BayesianNet(observed=observed)
# sample z ~ p(z)
normal = spt.Normal(mean=tf.zeros([1, config.z_dim]),
logstd=tf.zeros([1, config.z_dim]))
normal = normal.batch_ndims_to_value(1)
xi = tf.get_variable(name='xi', shape=(), initializer=tf.constant_initializer(config.initial_xi),
dtype=tf.float32, trainable=True)
# xi = tf.square(xi)
xi = tf.nn.sigmoid(xi) # TODO
pz = EnergyDistribution(normal, G=G_theta, D=D_psi, log_Z=log_Z, xi=xi, mcmc_iterator=mcmc_iterator,
initial_z=initial_z, mcmc_alpha=mcmc_alpha)
z = net.add('z', pz, n_samples=n_z)
x_mean = G_theta(z)
x_mean = tf.clip_by_value(x_mean, 1e-7, 1 - 1e-7)
logits = tf.log(x_mean) - tf.log1p(-x_mean)
bernouli = spt.Bernoulli(
logits=logits, dtype=tf.float32
)
# bernouli.mean = x_mean
x = net.add('x', bernouli, group_ndims=3)
return net