def g_mean_provided2(self, x, y_bar, args, reuse=False):
with tf.variable_scope('encoder', reuse=reuse), \
arg_scope([hem.conv2d],
reuse=reuse,
filter_size=5,
stride=2,
padding='VALID',
init=tf.contrib.layers.xavier_initializer,
activation=tf.nn.relu): # 65x65x3
x = tf.concat([x, tf.ones((args.batch_size, 1, 65, 65))], axis=1)
e1 = hem.conv2d(x, 4, 64, name='e1') # 31x31x64
# e1 = tf.concat([e1, tf.ones((args.batch_size, 1, 31, 31)) * y_bar], axis=1)
e2 = hem.conv2d(e1, 64, 128, name='e2') # 14x14x128
e3 = hem.conv2d(e2, 128, 256, name='e3') # 5x5x256
e4 = hem.conv2d(e3, 256, 512, name='e4') # 1x1x512
with tf.variable_scope('decoder', reuse=reuse), \
arg_scope([hem.deconv2d, hem.conv2d],
reuse=reuse,
filter_size=5,
stride=2,
init=tf.contrib.layers.xavier_initializer,
padding='VALID',
activation=lambda x: hem.lrelu(x, leak=0.2)): # 1x1x512
y_hat = hem.deconv2d(e4, 512, 256, output_shape=(args.batch_size, 256, 5, 5), name='d1') # 5x5x256
y_hat = tf.concat([y_hat, e3], axis=1) # 5x5x512
y_hat = hem.deconv2d(y_hat, 512, 128, output_shape=(args.batch_size, 128, 14, 14), name='d2') # 14x14x128
y_hat = tf.concat([y_hat, e2], axis=1) # 14x14x256
y_hat = hem.deconv2d(y_hat, 256, 64, output_shape=(args.batch_size, 64, 31, 31), name='d3') # 31x31x64
y_hat = tf.concat([y_hat, e1], axis=1) # 31x31x128
y_hat = hem.conv2d(y_hat, 128, 1, stride=1, filter_size=1, padding='SAME', activation=None,
name='d4') # 31x31x1
y_hat = hem.crop_to_bounding_box(y_hat, 0, 0, 29, 29) # 29x29x1
# y_hat = tf.maximum(y_hat, tf.zeros_like(y_hat))
return y_hat
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 g_baseline(self, x, args, reuse=False):
with tf.variable_scope('encoder', reuse=reuse), \
arg_scope([hem.conv2d],
reuse=reuse,
filter_size=5,
stride=2,
padding='VALID',
use_batch_norm=args.e_bn,
init=tf.contrib.layers.xavier_initializer,
activation=tf.nn.relu): # 65x65x3
if args.noise_layer == 'x':
noise = tf.random_uniform([args.batch_size, 1, 65, 65],
minval=0,
maxval=1)
e1 = hem.conv2d(tf.concat([x, noise], axis=1),
4,
64,
name='e1') # 31x31x64
else:
e1 = hem.conv2d(x, 3, 64, name='e1')
if args.noise_layer == 'e1':
noise = tf.random_uniform([args.batch_size, 1, 31, 31],
minval=0,
maxval=1)
e2 = hem.conv2d(tf.concat([e1, noise], axis=1),
65,
128,
name='e2') # 14x14x128
else:
e2 = hem.conv2d(e1, 64, 128, name='e2') # 14x14x128
if args.noise_layer == 'e2':
noise = tf.random_uniform([args.batch_size, 1, 14, 14],
minval=0,
maxval=1)
e3 = hem.conv2d(tf.concat([e2, noise], axis=1),
129,
256,
name='e3') # 5x5x256
else:
e3 = hem.conv2d(e2, 128, 256, name='e3') # 5x5x256
if args.noise_layer == 'e3':
noise = tf.random_uniform([args.batch_size, 1, 5, 5],
minval=0,
maxval=1)
e4 = hem.conv2d(tf.concat([e3, noise], axis=1),
257,
512,
name='e4') # 1x1x512
else:
e4 = hem.conv2d(e3, 256, 512, name='e4') # 1x1x512
with tf.variable_scope('decoder', reuse=reuse), \
arg_scope([hem.deconv2d, hem.conv2d],
reuse=reuse,
filter_size=5,
stride=2,
init=tf.contrib.layers.xavier_initializer,
padding='VALID',
activation=lambda x: hem.lrelu(x, leak=0.2)): # 1x1x512
# TODO: noise could be of size 512, instead of 1
if args.noise_layer == 'e4':
noise = tf.random_uniform([args.batch_size, 1, 1, 1],
minval=0,
maxval=1)
y_hat = hem.deconv2d(tf.concat([e4, noise], axis=1),
513,
256,
output_shape=(args.batch_size, 256, 5, 5),
name='d1') # 5x5x256
elif args.noise_layer == 'e4-512':
noise = tf.random_uniform([args.batch_size, 512, 1, 1],
minval=0,
maxval=1)
y_hat = hem.deconv2d(tf.concat([e4, noise], axis=1),
1024,
256,
output_shape=(args.batch_size, 256, 5, 5),
name='d1') # 5x5x256
else:
y_hat = hem.deconv2d(e4,
512,
256,
output_shape=(args.batch_size, 256, 5, 5),
name='d1') # 5x5x256
y_hat = tf.concat([y_hat, e3], axis=1) # 5x5x512
if args.noise_layer == 'd2':
noise = tf.random_uniform([args.batch_size, 1, 5, 5],
minval=0,
maxval=1)
y_hat = hem.deconv2d(tf.concat([y_hat, noise], axis=1),
513,
128,
output_shape=(args.batch_size, 128, 14,
14),
name='d2') # 14x14x128z
else:
y_hat = hem.deconv2d(y_hat,
512,
128,
output_shape=(args.batch_size, 128, 14,
14),
name='d2') # 14x14x128z
y_hat = tf.concat([y_hat, e2], axis=1) # 14x14x256
if args.noise_layer == 'd3':
noise = tf.random_uniform([args.batch_size, 1, 14, 14],
minval=0,
maxval=1)
y_hat = hem.deconv2d(tf.concat([y_hat, noise], axis=1),
257,
64,
output_shape=(args.batch_size, 64, 31,
31),
name='d3') # 31x31x64
else:
y_hat = hem.deconv2d(y_hat,
256,
64,
output_shape=(args.batch_size, 64, 31,
31),
name='d3') # 31x31x64
y_hat = tf.concat([y_hat, e1], axis=1) # 31x31x128
if args.noise_layer == 'd4':
noise = tf.random_uniform([args.batch_size, 1, 31, 31],
minval=0,
maxval=1)
y_hat = hem.conv2d(tf.concat([y_hat, noise], axis=1),
129,
1,
stride=1,
filter_size=1,
padding='SAME',
activation=None,
name='d4') # 31x31x1
else:
y_hat = hem.conv2d(y_hat,
128,
1,
stride=1,
filter_size=1,
padding='SAME',
activation=None,
name='d4') # 31x31x1
y_hat = hem.crop_to_bounding_box(y_hat, 0, 0, 29, 29) # 29x29x1
#y_hat = tf.maximum(y_hat, tf.zeros_like(y_hat))
return y_hat
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'))
