def discriminator(image, labels, df_dim, number_classes, act=tf.nn.relu):
"""Builds the discriminator graph.
Args:
image: The current batch of images to classify as fake or real.
labels: The corresponding labels for the images.
df_dim: The df dimension.
number_classes: The number of classes in the labels.
act: The activation function used in the discriminator.
Returns:
- A `Tensor` representing the logits of the discriminator.
- A list containing all trainable varaibles defined by the model.
"""
with tf.variable_scope('discriminator', reuse=tf.AUTO_REUSE) as dis_scope:
h0 = optimized_block(image, df_dim, 'd_optimized_block1',
act=act) # 64 * 64
h1 = block(h0, df_dim * 2, 'd_block2', act=act) # 32 * 32
h1 = ops.sn_non_local_block_sim(h1, name='d_ops') # 32 * 32
h2 = block(h1, df_dim * 4, 'd_block3', act=act) # 16 * 16
h3 = block(h2, df_dim * 8, 'd_block4', act=act) # 8 * 8
h4 = block(h3, df_dim * 16, 'd_block5', act=act) # 4 * 4
h5 = block(h4, df_dim * 16, 'd_block6', downsample=False, act=act)
h5_act = act(h5)
h6 = tf.reduce_sum(input_tensor=h5_act, axis=[1, 2])
output = ops.snlinear(h6, 1, name='d_sn_linear')
h_labels = ops.sn_embedding(labels,
number_classes,
df_dim * 16,
name='d_embedding')
output += tf.reduce_sum(input_tensor=h6 * h_labels,
axis=1,
keepdims=True)
var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
dis_scope.name)
return output, var_list
def test_sn_non_local_block_sim_shapes(self):
"""Tests that downsampling has the desired effect on shape."""
if tf.executing_eagerly():
# `compute_spectral_norm` doesn't work when executing eagerly.
return
image = tf.random.normal([10, 8, 8, 64])
big_image = ops.sn_non_local_block_sim(image, name='test_sa')
self.assertEqual([10, 8, 8, 64], big_image.shape.as_list())
def biggan_generator_128(z, target_class, gf_dim, num_classes, training=True):
"""...
TODO(ilyak): Fix, this does not work as is.
y is embedded, and skip concatenate with z
4th block has attention (64x64 resolution)
batch norm is conditional batch norm
no layer_norm
"""
# setables
embed_y_dim = 128
embed_bias = False
with tf.compat.v1.variable_scope(
'generator', reuse=tf.compat.v1.AUTO_REUSE) as gen_scope:
num_blocks = 5
# embedding of y that is shared
target_class_onehot = tf.one_hot(target_class, num_classes)
y = ops.linear(target_class_onehot,
embed_y_dim,
use_bias=embed_bias,
name="embed_y")
y_per_block = num_blocks * [y]
# skip z connections / hierarchical z
z_per_block = tf.split(z, num_blocks + 1, axis=1)
z0, z_per_block = z_per_block[0], z_per_block[1:]
y_per_block = [tf.concat([zi, y], 1) for zi in z_per_block]
act0 = ops.snlinear(z0,
gf_dim * 16 * 4 * 4,
training=training,
name='g_snh0')
act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])
# pylint: disable=line-too-long
act1 = biggan_block(act0, y_per_block[0], gf_dim * 16, num_classes,
'g_block1', training) # 8
act2 = biggan_block(act1, y_per_block[1], gf_dim * 8, num_classes,
'g_block2', training) # 16
act3 = biggan_block(act2, y_per_block[2], gf_dim * 4, num_classes,
'g_block3', training) # 32
act4 = biggan_block(act3, y_per_block[3], gf_dim * 2, num_classes,
'g_block4', training) # 64
act4 = ops.sn_non_local_block_sim(act4, training, name='g_ops') # 64
act5 = biggan_block(act4, y_per_block[4], gf_dim, num_classes,
'g_block5', training) # 128
act5 = tf.nn.relu(
tfgan.tpu.batch_norm(act5,
training,
conditional_class_labels=None,
name='g_bn'))
act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, training, 'g_snconv_last')
out = (tf.nn.tanh(act6) + 1.0) / 2.0
var_list = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, gen_scope.name)
return out, var_list
def generator(zs, target_class, gf_dim, num_classes, training=True):
"""Builds the generator segment of the graph, going from z -> G(z).
Args:
zs: Tensor representing the latent variables.
target_class: The class from which we seek to sample.
gf_dim: The gf dimension.
num_classes: Number of classes in the labels.
training: Whether in train mode or not. This affects things like batch
normalization and spectral normalization.
Returns:
- The output layer of the generator.
- A list containing all trainable varaibles defined by the model.
"""
with tf.compat.v1.variable_scope(
'generator', reuse=tf.compat.v1.AUTO_REUSE) as gen_scope:
act0 = ops.snlinear(zs,
gf_dim * 16 * 4 * 4,
training=training,
name='g_snh0')
act0 = tf.reshape(act0, [-1, 4, 4, gf_dim * 16])
# pylint: disable=line-too-long
act1 = block(act0, target_class, gf_dim * 16, num_classes, 'g_block1',
training) # 8
act2 = block(act1, target_class, gf_dim * 8, num_classes, 'g_block2',
training) # 16
act3 = block(act2, target_class, gf_dim * 4, num_classes, 'g_block3',
training) # 32
act3 = ops.sn_non_local_block_sim(act3, training, name='g_ops') # 32
act4 = block(act3, target_class, gf_dim * 2, num_classes, 'g_block4',
training) # 64
act5 = block(act4, target_class, gf_dim, num_classes, 'g_block5',
training) # 128
act5 = tf.nn.relu(
tfgan.tpu.batch_norm(act5,
training,
conditional_class_labels=None,
name='g_bn'))
act6 = ops.snconv2d(act5, 3, 3, 3, 1, 1, training, 'g_snconv_last')
out = tf.nn.tanh(act6)
var_list = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, gen_scope.name)
return out, var_list
def biggan_discriminator_128(image,
labels,
df_dim,
number_classes,
act=tf.nn.relu):
"""Builds the discriminator graph.
TODO(ilyak): debug, this implementation doesn't work as is.
...
Only position of the non local block changes
"""
with tf.compat.v1.variable_scope(
'discriminator', reuse=tf.compat.v1.AUTO_REUSE) as dis_scope:
h0 = optimized_block(image, df_dim, 'd_optimized_block1',
act=act) # 64 * 64
h0 = ops.sn_non_local_block_sim(h0, name='d_ops') # 64 * 64
h1 = block(h0, df_dim * 2, 'd_block2', act=act) # 32 * 32
h2 = block(h1, df_dim * 4, 'd_block3', act=act) # 16 * 16
h3 = block(h2, df_dim * 8, 'd_block4', act=act) # 8 * 8
h4 = block(h3, df_dim * 16, 'd_block5', act=act) # 4 * 4
h5 = block(h4, df_dim * 16, 'd_block6', downsample=False, act=act)
h5_act = act(h5)
h6 = tf.reduce_sum(input_tensor=h5_act, axis=[1, 2])
output = ops.snlinear(h6, 1, name='d_sn_linear')
classification_output = ops.snlinear(h6,
flags.FLAGS.num_classes,
name='d_sn_linear_class')
if labels is None:
pseudo_labels = tf.argmax(classification_output, axis=1)
h_labels = ops.sn_embedding(pseudo_labels,
number_classes,
df_dim * 16,
name='d_embedding')
else:
h_labels = ops.sn_embedding(labels,
number_classes,
df_dim * 16,
name='d_embedding')
output += tf.reduce_sum(input_tensor=h6 * h_labels,
axis=1,
keepdims=True)
var_list = tf.compat.v1.get_collection(
tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, dis_scope.name)
return output, classification_output, var_list
