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