def __init__(self, x_y, args):
g_opt = hem.init_optimizer(args)
d_opt = hem.init_optimizer(args)
q_opt = hem.init_optimizer(args)
x = hem.rescale(x_y[0], (0, 1), (-1, 1)) # 256x256x3
y = hem.rescale(x_y[1], (0, 1), (-1, 1)) # 256x256x1
z = tf.random_uniform((args.batch_size, 1, 256, 256)) # 256x256x1
with tf.variable_scope('generator') as scope:
g = info_gan.generator(z, x)
with tf.variable_scope('discriminator') as scope:
d_real = info_gan.discriminator(y)
d_fake = info_gan.discriminator(g, reuse=True)
with tf.variable_scope('predictor') as scope:
q = info_gan.predictor(g)
g_loss = -tf.reduce_mean(tf.log(d_fake + 1e-8))
d_loss = -tf.reduce_mean(tf.log(d_real + 1e-8) + tf.log(1 - d_fake + 1e-8))
cross_entropy = tf.reduce_mean(-tf.reduce_sum(tf.log(q + 1e-8) * x), axis=1)
entropy = tf.reduce_mean(-tf.reduce_sum(tf.log(x + 1e-8) * x), axis=1)
q_loss = cross_entropy + entropy
for l in [g_loss, d_loss, q_loss]:
tf.add_to_collection('losses', l)
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
g_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'generator')
d_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'discriminator')
q_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'predictor')
self.g_train_op = g_opt.minimize(g_loss, var_list=g_vars)
self.d_train_op = d_opt.minimize(d_loss, var_list=d_vars)
self.q_train_op = q_opt.minimize(q_loss, var_list=q_vars + g_vars)
def __init__(self, x_y, args):
x_opt = hem.init_optimizer(args)
y_opt = hem.init_optimizer(args)
x_decoder_tower_grads = []
y_decoder_tower_grads = []
global_step = tf.train.get_global_step()
for x_y, scope, gpu_id in hem.tower_scope_range(
x_y, args.n_gpus, args.batch_size):
x = hem.rescale(x_y[0], (0, 1), (-1, 1))
y = hem.rescale(x_y[1], (0, 1), (-1, 1))
with tf.variable_scope('encoder'):
e = artist.encoder(x, reuse=gpu_id > 0)
with tf.variable_scope('x_decoder'):
x_hat = artist.decoder(e,
args,
channel_output=3,
reuse=gpu_id > 0)
with tf.variable_scope('y_decoder'):
y_hat = artist.decoder(e,
args,
channel_output=1,
reuse=gpu_id > 0)
x_hat_loss, y_hat_loss = artist.losses(x, x_hat, y, y_hat,
gpu_id == 0)
encoder_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
'encoder')
x_decoder_vars = tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES, 'x_decoder')
y_decoder_vars = tf.get_collection(
tf.GraphKeys.TRAINABLE_VARIABLES, 'y_decoder')
# # train for x-reconstruction
# x_decoder_tower_grads.append(x_opt.compute_gradients(x_hat_loss, var_list=encoder_vars + x_decoder_vars))
# y_decoder_tower_grads.append(y_opt.compute_gradients(y_hat_loss, var_list=y_decoder_vars))
# train for y-reconstruction
x_decoder_tower_grads.append(
x_opt.compute_gradients(x_hat_loss, var_list=x_decoder_vars))
y_decoder_tower_grads.append(
y_opt.compute_gradients(y_hat_loss,
var_list=encoder_vars +
y_decoder_vars))
batchnorm_updates = tf.get_collection(tf.GraphKeys.UPDATE_OPS,
scope)
x_grads = hem.average_gradients(x_decoder_tower_grads)
y_grads = hem.average_gradients(y_decoder_tower_grads)
with tf.control_dependencies(batchnorm_updates):
self.x_train_op = x_opt.apply_gradients(x_grads,
global_step=global_step)
self.y_train_op = y_opt.apply_gradients(y_grads,
global_step=global_step)
self.x_hat = x_hat
self.y_hat = y_hat
self.x_hat_loss = x_hat_loss
self.y_hat_loss = y_hat_loss
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
artist.summaries(x, y, x_hat, y_hat, x_grads, y_grads, args)
def montage_summaries(x_rgb, x_depth, g, args, x_repeated_rgb,
x_repeated_depth, g_sampler):
"""Adds montage summaries to the graph for the images, generator, and sampler.
Args:
x_rgb: Tensor, the real RGB image input batch.
x_depth: Tensor, the real depth input batch.
g: Tensor, the generator's output.
args: Argparse struct.
x_repeated_rgb: Tensor, the real RGB image inputs to the sampler path.
x_repeated_depth: Tensor, the real RGB depth inputs to the sampler path.
g_sampler: Tensor, the sampler generator outputs.
Returns:
None.
"""
# print('x shape', x_rgb.shape)
# print('y shape', x_depth.shape)
# print('y_hat shape', g.shape)
# print('x_repeated shape', x_repeated_depth.shape)
# print('y_hat_repeated shape', g_sampler.shape)
# example image montages
n = math.floor(math.sqrt(args.examples))
# rescale
with tf.variable_scope('montage_preprocess'):
g = tf.reshape(g, tf.shape(x_depth)) # re-attach lost shape info
g = hem.rescale(g, (-1, 1), (0, 1))
x_depth = hem.rescale(x_depth, (-1, 1), (0, 1))
g_sampler = tf.reshape(g_sampler, tf.shape(x_depth))
g_sampler = hem.rescale(g_sampler, (-1, 1), (0, 1))
x_repeated_depth = hem.rescale(x_repeated_depth, (-1, 1), (0, 1))
# x_depth = tf.reshape(x_depth, (-1, 1, 65, 65))
# add image montages
with arg_scope([hem.montage], height=n, width=n):
# example batches
hem.montage(x_rgb[0:args.examples], name='real/images')
hem.montage(x_depth[0:args.examples],
name='real/depths',
colorize=True)
hem.montage(x_rgb[0:args.examples], name='fake/images')
hem.montage(g[0:args.examples], name='fake/depths', colorize=True)
# sampler
hem.montage(x_repeated_rgb[0:args.examples], name='sampler/images')
hem.montage(x_repeated_depth[0:args.examples],
name='sampler/depths',
colorize=True)
hem.montage(g_sampler[0:args.examples],
name='sampler/fake',
colorize=True)
def losses(x, x_hat, y, y_hat, add_to_collection=False):
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))
x_hat_loss = tf.reduce_mean(tf.square(x - x_hat), name='x_hat_loss')
y_hat_loss = tf.reduce_mean(tf.square(y - y_hat), name='y_hat_loss')
y_hat_rmse_loss = tf.sqrt(tf.reduce_mean(tf.square(y - y_hat)),
name='y_hat_rmse')
if add_to_collection:
tf.add_to_collection('losses', x_hat_loss)
tf.add_to_collection('losses', y_hat_loss)
tf.add_to_collection('losses', y_hat_rmse_loss)
return x_hat_loss, y_hat_loss
def loss(d_real,
d_real_logits,
d_fake,
d_fake_logits,
g,
x_depth,
args,
reuse=False):
"""Adds loss nodes to the graph.
Args:
d_real: Tensor, the discriminator's output with a real batch.
d_fake: Tensor, the discriminator's output with a fake batch.
g: Tensor, the generator's output.
x_depth: Tensor, the real depth maps.
reuse: Bool, whether to add the loss nodes to the loss collection
for later summary collection. Should only be True once.
Returns:
Tensors, the losses for the generator and discriminator, respectively.
"""
def xentropy(logits, labels):
return tf.nn.sigmoid_cross_entropy_with_logits(logits=logits,
labels=labels)
with tf.variable_scope('loss'):
g = hem.rescale(g, (-1, 1), (0, 1))
x_depth = hem.rescale(x_depth, (-1, 1), (0, 1))
# losses
with tf.variable_scope('generator'):
g_fake = tf.reduce_mean(xentropy(d_fake_logits,
tf.ones_like(d_fake)),
name='g_fake')
with tf.variable_scope('discriminator'):
d_real = tf.reduce_mean(xentropy(d_real_logits,
tf.ones_like(d_fake)),
name='d_real')
d_fake = tf.reduce_mean(xentropy(d_fake_logits,
tf.zeros_like(d_fake)),
name='d_fake')
d_total = tf.identity(d_real + d_fake, name='total')
rmse_loss = hem.rmse(x_depth, g)
l1_loss = tf.reduce_mean(tf.abs(x_depth - g), name='l1')
# only add these to the collection once
if not reuse:
hem.add_to_collection(
'losses',
[g_fake, d_real, d_fake, d_total, rmse_loss, l1_loss])
return g_fake, d_total
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 cnn(x, args):
"""Initialize a standard convolutional autoencoder.
Args:
x: Tensor, input tensor representing the images.
args: Argparse struct.
"""
opt = hem.init_optimizer(args)
tower_grads = []
x = hem.rescale(x, (0, 1), (-1, 1))
for x, scope, gpu_id in hem.tower_scope_range(x, args.n_gpus,
args.batch_size):
# create model
with tf.variable_scope('encoder'):
e = encoder(x, args, reuse=(gpu_id > 0))
with tf.variable_scope('latent'):
z = latent(e, args, reuse=(gpu_id > 0))
with tf.variable_scope('decoder'):
d = decoder(z, args, reuse=(gpu_id > 0))
with tf.variable_scope('loss'):
d_loss = loss(x, d, reuse=(gpu_id > 0))
# compute gradients
tower_grads.append(opt.compute_gradients(d_loss))
# training
avg_grads = hem.average_gradients(tower_grads)
train_op = opt.apply_gradients(avg_grads,
global_step=tf.train.get_global_step())
train_func = hem.default_training(train_op)
# add summaries
summaries(x, d, avg_grads, args)
return train_func
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 sampler_summaries(x_repeated_depth, g_sampler, args):
with tf.variable_scope('sampler_metrics'):
# mean and var metrics for sampler
g_sampler = tf.reshape(
g_sampler, x_repeated_depth.shape) # reattach shape info
sampler_gan.summarize_moments(g_sampler, 'depth', args)
sampler_gan.summarize_moments(
g_sampler - tf.reduce_mean(g_sampler), 'depth_normalized',
args)
# ground truth example
ic = tf.expand_dims(x_repeated_depth[0], axis=0)
ic = hem.rescale(ic, (-1, 1), (0, 1))
ic = hem.colorize(ic)
tf.summary.image('real_depth', ic)
# mean and min l2 loss from sampler
sample_l2_loss = tf.reduce_mean(tf.square(x_repeated_depth -
g_sampler),
axis=[1, 2, 3])
mean_l2_loss = tf.reduce_mean(sample_l2_loss)
min_l2_loss = tf.reduce_min(sample_l2_loss)
tf.summary.scalar('mean sample l2', mean_l2_loss)
tf.summary.scalar('min sample l2', min_l2_loss)
sample_rmse_loss = tf.reduce_mean(tf.sqrt(
tf.square(x_repeated_depth - g_sampler)),
axis=[1, 2, 3])
mean_rmse_loss = tf.reduce_mean(sample_rmse_loss)
min_rmse_loss = tf.reduce_min(sample_rmse_loss)
tf.summary.scalar('mean sample rmse', mean_rmse_loss)
tf.summary.scalar('min sample rmse', min_rmse_loss)
def montage_summaries(x, y, g,
x_sample, y_sample, g_sampler,
x_noise, g_noise,
x_shuffled, y_shuffled, g_shuffle,
args):
n = math.floor(math.sqrt(args.examples))
if args.g_arch in ['E2']:
x, x_loc, y_loc, mean_vec = tf.split(x, num_or_size_splits=[3, 1, 1, 1], axis=1)
x_noise, x_loc_noise, y_loc_noise, mean_vec_noise = tf.split(x_noise, num_or_size_splits=[3, 1, 1, 1], axis=1)
x_shuffled, x_loc_shuffled, y_loc_shuffled, mean_vec_shuffled = tf.split(x_shuffled, num_or_size_splits=[3, 1, 1, 1], axis=1)
x_sample, x_loc_sample, y_loc_sample, mean_vec_sample = tf.split(x_sample, num_or_size_splits=[3, 1, 1, 1], axis=1)
with tf.variable_scope('model'):
hem.montage(x_loc, num_examples=args.examples, height=n, width=n, name='x_loc', colorize=True)
hem.montage(y_loc, num_examples=args.examples, height=n, width=n, name='y_loc', colorize=True)
hem.montage(mean_vec, num_examples=args.examples, height=n, width=n, name='mean_vec', colorize=True)
with tf.variable_scope('montage_preprocess'):
g = tf.reshape(g, tf.shape(y)) # reattach
g = hem.rescale(g, (-1, 1), (0, 1))
y = hem.rescale(y, (-1, 1), (0, 1))
g_sampler = tf.reshape(g_sampler, tf.shape(y)) # reattach
g_sampler = hem.rescale(g_sampler, (-1, 1), (0, 1))
y_sample = hem.rescale(y_sample, (-1, 1), (0, 1))
# add image montages
with arg_scope([hem.montage],
num_examples=args.examples,
height=n,
width=n,
colorize=True):
with tf.variable_scope('model'):
hem.montage(x, name='images')
hem.montage(y, name='real_depths')
hem.montage(g, name='fake_depths')
with tf.variable_scope('sampler'):
hem.montage(x_sample, name='images')
hem.montage(y_sample, name='real_depths')
hem.montage(g_sampler, name='fake_depths')
with tf.variable_scope('shuffled'):
hem.montage(x_shuffled, name='images')
hem.montage(g_shuffle, name='fake_depths')
with tf.variable_scope('noise'):
hem.montage(x_noise, name='images')
hem.montage(g_noise, name='fake_depths')
def cgan(x, args):
"""Create conditional GAN ('pix2pix') model on the graph.
Args:
x: Tensor, the real images.
args: Argparse structure
Returns:
Function, the training function. Call for one iteration of training.
"""
# 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()
# rescale to [-1, 1]
x = hem.rescale(x, (0, 1), (-1, 1))
for x, scope, gpu_id in hem.tower_scope_range(x, args.n_gpus,
args.batch_size):
# model
x_rgb, x_depth = _split(x)
# generator
with tf.variable_scope('generator'):
g = generator(x_rgb, args, reuse=(gpu_id > 0))
# discriminator
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, reuse=(gpu_id > 0))
# 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)
# average and apply gradients
g_grads = hem.average_gradients(g_tower_grads)
d_grads = hem.average_gradients(d_tower_grads)
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
_summaries(g, x, args)
hem.add_basic_summaries(g_grads + d_grads)
# training
return _train_cgan(g_apply_grads, d_apply_grads, batchnorm_updates)
def loss(d_real, d_real_logits, d_fake, d_fake_logits, x, g, x_depth, q, args, reuse=False):
def xentropy(logits, labels):
return tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=labels)
with tf.variable_scope('loss'):
g = hem.rescale(g, (-1, 1), (0, 1))
x_depth = hem.rescale(x_depth, (-1, 1), (0, 1))
rmse_loss = hem.rmse(x_depth, g)
l1_loss = tf.reduce_mean(tf.abs(x_depth - g), name='l1')
# losses
with tf.variable_scope('generator'):
g_fake = tf.reduce_mean(xentropy(d_fake_logits, tf.ones_like(d_fake)), name='g_fake')
# if args.g_sparsity:
# layers = tf.get_collection('conv_layers')
# for l in layers:
# if 'e5' in l.name:
# sparsity_term = tf.nn.zero_fraction(l, name='sparsity_term')
# lambda_term = 1.0
# if args.g_rmse:
# g_total = tf.identity(g_fake - lambda_term * sparsity_term + rmse_loss, name='g_total')
# else:
# g_total = tf.identity(g_fake - lambda_term * sparsity_term, name='g_total')
# elif args.g_rmse:
# g_total = tf.identity(g_fake + rmse_loss, name='g_total')
with tf.variable_scope('discriminator'):
d_real = tf.reduce_mean(xentropy(d_real_logits, tf.ones_like(d_real)), name='d_real')
d_fake = tf.reduce_mean(xentropy(d_fake_logits, tf.zeros_like(d_fake)), name='d_fake')
d_total = tf.identity(d_real + d_fake, name='d_total')
# only add these to the collection once
if not reuse:
hem.add_to_collection('losses', [g_fake, d_real, d_fake, d_total, rmse_loss, l1_loss])
# if args.g_sparsity:
# hem.add_to_collection('losses', [g_total, sparsity_term])
# elif args.g_rmse:
# hem.add_to_collection('losses', [g_total])
# if args.g_sparsity or args.g_rmse:
# return g_total, d_total
# else:
return g_fake, d_total
def _summaries(g, x, args):
"""Add specialized summaries to the graph.
This adds summaries that:
- Track variability in generator samples given random noise vectors.
- Track how much of the noise vector is used.
- Generate examples from the real and learned distributions.
Args:
g: Tensor, the generator's output. i.e., the fake images.
x: Tensor, the real (input) images.
args: Argparse structure.
Returns:
None
"""
# 1. generate multiple samples using a single image
with tf.variable_scope(tf.get_variable_scope()):
gpu_id = 0
with tf.device(hem.variables_on_cpu(gpu_id)):
with tf.name_scope('tower_{}'.format(gpu_id)) as scope:
# print('input shape', x.shape)
with tf.variable_scope('generator'):
# using first image in batch, form new one with just this image
x_repeated = tf.stack([x[0]] * args.batch_size)
x_rgb, x_depth = _split(x_repeated)
# then create a new path for the generator using just this dataset
import copy
args_copy = copy.copy(args)
args_copy.batch_norm = False
d = generator(x_rgb, args_copy, reuse=True)
# scale to [0, 1]
sampler_rgb, sampler_depth = _split(d)
sampler_rgb = hem.rescale(sampler_rgb, (-1, 1), (0, 1))
sampler_depth = hem.rescale(sampler_depth, (-1, 1), (0, 1))
with tf.variable_scope('sampler'):
hem.montage(sampler_rgb[0:args.examples], 8, 8, name='images')
hem.montage(sampler_depth[0:args.examples], 8, 8, name='depths')
# and track the variance of the depth predictions
mean, var = tf.nn.moments(sampler_depth, axes=[0])
hem.scalars(('predicted_depth_mean', tf.reduce_mean(mean)),
('predicted_depth_var', tf.reduce_mean(var)))
# and the images (temporary, for calibration/sanity check)
x_rgb, x_depth = _split(x)
mean, var = tf.nn.moments(x_depth, axes=[0])
hem.scalars(('real_depth_mean', tf.reduce_mean(mean)),
('real_depth_var', tf.reduce_mean(var)))
sampler_rgb = tf.transpose(sampler_rgb, [0, 2, 3, 1])
sampler_depth = tf.transpose(sampler_depth, [0, 2, 3, 1])
# mean, var = tf.nn.moments(tf.image.rgb_to_grayscale(sampler_rgb), axes=[0])
axis = [0, -1]
mean, var = tf.nn.moments(tf.image.rgb_to_grayscale(sampler_rgb),
axes=[0])
hem.scalars(('image_mean', tf.reduce_mean(mean)),
('image_var', tf.reduce_mean(var)))
mean, var = tf.nn.moments(sampler_depth, axes=[0])
hem.scalars(('depth_mean', tf.reduce_mean(mean)),
('depth_var', tf.reduce_mean(var)))
# 2. generate summaries for real and fake images
with tf.variable_scope('examples'):
hem.histograms(('fake', g), ('real', x))
# rescale, split, and colorize
x = hem.rescale(x[0:args.examples], (-1, 1), (0, 1))
g = hem.rescale(g[0:args.examples], (-1, 1), (0, 1))
hem.histograms(('fake_rescaled', g), ('real_rescaled', x))
x_rgb, x_depth = _split(x)
g_rgb, g_depth = _split(g)
# note: these are rescaled to (0, 1)
hem.histograms(('real_depth', x_depth), ('fake_depth', g_depth),
('real_rgb', x_rgb), ('fake_rgb', g_rgb))
# add montages
# TODO shouldn't be fixed to 8, but to ceil(sqrt(args.examples))
hem.montage(x_rgb, 8, 8, name='real/images')
hem.montage(x_depth, 8, 8, name='real/depths')
hem.montage(g_rgb, 8, 8, name='fake/images')
hem.montage(g_depth, 8, 8, name='fake/depths')
# 3. add additional summaries for weights and biases in e_c1 (the initial noise layer)
# TODO don't iterate through list, but grab directly by full name
with tf.variable_scope('noise'):
for l in tf.get_collection('weights'):
if 'e_c1' in hem.tensor_name(l):
print(l, hem.tensor_name(l))
x_rgb, x_noise = _split(l)
hem.histograms(('rgb/weights', x_rgb),
('noise/weights', x_noise))
hem.scalars(
('rgb/weights/sparsity', tf.nn.zero_fraction(x_rgb)),
('noise/weights/sparsity', tf.nn.zero_fraction(x_noise)),
('noise/weights/mean', tf.reduce_mean(x_noise)),
('rgb/weights/mean', tf.reduce_mean(x_rgb)))
break
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):
"""Create conditional GAN ('pix2pix') model on the graph.
Args:
x: Tensor, the real images.
args: Argparse structure
Returns:
Function, the training function. Call for one iteration of training.
"""
# 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()
# rescale to [-1, 1]
# x_y = hem.rescale(x_y, (0, 1), (-1, 1))
# foreach gpu...
for x_y, scope, gpu_id in hem.tower_scope_range(x_y, args.n_gpus, args.batch_size):
# split inputs and scale to [-1, 1]
x = hem.rescale(x_y[0], (0, 1), (-1, 1))
y = hem.rescale(x_y[1], (0, 1), (-1, 1))
# x, y = tf.split(x_y, num_or_size_splits=[3, 1], axis=1)
# repeated image tensor for sampling
x_sample = tf.stack([x[0]] * args.examples)
y_sample = tf.stack([y[0]] * args.examples)
# create model
with tf.variable_scope('generator'):
g = pix2pix.generator(x, args, reuse=(gpu_id > 0))
g_sampler = pix2pix.generator(x_sample, args, reuse=True)
with tf.variable_scope('discriminator'):
d_real, d_real_logits = pix2pix.discriminator(x, y, args, reuse=(gpu_id > 0))
d_fake, d_fake_logits = pix2pix.discriminator(x, g, args, reuse=True)
# losses
g_loss, d_loss = pix2pix.loss(d_real, d_real_logits, d_fake, d_fake_logits, g, y, args, reuse=(gpu_id > 0))
# 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)
# 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
pix2pix.montage_summaries(x, y, g, args, x_sample, y_sample, g_sampler, d_real, d_fake)
pix2pix.activation_summaries()
pix2pix.loss_summaries()
pix2pix.gradient_summaries(g_grads, 'generator_gradients')
pix2pix.gradient_summaries(d_grads, 'discriminator_gradients')
pix2pix.sampler_summaries(y_sample, g_sampler, args)
# training ops
with tf.control_dependencies(batchnorm_updates):
self.d_train_op = d_apply_grads
self.g_train_op = g_apply_grads
self.all_losses = hem.collection_to_dict(tf.get_collection('losses'))
