def block(x, labels, out_channels, num_classes, name, training=True): """Builds the residual blocks used in the generator. Args: x: The 4D input tensor. labels: The labels of the class we seek to sample from. out_channels: Integer number of features in the output layer. num_classes: Integer number of classes in the labels. name: The variable scope name for the block. training: Whether this block is for training or not. Returns: A `Tensor` representing the output of the operation. """ with tf.compat.v1.variable_scope(name): labels_onehot = tf.one_hot(labels, num_classes) x_0 = x x = tf.nn.relu( tfgan.tpu.batch_norm(x, training, labels_onehot, name='cbn_0')) x = usample(x) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv1') x = tf.nn.relu( tfgan.tpu.batch_norm(x, training, labels_onehot, name='cbn_1')) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv2') x_0 = usample(x_0) x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, training, 'snconv3') return x_0 + x
def block(x, out_channels, name, downsample=True, act=tf.nn.relu): """Builds the residual blocks used in the discriminator. Args: x: The 4D input vector. out_channels: Number of features in the output layer. name: The variable scope name for the block. downsample: If True, downsample the spatial size the input tensor by a factor of 2 on each side. If False, the spatial size of the input tensor is unchanged. act: The activation function used in the block. Returns: A `Tensor` representing the output of the operation. """ with tf.variable_scope(name): input_channels = x.shape.as_list()[-1] x_0 = x x = act(x) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv1') x = act(x) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv2') if downsample: x = dsample(x) if downsample or input_channels != out_channels: x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='sn_conv3') if downsample: x_0 = dsample(x_0) return x_0 + x
def biggan_block(x, y, out_channels, num_classes, name, training=True): """Builds the residual blocks used in the generator. ... """ with tf.compat.v1.variable_scope(name): x_0 = x x = tf.nn.relu(conditional_batch_norm(x, y, training, name='cbn_0')) x = usample(x) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv1') x = tf.nn.relu(conditional_batch_norm(x, y, training, name='cbn_1')) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, training, 'snconv2') x_0 = usample(x_0) x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, training, 'snconv3') return x_0 + x
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 test_snconv2d_shapes(self): """Tests the spectrally normalized 2d conv function. This is a minimal test to make sure that shapes are OK. The image shape should match after snconv is applied. """ if tf.executing_eagerly(): # `compute_spectral_norm` doesn't work when executing eagerly. return image = tf.random.normal([10, 32, 32, 3]) snconv_image = ops.snconv2d(image, 3, k_h=3, k_w=3, d_h=1, d_w=1) self.assertEqual([10, 32, 32, 3], snconv_image.shape.as_list())
def optimized_block(x, out_channels, name, act=tf.nn.relu): """Builds optimized residual blocks for downsampling. Compared with block, optimized_block always downsamples the spatial resolution by a factor of 2 on each side. Args: x: The 4D input vector. out_channels: Number of features in the output layer. name: The variable scope name for the block. act: The activation function used in the block. Returns: A `Tensor` representing the output of the operation. """ with tf.variable_scope(name): x_0 = x x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv1') x = act(x) x = ops.snconv2d(x, out_channels, 3, 3, 1, 1, name='sn_conv2') x = dsample(x) x_0 = dsample(x_0) x_0 = ops.snconv2d(x_0, out_channels, 1, 1, 1, 1, name='sn_conv3') return x + x_0
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