def generate(batch_size,
style_noises,
noise_seed,
mix_after,
truncation_psi=0.5):
"""
given style noises, noise seed and truncation value, generate an image.
"""
# normalize noise inputs
style_noises_normalized = []
for style_noise in style_noises:
noise_std = (F.mean(style_noise**2., axis=1, keepdims=True) +
1e-8)**0.5
style_noise_normalized = F.div2(style_noise, noise_std)
style_noises_normalized.append(style_noise_normalized)
# get latent code
w = [mapping_network(_, outmaps=512) for _ in style_noises_normalized]
# truncation trick
dlatent_avg = nn.parameter.get_parameter_or_create(name="dlatent_avg",
shape=(1, 512))
w = [lerp(dlatent_avg, _, truncation_psi) for _ in w]
constant = nn.parameter.get_parameter_or_create(
name="G_synthesis/4x4/Const/const", shape=(1, 512, 4, 4))
constant_bc = F.broadcast(constant, (batch_size, ) + constant.shape[1:])
rgb_output = synthesis(w, constant_bc, noise_seed, mix_after)
return rgb_output
def smooth_crossfade(images, alpha):
s = tf.shape(images)
y = tf.reshape(images, [-1, s[1] // 2, 2, s[2] // 2, 2, s[3]])
y = tf.reduce_mean(y, axis=[2, 4], keepdims=True)
y = tf.tile(y, [1, 1, 2, 1, 2, 1])
y = tf.reshape(y, [-1, s[1], s[2], s[3]])
images = lerp(images, y, alpha)
return images
def model_fn(features, labels, mode, cfg):
del labels
resolution = features['resolution']
if mode == 'PREDICT':
random_noise = features['random_noise'] * cfg.temperature
return models.generator(random_noise,
resolution,
cfg,
is_training=False)
real_images_1 = features['real_images']
if cfg.data_format == 'NCHW':
real_images_1 = utils.nchw_to_nhwc(real_images_1)
real_images_2 = tf.image.flip_left_right(real_images_1)
real_images_1 = utils.nhwc_to_nchw(real_images_1)
real_images_2 = utils.nhwc_to_nchw(real_images_2)
else:
real_images_2 = tf.image.flip_left_right(real_images_1)
random_noise_1 = features['random_noise_1']
fake_images_out_1 = models.generator(random_noise_1,
resolution,
cfg,
is_training=True)
real_scores_out = models.discriminator(real_images_1, resolution, cfg)
fake_scores_out = models.discriminator(fake_images_out_1, resolution, cfg)
#fake_scores_out_g = models.discriminator(fake_images_out_2, resolution, cfg)
with tf.name_scope('Penalties'):
d_loss = fake_scores_out - real_scores_out
g_loss = -1.0 * fake_scores_out
with tf.name_scope('GradientPenalty'):
mixing_factors = tf.random_uniform(
[int(real_images_1.get_shape()[0]), 1, 1, 1],
0.0,
1.0,
dtype=fake_images_out_1.dtype)
mixed_images_out = ops.lerp(real_images_1, real_images_2,
mixing_factors)
mixed_scores_out = models.discriminator(mixed_images_out,
resolution, cfg)
mixed_loss = tf.reduce_sum(mixed_scores_out)
mixed_grads = tf.gradients(mixed_loss, [mixed_images_out])[0]
mixed_norms = tf.sqrt(
1e-8 + tf.reduce_sum(tf.square(mixed_grads), axis=[1, 2, 3]))
gradient_penalty = tf.square(mixed_norms - 1.0)
d_loss += gradient_penalty * 10.0
with tf.name_scope('EpsilonPenalty'):
epsilon_penalty = tf.square(real_scores_out)
d_loss += epsilon_penalty * 0.001
resolution_step = utils.get_or_create_resolution_step()
fadein_rate = tf.minimum(
tf.cast(resolution_step, tf.float32) / float(cfg.fadein_steps), 1.0)
learning_rate = cfg.base_learning_rate
d_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate,
beta1=cfg.beta1,
beta2=cfg.beta2,
epsilon=cfg.eps,
name="AdamD")
g_optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate,
beta1=cfg.beta1,
beta2=cfg.beta2,
epsilon=cfg.eps,
name="AdamG")
if cfg.data_format == 'NCHW':
fake_images_out_1 = utils.nchw_to_nhwc(fake_images_out_1)
real_images_1 = utils.nchw_to_nhwc(real_images_1)
real_images_2 = utils.nchw_to_nhwc(real_images_2)
mixed_images_out = utils.nchw_to_nhwc(mixed_images_out)
tf.summary.image('generated_images', fake_images_out_1)
tf.summary.image('real_images_1', real_images_1)
tf.summary.image('real_images_2', real_images_2)
tf.summary.image('mixed_images', mixed_images_out)
with tf.variable_scope("Loss"):
tf.summary.scalar('real_scores_out', tf.reduce_mean(real_scores_out))
tf.summary.scalar('fake_scores_out', tf.reduce_mean(fake_scores_out))
tf.summary.scalar('epsilon_penalty', tf.reduce_mean(epsilon_penalty))
tf.summary.scalar('mixed_norms', tf.reduce_mean(mixed_norms))
with tf.variable_scope("Rate"):
tf.summary.scalar('fadein', fadein_rate)
g_loss = tf.reduce_mean(g_loss)
d_loss = tf.reduce_mean(d_loss)
with tf.name_scope('TrainOps'):
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
d_step = d_optimizer.minimize(d_loss,
var_list=tf.get_collection(
tf.GraphKeys.GLOBAL_VARIABLES,
scope='Discriminator'))
g_step = g_optimizer.minimize(g_loss,
var_list=tf.get_collection(
tf.GraphKeys.GLOBAL_VARIABLES,
scope='Generator'))
with tf.control_dependencies([g_step, d_step]):
increment_global_step = tf.assign_add(
tf.train.get_or_create_global_step(), 1)
increment_resolution_step = tf.assign_add(
utils.get_or_create_resolution_step(), 1)
if resolution >= cfg.starting_resolution * 2:
with tf.control_dependencies(
[increment_global_step, increment_resolution_step]):
lerp_ops = lerp_update_ops(resolution, fadein_rate)
joint_op = tf.group([
d_step, g_step, lerp_ops[0], lerp_ops[1],
increment_global_step, increment_resolution_step
])
else:
joint_op = tf.group([
d_step, g_step, increment_global_step,
increment_resolution_step
])
return joint_op, [g_loss, d_loss], [g_optimizer, d_optimizer]