def vae(x, args):
opt = hem.init_optimizer(args)
tower_grads = []
for x, scope, gpu_id in hem.tower_scope_range(x, args.n_gpus, args.batch_size):
# model
with tf.variable_scope('encoder'):
e = encoder(x, args, reuse=(gpu_id>0))
with tf.variable_scope('latent'):
samples, z, z_mean, z_stddev = latent(e, args.batch_size, args.latent_size, reuse=(gpu_id>0))
with tf.variable_scope('decoder'):
d_real = decoder(z, args, reuse=(gpu_id>0))
d_fake = decoder(samples, args, reuse=True)
# losses
d_loss, l_loss, t_loss = losses(x, z_mean, z_stddev, d_real)
# gradients
tower_grads.append(opt.compute_gradients(d_loss))
# summaries
summaries(x, d_fake, d_real, args)
# training
avg_grads = hem.average_gradients(tower_grads)
hem.summarize_gradients(avg_grads)
train_op = opt.apply_gradients(avg_grads, global_step=tf.train.get_global_step())
return hem.default_training(train_op)
def __init__(self, x_y, args):
# init/setup
m_opt = hem.init_optimizer(args)
m_tower_grads = []
global_step = tf.train.get_global_step()
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(x_y, args.n_gpus, args.batch_size):
# for i in range(len(x_y)):
# print('estimator', i, x_y[i])
# x = x_y[0]
# y = x_y[1]
m_arch = {'E2': mean_depth_estimator.E2}
x = x_y[4]
x = tf.reshape(x, (-1, 3, 53, 70))
# print('estimator x shape', x)
y = x_y[5]
with tf.variable_scope('model'):
m_func = m_arch[args.m_arch]
m = m_func(x, args, reuse=(gpu_id>0))
self.output_layer = m
# calculate losses
m_loss = mean_depth_estimator.loss(m, x, y, args, reuse=(gpu_id > 0))
# calculate gradients
m_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'model')
m_tower_grads.append(m_opt.compute_gradients(m_loss, var_list=m_params))
# only need one batchnorm update (ends up being updates for last tower)
batchnorm_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS, scope)
# TODO: do we need to do this update for batchrenorm? for instance renorm?
# average and apply gradients
m_grads = hem.average_gradients(m_tower_grads, check_numerics=args.check_numerics)
m_apply_grads = m_opt.apply_gradients(m_grads, global_step=global_step)
# add summaries
hem.summarize_losses()
hem.summarize_gradients(m_grads, name='m_gradients')
hem.summarize_layers('m_activations', [l for l in tf.get_collection('conv_layers') if 'model' in l.name], montage=True)
mean_depth_estimator.montage_summaries(x, y, m, args)
# improved_sampler.montage_summarpies(x, y, g, x_sample, y_sample, g_sampler, x_noise, g_noise, x_shuffled, y_shuffled, g_shuffle, args)
# improved_sampler.sampler_summaries(y_sample, g_sampler, g_noise, y_shuffled, g_shuffle, args)
# training ops
with tf.control_dependencies(batchnorm_updates):
self.m_train_op = m_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
def summaries(x, d, avg_grads, args):
with tf.variable_scope('examples'):
n = int(sqrt(args.examples))
hem.montage(hem.rescale(x[0:args.examples], (-1, 1), (0, 1)),
n,
n,
name='inputs')
hem.montage(hem.rescale(d[0:args.examples], (-1, 1), (0, 1)),
n,
n,
name='outputs')
hem.summarize_activations()
hem.summarize_losses()
hem.summarize_weights_biases()
hem.summarize_gradients(avg_grads)
def summaries(x, y, x_hat, y_hat, x_grads, y_grads, args):
n = math.floor(math.sqrt(args.examples))
with arg_scope([hem.montage], height=n, width=n):
x = hem.rescale(x, (-1, 1), (0, 1))
y = hem.rescale(y, (-1, 1), (0, 1))
x_hat = hem.rescale(x_hat, (-1, 1), (0, 1))
y_hat = hem.rescale(y_hat, (-1, 1), (0, 1))
hem.montage(x[0:args.examples], name='x')
hem.montage(y[0:args.examples], name='y', colorize=True)
hem.montage(x_hat[0:args.examples], name='x_hat')
hem.montage(y_hat[0:args.examples], name='y_hat', colorize=True)
hem.summarize_gradients(x_grads)
hem.summarize_gradients(y_grads)
hem.summarize_losses()
hem.summarize_activations(False)
hem.summarize_weights_biases()
def __init__(self, x_y, args):
# init/setup
g_opt = tf.train.AdamOptimizer(args.g_lr, args.g_beta1, args.g_beta2)
d_opt = tf.train.AdamOptimizer(args.d_lr, args.d_beta1, args.d_beta2)
g_tower_grads = []
d_tower_grads = []
global_step = tf.train.get_global_step()
self.mean_image_placeholder = tf.placeholder(dtype=tf.float32,
shape=(1, 29, 29))
# self.var_image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 29, 29))
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(
x_y, args.n_gpus, args.batch_size):
with tf.variable_scope('input_preprocess'):
# split inputs and rescale
x = x_y[0]
y = x_y[1]
# re-attach shape info
x = tf.reshape(x, (args.batch_size, 3, 65, 65))
# rescale from [0,1] to actual world depth
y = y * 10.0
y = hem.crop_to_bounding_box(y, 17, 17, 29, 29)
# re-attach shape info
y = tf.reshape(y, (args.batch_size, 1, 29, 29))
y_bar = tf.reduce_mean(y, axis=[2, 3], keep_dims=True)
x_sample = tf.stack([x[0]] * args.batch_size)
y_sample = tf.stack([y[0]] * args.batch_size)
# create model
with tf.variable_scope('generator'):
if args.model_version == 'baseline':
g = self.g_baseline(x, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
g_sampler = self.g_baseline(x_sample, args, reuse=True)
y_sample_bar = tf.reduce_mean(y_sample,
axis=[2, 3],
keep_dims=True)
y_sampler = g_sampler + y_sample_bar
with tf.variable_scope('discriminator'):
if args.model_version == 'baseline':
# this is the 'mean_adjusted' model from paper_baseline_sampler.py
d_fake, d_fake_logits = self.d_baseline(x,
y_hat - y_bar,
args,
reuse=(gpu_id > 0))
d_real, d_real_logits = self.d_baseline(x,
y - y_bar,
args,
reuse=True)
# calculate losses
g_loss, d_loss = self.loss(d_real,
d_real_logits,
d_fake,
d_fake_logits,
reuse=(gpu_id > 0))
# calculate 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))
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads,
check_numerics=args.check_numerics)
d_grads = hem.average_gradients(d_tower_grads,
check_numerics=args.check_numerics)
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)
# add summaries
hem.summarize_losses()
hem.summarize_gradients(g_grads, name='g_gradients')
hem.summarize_gradients(d_grads, name='d_gradients')
generator_layers = [
l for l in tf.get_collection('conv_layers')
if 'generator' in l.name
]
discriminator_layers = [
l for l in tf.get_collection('conv_layers')
if 'discriminator' in l.name
]
hem.summarize_layers('g_activations', generator_layers, montage=True)
hem.summarize_layers('d_activations',
discriminator_layers,
montage=True)
self.montage_summaries(x, y, g, y_hat, args, name='y_hat')
self.metric_summaries(x, y, g, y_hat, args, name='y_hat')
self.metric_summaries(x, y, g_0, y_0, args, name='y_0')
self.metric_summaries(x,
y,
g,
self.mean_image_placeholder * 10.0,
args,
name='y_mean')
self.metric_summaries(x_sample,
y_sample,
g_sampler,
y_sampler,
args,
name='y_sampler')
self.montage_summaries(x_sample,
y_sample,
g_sampler,
y_sampler,
args,
name='y_sampler')
# training ops
self.g_train_op = g_apply_grads
self.d_train_op = d_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
def __init__(self, x_y, args):
# init/setup
g_opt = tf.train.AdamOptimizer(args.g_lr, args.g_beta1, args.g_beta2)
g_tower_grads = []
global_step = tf.train.get_global_step()
self.mean_image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 29, 29))
# self.var_image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 29, 29))
self.x = []
self.y = []
self.y_hat = []
self.g = []
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(x_y, args.n_gpus, args.batch_size):
with tf.variable_scope('input_preprocess'):
# split inputs and rescale
x = tf.identity(x_y[0], name='tower_{}_x'.format(gpu_id))
y = tf.identity(x_y[1], name='tower_{}_y'.format(gpu_id))
# re-attach shape info
x = tf.reshape(x, (args.batch_size, 3, 65, 65))
# rescale from [0,1] to actual world depth
y = y * 10.0
y = hem.crop_to_bounding_box(y, 17, 17, 29, 29)
# re-attach shape info
y = tf.reshape(y, (args.batch_size, 1, 29, 29))
y_bar = tf.reduce_mean(y, axis=[2, 3], keep_dims=True)
y_bar = tf.identity(y_bar, name='tower_{}_y_bar'.format(gpu_id))
# create model
with tf.variable_scope('generator'):
if args.model_version == 'baseline':
g = self.g_baseline(x, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g
y_0 = g_0
elif args.model_version == 'mean_adjusted':
g = self.g_baseline(x, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
elif args.model_version == 'mean_provided':
g = self.g_mean_provided(x, y_bar, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
elif args.model_version == 'mean_provided2':
g = self.g_mean_provided2(x, y_bar, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
g = tf.identity(g, 'tower_{}_g'.format(gpu_id))
g_0 = tf.identity(g_0, 'tower_{}_g0'.format(gpu_id))
y_hat = tf.identity(y_hat, 'tower_{}_y_hat'.format(gpu_id))
y_0 = tf.identity(y_0, 'tower_{}_y0'.format(gpu_id))
# if gpu_id == 0:
# tf.summary.histogram('g', g)
# tf.summary.histogram('y_hat', y_hat)
# tf.summary.histogram('y_0', y_0)
# hem.montage(g, num_examples=64, width=8, height=8, name='g')
# hem.montage(y_hat, num_examples=64, width=8, height=8, name='y_hat')
# hem.montage(y_0, num_examples=64, width=8, height=8, name='y_0')
self.g.append(g)
self.y_hat.append(y_hat)
self.y.append(y)
self.x.append(x)
# calculate losses
g_loss = self.loss(x, y, y_hat, args, reuse=(gpu_id > 0))
# calculate gradients
g_params = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'generator')
g_tower_grads.append(g_opt.compute_gradients(g_loss, var_list=g_params))
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads, check_numerics=args.check_numerics)
g_apply_grads = g_opt.apply_gradients(g_grads, global_step=global_step)
# add summaries
hem.summarize_losses()
hem.summarize_gradients(g_grads, name='g_gradients')
generator_layers = [l for l in tf.get_collection('conv_layers') if 'generator' in l.name]
hem.summarize_layers('g_activations', generator_layers, montage=True)
self.montage_summaries(x, y, g, y_hat, args)
self.metric_summaries(x, y, g, y_hat, args, name='y_hat')
self.metric_summaries(x, y, g_0, y_0, args, name='y_0')
self.metric_summaries(x, y, g, self.mean_image_placeholder * 10.0, args, name='y_mean')
# training ops
self.g_train_op = g_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
def __init__(self, x_y, estimator, args):
# init/setup
g_opt = hem.init_optimizer(args)
d_opt = hem.init_optimizer(args)
g_tower_grads = []
d_tower_grads = []
global_step = tf.train.get_global_step()
# sess = tf.Session(config = tf.ConfigProto(allow_soft_placement=True, log_device_placement=False))
# new_saver = tf.train.import_meta_graph(
# '/mnt/research/projects/autoencoders/workspace/improved_sampler/experimentE/meandepth.e1/checkpoint-4.meta', import_scope='estimator')
# new_saver.restore(sess, tf.train.latest_checkpoint('/mnt/research/projects/autoencoders/workspace/improved_sampler/experimentE/meandepth.e1'))
#
# # estimator_vars = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
# # print('estimator_vars:', estimator_vars)
#
# # print('all ops!')
# # for op in tf.get_default_graph().get_operations():
# # if 'l8' in str(op.name):
# # print(str(op.name))
#
#
# # sess.graph
# # graph = tf.get_default_graph()
# estimator_tower0 = sess.graph.as_graph_element('estimator/tower_0/model/l8/add').outputs[0]
# estimator_tower1 = sess.graph.as_graph_element('estimator/tower_1/model/l8/add').outputs[0]
# self.estimator_placeholder = sess.graph.as_graph_element('estimator/input_pipeline/Placeholder') #.outputs[0]
# print('PLACEHOLDER:', self.estimator_placeholder)
# print('estimator_tower0:', estimator_tower0)
# print('estimator_tower1:', estimator_tower1)
# estimator_tower0 = tf.stop_gradient(estimator_tower0)
# estimator_tower1 = tf.stop_gradient(estimator_tower1)
# sess.close()
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(x_y, args.n_gpus, args.batch_size):
with tf.variable_scope('input_preprocess'):
# split inputs and rescale
x = hem.rescale(x_y[0], (0, 1), (-1, 1))
y = hem.rescale(x_y[1], (0, 1), (-1, 1))
# if args.g_arch == 'E2':
y = hem.crop_to_bounding_box(y, 16, 16, 32, 32)
y = tf.reshape(y, (-1, 1, 32, 32))
x_loc = x_y[2]
y_loc = x_y[3]
scene_image = x_y[4]
mean_depth = tf.stop_gradient(estimator.output_layer)
# print('mean_depth_layer:', estimator.output_layer)
# mean_depth = estimator(scene_image)
# mean_depth = tf.expand_dims(mean_depth, axis=-1)
# mean_depth = tf.expand_dims(mean_depth, axis=-1)
mean_depth_channel = tf.stack([mean_depth] * 64, axis=2)
mean_depth_channel = tf.stack([mean_depth_channel] * 64, axis=3)
# mean_depth_channel = tf.squeeze(mean_depth_channel)
# print('mean_depth_layer99:', mean_depth_channel)
# mean_depth_channel = tf.ones_like(x_loc) * mean_depth
# print('x', x)
# print('x_loc', x_loc)
# print('y_loc', y_loc)
# print('mean_depth_channel', mean_depth_channel)
x = tf.concat([x, x_loc, y_loc, mean_depth_channel], axis=1)
# print('x shape:', x)
# create repeated image tensors for sampling
x_sample = tf.stack([x[0]] * args.batch_size)
y_sample = tf.stack([y[0]] * args.batch_size)
# shuffled x for variance calculation
x_shuffled = tf.random_shuffle(x)
y_shuffled = y
# noise vector for testing
x_noise = tf.random_uniform(tf.stack([tf.Dimension(args.batch_size), x.shape[1], x.shape[2], x.shape[3]]), minval=-1.0, maxval=1.0)
g_arch = {'E2': experimental_sampler.generatorE2}
d_arch = {'E2': experimental_sampler.discriminatorE2}
# create model
with tf.variable_scope('generator'):
g_func = g_arch[args.g_arch]
g = g_func(x, args, reuse=(gpu_id > 0))
g_sampler = g_func(x_sample, args, reuse=True)
g_shuffle = g_func(x_shuffled, args, reuse=True)
g_noise = g_func(x_noise, args, reuse=True)
with tf.variable_scope('discriminator'):
d_func = d_arch[args.d_arch]
d_real, d_real_logits = d_func(x, y, args, reuse=(gpu_id > 0))
d_fake, d_fake_logits = d_func(x, g, args, reuse=True)
# calculate losses
g_loss, d_loss = experimental_sampler.loss(d_real, d_real_logits, d_fake, d_fake_logits, x, g, y, None, args, reuse=(gpu_id > 0))
# calculate 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)
# TODO: do we need to do this update for batchrenorm? for instance renorm?
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads, check_numerics=args.check_numerics)
d_grads = hem.average_gradients(d_tower_grads, check_numerics=args.check_numerics)
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)
# add summaries
hem.summarize_losses()
hem.summarize_gradients(g_grads, name='g_gradients')
hem.summarize_gradients(d_grads, name='d_gradients')
hem.summarize_layers('g_activations', [l for l in tf.get_collection('conv_layers') if 'generator' in l.name], montage=True)
hem.summarize_layers('d_activations', [l for l in tf.get_collection('conv_layers') if 'discriminator' in l.name], montage=True)
experimental_sampler.montage_summaries(x, y, g, x_sample, y_sample, g_sampler, x_noise, g_noise, x_shuffled, y_shuffled, g_shuffle, args)
experimental_sampler.sampler_summaries(y_sample, g_sampler, g_noise, y_shuffled, g_shuffle, args)
# training ops
with tf.control_dependencies(batchnorm_updates):
self.g_train_op = g_apply_grads
self.d_train_op = d_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
def __init__(self, x_y, args):
# init/setup
# wgan training
if args.training_version == 'wgan':
g_opt = tf.train.RMSPropOptimizer(args.g_lr)
d_opt = tf.train.AdamOptimizer(args.d_lr)
else:
g_opt = tf.train.AdamOptimizer(args.g_lr, args.g_beta1, args.g_beta2)
d_opt = tf.train.AdamOptimizer(args.d_lr, args.d_beta1, args.d_beta2)
g_tower_grads = []
d_tower_grads = []
global_step = tf.train.get_global_step()
self.x = []
self.y = []
self.y_hat = []
self.g = []
self.mean_image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 29, 29))
# self.var_image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 29, 29))
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(x_y, args.n_gpus, args.batch_size):
with tf.variable_scope('input_preprocess'):
# split inputs and rescale
x = tf.identity(x_y[0], name='tower_{}_x'.format(gpu_id))
y = tf.identity(x_y[1], name='tower_{}_y'.format(gpu_id))
# re-attach shape info
x = tf.reshape(x, (args.batch_size, 3, 65, 65))
# rescale from [0,1] to actual world depth
y = y * 10.0
y = hem.crop_to_bounding_box(y, 17, 17, 29, 29)
# re-attach shape info
y = tf.reshape(y, (args.batch_size, 1, 29, 29))
y_bar = tf.reduce_mean(y, axis=[2, 3], keep_dims=True)
y_bar = tf.identity(y_bar, name='tower_{}_y_bar'.format(gpu_id))
# create model
with tf.variable_scope('generator'):
if args.model_version == 'baseline':
g = self.g_baseline(x, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g
y_0 = g_0
elif args.model_version == 'mean_adjusted':
g = self.g_baseline(x, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
elif args.model_version == 'mean_provided':
g = self.g_mean_provided(x, y_bar, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
elif args.model_version == 'mean_provided2':
g = self.g_mean_provided2(x, y_bar, args, reuse=(gpu_id > 0))
g_0 = tf.zeros_like(g)
y_hat = g + y_bar
y_0 = g_0 + y_bar
g = tf.identity(g, 'tower_{}_g'.format(gpu_id))
g_0 = tf.identity(g_0, 'tower_{}_g0'.format(gpu_id))
y_hat = tf.identity(y_hat, 'tower_{}_y_hat'.format(gpu_id))
y_0 = tf.identity(y_0, 'tower_{}_y0'.format(gpu_id))
with tf.variable_scope('discriminator'):
if args.model_version == 'baseline':
d_fake, d_fake_logits = self.d_baseline(x, y_hat, args, reuse=(gpu_id > 0))
d_real, d_real_logits = self.d_baseline(x, y, args, reuse=True)
elif args.model_version == 'mean_adjusted':
d_fake, d_fake_logits = self.d_baseline(x, y_hat - y_bar, args, reuse=(gpu_id > 0))
d_real, d_real_logits = self.d_baseline(x, y - y_bar, args, reuse=True)
elif args.model_version == 'mean_provided':
d_fake, d_fake_logits = self.d_mean_provided(x, y_hat - y_bar, y_bar, args, reuse=(gpu_id > 0))
d_real, d_real_logits = self.d_mean_provided(x, y - y_bar, y_bar, args, reuse=True)
elif args.model_version == 'mean_provided2':
d_fake, d_fake_logits = self.d_mean_provided2(x, y_hat - y_bar, y_bar, args, reuse=(gpu_id > 0))
d_real, d_real_logits = self.d_mean_provided2(x, y - y_bar, y_bar, args, reuse=True)
d_fake = tf.identity(d_fake, 'tower_{}_d_fake'.format(gpu_id))
d_real = tf.identity(d_real, 'tower_{}_d_real'.format(gpu_id))
d_fake_logits = tf.identity(d_fake_logits, 'tower_{}_d_fake_logits'.format(gpu_id))
d_real_logits = tf.identity(d_real_logits, 'tower_{}_d_real_logits'.format(gpu_id))
self.g.append(g)
self.y_hat.append(y_hat)
self.y.append(y)
self.x.append(x)
# calculate losses
g_loss, d_loss = self.loss(d_real, d_real_logits, d_fake, d_fake_logits, args, reuse=(gpu_id > 0))
# calculate 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))
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads, check_numerics=args.check_numerics)
d_grads = hem.average_gradients(d_tower_grads, check_numerics=args.check_numerics)
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)
# add summaries
hem.summarize_losses()
hem.summarize_gradients(g_grads, name='g_gradients')
hem.summarize_gradients(d_grads, name='d_gradients')
generator_layers = [l for l in tf.get_collection('conv_layers') if 'generator' in l.name]
discriminator_layers = [l for l in tf.get_collection('conv_layers') if 'discriminator' in l.name]
hem.summarize_layers('g_activations', generator_layers, montage=True)
hem.summarize_layers('d_activations', discriminator_layers, montage=True)
self.montage_summaries(x, y, g, y_hat, args)
self.metric_summaries(x, y, g, y_hat, args, name='y_hat')
self.metric_summaries(x, y, g_0, y_0, args, name='y_0')
self.metric_summaries(x, y, g, self.mean_image_placeholder * 10.0, args, name='y_mean')
# training ops
if args.training_version == 'wgan':
clip_D = [p.assign(tf.clip_by_value(p, -0.01, 0.01)) for p in d_params]
clip_G = [p.assign(tf.clip_by_value(p, -0.01, 0.01)) for p in g_params]
with tf.control_dependencies(clip_D):
self.d_train_op = d_apply_grads
with tf.control_dependencies(clip_G):
self.g_train_op = g_apply_grads
else:
self.g_train_op = g_apply_grads
self.d_train_op = d_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
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