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