def h_for_q_z(x, is_training, channels_last):
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)
return {
'mean': dense(h_x, config.z_dim, name='z_mean'),
'logstd': dense(h_x, config.z_dim, name='z_logstd'),
}
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 baseline_net(x):
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)
return tf.squeeze(dense(h_x, 1), -1)
def reinforce_baseline_net(config, x):
x, s1, s2 = flatten(tf.to_float(x), 2)
with arg_scope([dense],
kernel_regularizer=l2_regularizer(config.l2_reg),
activation_fn=tf.nn.leaky_relu):
h_x = dense(x, 500)
h_x = unflatten(tf.reshape(dense(h_x, 1), [-1]), s1, s2)
return h_x
def model(x, is_training):
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_x = x
h_x = dense(h_x, 500)
h_x = dense(h_x, 500)
logits = dense(h_x, 10, name='logits')
return logits
def h_for_p_x(z, is_training):
with arg_scope([dense],
activation_fn=tf.nn.leaky_relu,
kernel_regularizer=l2_regularizer(config.l2_reg)):
h_z, s1, s2 = flatten(z, 2)
h_z = dense(h_z, 500)
h_z = dense(h_z, 500)
return {
'logits': unflatten(dense(h_z, config.x_dim, name='x_logits'), s1, s2)
}
def h_for_q_z(x, is_training):
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)
return {
'logits': dense(h_x, config.z_dim, name='z_logits'),
}
def p_net(observed=None,
n_y=None,
n_z=None,
tau=None,
is_training=True,
n_samples=None):
if n_samples is not None:
warnings.warn('`n_samples` is deprecated, use `n_y` instead.')
n_y = n_samples
use_concrete = config.use_concrete_distribution and tau is not None
logging.info('p_net builder: %r', locals())
net = BayesianNet(observed=observed)
# sample y
if use_concrete:
y = net.add('y',
ExpConcrete(tau, tf.zeros([1, config.n_clusters])),
n_samples=n_y,
is_reparameterized=True)
else:
y = net.add('y',
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,
use_concrete=use_concrete),
group_ndims=1,
n_samples=n_z,
is_reparameterized=use_concrete)
# compute the hidden features for x
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 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_logits = dense(h_x, config.z_dim, name='z_logits')
z = net.add('z', Bernoulli(logits=z_logits), n_samples=n_z, group_ndims=1)
return net
def model(config, x, is_training, channels_last):
with arg_scope([resnet_block],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=functools.partial(
batch_norm_2d,
channels_last=channels_last,
training=is_training,
),
dropout_fn=functools.partial(
tf.layers.dropout,
training=is_training
),
kernel_regularizer=l2_regularizer(config.l2_reg),
channels_last=channels_last):
h_x = tf.reshape(
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 = global_average_pooling(
h_x, channels_last=channels_last) # output: (64, 1, 1)
h_x = tf.reshape(h_x, [-1, 64])
logits = dense(h_x, 10, name='logits')
return logits
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,
flow=posterior_flow())
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
def p_net(config, observed=None, n_z=None, is_training=True):
net = BayesianNet(observed=observed)
# sample z ~ p(z)
z = net.add('z',
Bernoulli(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)):
z = tf.to_float(z)
h_z, s1, s2 = flatten(z, 2)
h_z = dense(h_z, 500)
h_z = dense(h_z, 500)
# sample x ~ p(x|z)
x_logits = unflatten(dense(h_z, config.x_dim, name='x_logits'), s1, s2)
x = net.add('x', Bernoulli(logits=x_logits), group_ndims=1)
return net
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 h_for_p_x(z, is_training, channels_last):
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])
x_logits = unflatten(h_z, s1, s2)
return {'logits': x_logits}
def model(x, is_training, channels_last, k=4, n=2):
with arg_scope([resnet_block],
activation_fn=tf.nn.leaky_relu,
normalizer_fn=functools.partial(
batch_norm_2d,
channels_last=channels_last,
training=is_training,
),
dropout_fn=functools.partial(
tf.layers.dropout,
rate=config.dropout,
training=is_training
),
kernel_regularizer=l2_regularizer(config.l2_reg),
channels_last=channels_last):
if not channels_last:
h_x = x
else:
h_x = tf.transpose(x, [-1, 2, 3, 1])
h_x = conv2d(h_x, 16 * k, (1, 1), channels_last=channels_last)
# 1st group, (16 * k, 32, 32)
for i in range(n):
h_x = resnet_block(h_x, 16 * k)
# 2nd group, (32 * k, 16, 16)
h_x = resnet_block(h_x, 32 * k, strides=2)
for i in range(n):
h_x = resnet_block(h_x, 32 * k)
# 3rd group, (64 * k, 8, 8)
h_x = resnet_block(h_x, 64 * k, strides=2)
for i in range(n):
h_x = resnet_block(h_x, 64 * k)
h_x = global_average_pooling(
h_x, channels_last=channels_last) # output: (64 * k, 1, 1)
h_x = tf.reshape(h_x, [-1, 64 * k])
logits = dense(h_x, 10, name='logits')
return logits
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
