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 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 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 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 __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 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
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'))