def E2(x, args, reuse=False):
with arg_scope([hem.conv2d],
reuse=reuse,
stride=2,
padding='SAME',
filter_size=5,
init=tf.contrib.layers.xavier_initializer,
activation=tf.nn.relu): # 427 x 561 x 3
l1 = hem.conv2d(x, 3, 64, name='l1') # 214 x 281 x 64
# print('l1', l1)
l2 = hem.conv2d(l1, 64, 128, name='l2') # 107, 141 x 128
# print('l2', l2)
l3 = hem.conv2d(l2, 128, 256, name='l3') # 54 x 71 x 256
# print('l3', l3)
l4 = hem.conv2d(l3, 256, 512, name='l4') # 27 x 36 x 512
# print('l4', l4)
l5 = hem.conv2d(l4, 512, 1024, name='l5') # 14 x 18 x 1024
# print('l5', l5)
l6 = hem.conv2d(l5, 1024, 2048, name='l6') # 4 x 5 x 2048
# print('l6', l6)
f = hem.flatten(l6) # 4096
# print('f', f)
l7 = hem.dense(f, 4096, 2048, name='l7') # 2048
# print('l7', l7)
l8 = hem.dense(l7, 2048, 1, activation=tf.nn.sigmoid, name='l8') # 1
# print('l8', l8)
return l8
def latent(x, args, reuse=False):
"""Add latant nodes to the graph.
Args:
x: Tensor, output from encoder.
args: Argparse struct, containing:
latent_size, integer describing size of the latent vector
reuse: Boolean, whether to reuse variables.
"""
with arg_scope([hem.dense], reuse=reuse):
x = hem.flatten(x)
x = hem.dense(x, 32 * 4 * 4, args.latent_size, name='d1')
return x
def _gradient_penalty(x, g, args):
"""Calculate gradient penalty for discriminator.
Source: `Improved Training of Wasserstein GANs`
https://arxiv.org/abs/1704.00028
Args:
x: Tensor, the real (input) images.
g: Tensor, the fake (generator) images.
args: Argparse struct.
Returns:
Tensor, the gradient penalty term.
"""
x = hem.flatten(x)
g = hem.flatten(g)
alpha = tf.random_uniform(shape=[args.batch_size, 1], minval=0, maxval=1.0)
interpolates = x + (alpha * (g - x))
interpolates = tf.reshape(interpolates, (-1, 4, 64, 64))
d = discriminator(interpolates, args, reuse=True)
gradients = tf.gradients(d, [interpolates])[0]
slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients)))
penalty = tf.reduce_mean((slopes - 1.0)**2)
return penalty
def latent(x, batch_size, latent_size, reuse=False):
"""Adds latent nodes for sampling and reparamaterization.
Args:
x: Tensor, input images.
batch_size: Integer, batch size.
latent_size: Integer, size of latent vector.
reuse: Boolean, whether to reuse variables.
"""
with arg_scope([hem.dense],
reuse = reuse):
flat = hem.flatten(x)
z_mean = hem.dense(flat, 32*4*4, latent_size, name='d1')
z_stddev = hem.dense(flat, 32*4*4, latent_size, name='d2')
samples = tf.random_normal([batch_size, latent_size], 0, 1)
z = (z_mean + (z_stddev * samples))
return (samples, z, z_mean, z_stddev)
def gan(x, args):
"""Initialize model.
Args:
x: Tensor, the real images.
args: Argparse structure.
"""
g_opt, d_opt = hem.init_optimizer(args), hem.init_optimizer(args)
g_tower_grads, d_tower_grads = [], []
x = hem.flatten(x)
x = 2 * (x - 0.5)
# rescale [0,1] to [-1,1] depending on model
# if args.model in ['wgan', 'iwgan']:
# x = 2 * (x - 0.5)
for x, scope, gpu_id in hem.tower_scope_range(x, args.n_gpus,
args.batch_size):
# model
with tf.variable_scope('generator'):
g = generator(args.batch_size,
args.latent_size,
args,
reuse=(gpu_id > 0))
with tf.variable_scope('discriminator'):
d_real = discriminator(x, args, reuse=(gpu_id > 0))
d_fake = discriminator(g, args, reuse=True)
# losses
g_loss, d_loss = losses(x, g, d_fake, d_real, args)
# compute gradients
g_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
'generator')
d_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
'discriminator')
g_tower_grads.append(g_opt.compute_gradients(g_loss,
var_list=g_params))
d_tower_grads.append(d_opt.compute_gradients(d_loss,
var_list=d_params))
# only need one batchnorm update (ends up being updates for last tower)
batchnorm_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)
# summaries
summaries(g, x, args)
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads)
d_grads = hem.average_gradients(d_tower_grads)
hem.summarize_gradients(g_grads + d_grads)
global_step = tf.train.get_global_step()
g_apply_grads = g_opt.apply_gradients(g_grads, global_step=global_step)
d_apply_grads = d_opt.apply_gradients(d_grads, global_step=global_step)
# training
if args.model == 'gan':
train_func = _train_gan(g_apply_grads, d_apply_grads,
batchnorm_updates)
elif args.model == 'wgan':
train_func = _train_wgan(g_apply_grads, g_params, d_apply_grads,
d_params, batchnorm_updates)
elif args.model == 'iwgan':
train_func = _train_iwgan(g_apply_grads, d_apply_grads,
batchnorm_updates)
return train_func