# Copyright 2017 Max Planck Society # Distributed under the BSD-3 Software license, # (See accompanying file ./LICENSE.txt or copy at # https://opensource.org/licenses/BSD-3-Clause) """This class implements Generative Adversarial Networks training. """ import logging import tensorflow as tf import utils from utils import ProgressBar from utils import TQDM import numpy as np import ops from metrics import Metrics class Gan(object): """A base class for running individual GANs. This class announces all the necessary bits for running individual GAN trainers. It is assumed that a GAN trainer should receive the data points and the corresponding weights, which are used for importance sampling of minibatches during the training. All the methods should be implemented in the subclasses. """ def __init__(self, opts, data, weights): # Create a new session with session.graph = default graph self._session = tf.Session() self._trained = False self._data = data self._data_weights = np.copy(weights) # Latent noise sampled ones to apply G while training self._noise_for_plots = utils.generate_noise(opts, 500) # Placeholders self._real_points_ph = None self._fake_points_ph = None self._noise_ph = None self._inv_target_ph = None # Main operations self._G = None # Generator function self._d_loss = None # Loss of discriminator self._g_loss = None # Loss of generator self._c_loss = None # Loss of mixture discriminator self._c_training = None # Outputs of the mixture discriminator on data self._inv_loss = None self._inv_loss_per_point = None # Variables self._inv_z = None # Optimizers self._g_optim = None self._d_optim = None self._c_optim = None self._inv_optim = None with self._session.as_default(), self._session.graph.as_default(): logging.debug('Building the graph...') self._build_model_internal(opts) if opts['inverse_metric']: assert opts['dataset'] in ('mnist', 'mnist3', 'guitars'),\ 'Invertion currently supported only for mnist, mnist3, guitars' logging.debug('Adding inversion ops to the graph...') self._add_inversion_ops(opts) # Make sure AdamOptimizer, if used in the Graph, is defined before # calling global_variables_initializer(). init = tf.global_variables_initializer() self._session.run(init) def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): # Cleaning the whole default Graph logging.debug('Cleaning the graph...') tf.reset_default_graph() logging.debug('Closing the session...') # Finishing the session self._session.close() def train(self, opts): """Train a GAN model. """ with self._session.as_default(), self._session.graph.as_default(): self._train_internal(opts) self._trained = True def sample(self, opts, num=100): """Sample points from the trained GAN model. """ assert self._trained, 'Can not sample from the un-trained GAN' with self._session.as_default(), self._session.graph.as_default(): return self._sample_internal(opts, num) def train_mixture_discriminator(self, opts, fake_images): """Train classifier separating true data from points in fake_images. Return: prob_real: probabilities of the points from training data being the real points according to the trained mixture classifier. Numpy vector of shape (self._data.num_points,) prob_fake: probabilities of the points from fake_images being the real points according to the trained mixture classifier. Numpy vector of shape (len(fake_images),) """ with self._session.as_default(), self._session.graph.as_default(): return self._train_mixture_discriminator_internal(opts, fake_images) def invert_points(self, opts, images): """Invert the learned generator function for every image in images. Args: images: numpy array of shape [num_points] + data_shape """ assert self._trained, 'Can not invert, not trained yet.' assert len(images) == opts['inverse_num'],\ 'Currently inversion works only for fixed number of images' with self._session.as_default(), self._session.graph.as_default(): target_ph = self._inv_target_ph z = self._inv_z loss_per_point = self._inv_loss_per_point optim = self._inv_optim norms = self._inv_norms val_list = [] err_per_point_list = [] z_list = [] norms_list = [] for _start in xrange(5): inv_vars = tf.get_collection( tf.GraphKeys.GLOBAL_VARIABLES, scope="inversion") # Initialize z and optimizer's variables randomly self._session.run(tf.variables_initializer(inv_vars)) prev_val = 100. check_every = 100 steps = 1 while True: # Stopping criterion: relative improvement of the maximal # per point mse gets smaller than a threshold self._session.run( optim, feed_dict={target_ph:images}) if steps % check_every == 0: err_per_point = loss_per_point.eval( feed_dict={target_ph:images}) err_max = np.max(err_per_point) err = np.mean(err_per_point) logging.debug('Init %02d, steps %d, loss %f, max mse %f' %\ (_start, steps, err, err_max)) relative_improvement = np.abs(prev_val - err) / prev_val if relative_improvement < 1e-3 or steps > 10000: val_list.append(err) err_per_point_list.append(err_per_point) z_list.append(self._session.run(z)) norms_list.append(self._session.run(norms)) break prev_val = err steps += 1 # Choose the run where we got the best (i.e. minimal) maximal # per point mse best_id = sorted(zip(val_list, range(len(val_list))))[0][1] best_err_per_point = err_per_point_list[best_id] best_z = z_list[best_id] best_norms = norms_list[best_id] best_reconstructions = self._G.eval( feed_dict={self._noise_ph:best_z, self._is_training_ph:False}) return best_reconstructions, best_z, best_err_per_point, best_norms def _add_inversion_ops(self, opts): data_shape = self._data.data_shape with tf.variable_scope("inversion"): target_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='target_ph') z = tf.get_variable( "inverted", [opts['inverse_num'], opts['latent_space_dim']], tf.float32, tf.random_normal_initializer(stddev=1.)) reconstructed_images = self.generator( opts, z, is_training=False, reuse=True) with tf.variable_scope("inversion"): loss_per_point = tf.reduce_mean( tf.square(tf.subtract(reconstructed_images, target_ph)), axis=[1, 2, 3]) loss = tf.reduce_mean(loss_per_point) norms = tf.reduce_sum(tf.square(z), axis=[1]) optim = tf.train.AdamOptimizer(0.01, 0.9) optim = optim.minimize(loss, var_list=[z]) self._inv_target_ph = target_ph self._inv_z = z self._inv_optim = optim self._inv_loss = loss self._inv_loss_per_point = loss_per_point self._inv_norms = norms def _run_batch(self, opts, operation, placeholder, feed, placeholder2=None, feed2=None): """Wrapper around session.run to process huge data. It is asumed that (a) first dimension of placeholder enumerates separate points, and (b) that operation is independently applied to every point, i.e. we can split it point-wisely and then merge the results. The second placeholder is meant either for is_train flag for batch-norm or probabilities of dropout. TODO: write util function which will be called both from this method and MNIST classification evaluation as well. """ assert len(feed.shape) > 0, 'Empry feed.' num_points = feed.shape[0] batch_size = opts['tf_run_batch_size'] batches_num = int(np.ceil((num_points + 0.) / batch_size)) result = [] for idx in xrange(batches_num): if idx == batches_num - 1: if feed2 is None: res = self._session.run( operation, feed_dict={placeholder: feed[idx * batch_size:]}) else: res = self._session.run( operation, feed_dict={placeholder: feed[idx * batch_size:], placeholder2: feed2}) else: if feed2 is None: res = self._session.run( operation, feed_dict={placeholder: feed[idx * batch_size: (idx + 1) * batch_size]}) else: res = self._session.run( operation, feed_dict={placeholder: feed[idx * batch_size: (idx + 1) * batch_size], placeholder2: feed2}) if len(res.shape) == 1: # convert (n,) vector to (n,1) array res = np.reshape(res, [-1, 1]) result.append(res) result = np.vstack(result) assert len(result) == num_points return result def _build_model_internal(self, opts): """Build a TensorFlow graph with all the necessary ops. """ assert False, 'Gan base class has no build_model method defined.' def _train_internal(self, opts): assert False, 'Gan base class has no train method defined.' def _sample_internal(self, opts, num): assert False, 'Gan base class has no sample method defined.' def _train_mixture_discriminator_internal(self, opts, fake_images): assert False, 'Gan base class has no mixture discriminator method defined.' class ToyGan(Gan): """A simple GAN implementation, suitable for toy datasets. """ def generator(self, opts, noise, reuse=False): """Generator function, suitable for simple toy experiments. Args: noise: [num_points, dim] array, where dim is dimensionality of the latent noise space. Returns: [num_points, dim1, dim2, dim3] array, where the first coordinate indexes the points, which all are of the shape (dim1, dim2, dim3). """ output_shape = self._data.data_shape num_filters = opts['g_num_filters'] with tf.variable_scope("GENERATOR", reuse=reuse): h0 = ops.linear(opts, noise, num_filters, 'h0_lin') h0 = tf.nn.relu(h0) h1 = ops.linear(opts, h0, num_filters, 'h1_lin') h1 = tf.nn.relu(h1) h2 = ops.linear(opts, h1, np.prod(output_shape), 'h2_lin') h2 = tf.reshape(h2, [-1] + list(output_shape)) if opts['input_normalize_sym']: return tf.nn.tanh(h2) else: return h2 def discriminator(self, opts, input_, prefix='DISCRIMINATOR', reuse=False): """Discriminator function, suitable for simple toy experiments. """ shape = input_.get_shape().as_list() num_filters = opts['d_num_filters'] assert len(shape) > 0, 'No inputs to discriminate.' with tf.variable_scope(prefix, reuse=reuse): h0 = ops.linear(opts, input_, num_filters, 'h0_lin') h0 = tf.nn.relu(h0) h1 = ops.linear(opts, h0, num_filters, 'h1_lin') h1 = tf.nn.relu(h1) h2 = ops.linear(opts, h1, 1, 'h2_lin') return h2 def _build_model_internal(self, opts): """Build the Graph corresponding to GAN implementation. """ data_shape = self._data.data_shape # Placeholders real_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') fake_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='fake_points_ph') noise_ph = tf.placeholder( tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph') # Operations G = self.generator(opts, noise_ph) d_logits_real = self.discriminator(opts, real_points_ph) d_logits_fake = self.discriminator(opts, G, reuse=True) c_logits_real = self.discriminator( opts, real_points_ph, prefix='CLASSIFIER') c_logits_fake = self.discriminator( opts, fake_points_ph, prefix='CLASSIFIER', reuse=True) c_training = tf.nn.sigmoid( self.discriminator(opts, real_points_ph, prefix='CLASSIFIER', reuse=True)) d_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real, labels=tf.ones_like(d_logits_real))) d_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.zeros_like(d_logits_fake))) d_loss = d_loss_real + d_loss_fake g_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.ones_like(d_logits_fake))) c_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_real, labels=tf.ones_like(c_logits_real))) c_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_fake, labels=tf.zeros_like(c_logits_fake))) c_loss = c_loss_real + c_loss_fake t_vars = tf.trainable_variables() d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name] g_vars = [var for var in t_vars if 'GENERATOR/' in var.name] d_optim = ops.optimizer(opts, 'd').minimize(d_loss, var_list=d_vars) g_optim = ops.optimizer(opts, 'g').minimize(g_loss, var_list=g_vars) c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) self._real_points_ph = real_points_ph self._fake_points_ph = fake_points_ph self._noise_ph = noise_ph self._G = G self._d_loss = d_loss self._g_loss = g_loss self._c_loss = c_loss self._c_training = c_training self._g_optim = g_optim self._d_optim = d_optim self._c_optim = c_optim def _train_internal(self, opts): """Train a GAN model. """ batches_num = self._data.num_points / opts['batch_size'] train_size = self._data.num_points counter = 0 logging.debug('Training GAN') for _epoch in xrange(opts["gan_epoch_num"]): for _idx in xrange(batches_num): data_ids = np.random.choice(train_size, opts['batch_size'], replace=False, p=self._data_weights) batch_images = self._data.data[data_ids].astype(np.float) batch_noise = utils.generate_noise(opts, opts['batch_size']) # Update discriminator parameters for _iter in xrange(opts['d_steps']): _ = self._session.run( self._d_optim, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise}) # Update generator parameters for _iter in xrange(opts['g_steps']): _ = self._session.run( self._g_optim, feed_dict={self._noise_ph: batch_noise}) counter += 1 if opts['verbose'] and counter % opts['plot_every'] == 0: metrics = Metrics() points_to_plot = self._run_batch( opts, self._G, self._noise_ph, self._noise_for_plots[0:320]) data_ids = np.random.choice(train_size, 320, replace=False, p=self._data_weights) metrics.make_plots( opts, counter, self._data.data[data_ids], points_to_plot, prefix='sample_e%04d_mb%05d_' % (_epoch, _idx)) def _sample_internal(self, opts, num): """Sample from the trained GAN model. """ noise = utils.generate_noise(opts, num) sample = self._run_batch(opts, self._G, self._noise_ph, noise) # sample = self._session.run( # self._G, feed_dict={self._noise_ph: noise}) return sample def _train_mixture_discriminator_internal(self, opts, fake_images): """Train a classifier separating true data from points in fake_images. """ batches_num = self._data.num_points / opts['batch_size'] logging.debug('Training a mixture discriminator') for epoch in xrange(opts["mixture_c_epoch_num"]): for idx in xrange(batches_num): ids = np.random.choice(len(fake_images), opts['batch_size'], replace=False) batch_fake_images = fake_images[ids] ids = np.random.choice(self._data.num_points, opts['batch_size'], replace=False) batch_real_images = self._data.data[ids] _ = self._session.run( self._c_optim, feed_dict={self._real_points_ph: batch_real_images, self._fake_points_ph: batch_fake_images}) res = self._run_batch( opts, self._c_training, self._real_points_ph, self._data.data) return res, None class ToyUnrolledGan(ToyGan): """A simple GAN implementation, suitable for toy datasets. """ def __init__(self, opts, data, weights): # Losses of the copied discriminator network self._d_loss_cp = None self._d_optim_cp = None # Rolling back ops (assign variable values fo true # to copied discriminator network) self._roll_back = None Gan.__init__(self, opts, data, weights) # Architecture used in unrolled gan paper def generator(self, opts, noise, reuse=False): """Generator function, suitable for simple toy experiments. Args: noise: [num_points, dim] array, where dim is dimensionality of the latent noise space. Returns: [num_points, dim1, dim2, dim3] array, where the first coordinate indexes the points, which all are of the shape (dim1, dim2, dim3). """ output_shape = self._data.data_shape num_filters = opts['g_num_filters'] with tf.variable_scope("GENERATOR", reuse=reuse): h0 = ops.linear(opts, noise, num_filters, 'h0_lin') h0 = tf.nn.tanh(h0) h1 = ops.linear(opts, h0, num_filters, 'h1_lin') h1 = tf.nn.tanh(h1) h2 = ops.linear(opts, h1, np.prod(output_shape), 'h2_lin') h2 = tf.reshape(h2, [-1] + list(output_shape)) if opts['input_normalize_sym']: return tf.nn.tanh(h2) else: return h2 def discriminator(self, opts, input_, prefix='DISCRIMINATOR', reuse=False): """Discriminator function, suitable for simple toy experiments. """ shape = input_.get_shape().as_list() num_filters = opts['d_num_filters'] assert len(shape) > 0, 'No inputs to discriminate.' with tf.variable_scope(prefix, reuse=reuse): h0 = ops.linear(opts, input_, num_filters, 'h0_lin') h0 = tf.nn.tanh(h0) h1 = ops.linear(opts, h0, num_filters, 'h1_lin') h1 = tf.nn.tanh(h1) h2 = ops.linear(opts, h1, 1, 'h2_lin') return h2 def _build_model_internal(self, opts): """Build the Graph corresponding to GAN implementation. """ data_shape = self._data.data_shape # Placeholders real_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') fake_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='fake_points_ph') noise_ph = tf.placeholder( tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph') # Operations G = self.generator(opts, noise_ph) d_logits_real = self.discriminator(opts, real_points_ph) d_logits_fake = self.discriminator(opts, G, reuse=True) # Disccriminator copy for the unrolling steps d_logits_real_cp = self.discriminator( opts, real_points_ph, prefix='DISCRIMINATOR_CP') d_logits_fake_cp = self.discriminator( opts, G, prefix='DISCRIMINATOR_CP', reuse=True) c_logits_real = self.discriminator( opts, real_points_ph, prefix='CLASSIFIER') c_logits_fake = self.discriminator( opts, fake_points_ph, prefix='CLASSIFIER', reuse=True) c_training = tf.nn.sigmoid( self.discriminator(opts, real_points_ph, prefix='CLASSIFIER', reuse=True)) d_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real, labels=tf.ones_like(d_logits_real))) d_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.zeros_like(d_logits_fake))) d_loss = d_loss_real + d_loss_fake d_loss_real_cp = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real_cp, labels=tf.ones_like(d_logits_real_cp))) d_loss_fake_cp = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake_cp, labels=tf.zeros_like(d_logits_fake_cp))) d_loss_cp = d_loss_real_cp + d_loss_fake_cp if opts['objective'] == 'JS': g_loss = - d_loss_cp elif opts['objective'] == 'JS_modified': g_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake_cp, labels=tf.ones_like(d_logits_fake_cp))) else: assert False, 'No objective %r implemented' % opts['objective'] c_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_real, labels=tf.ones_like(c_logits_real))) c_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_fake, labels=tf.zeros_like(c_logits_fake))) c_loss = c_loss_real + c_loss_fake t_vars = tf.trainable_variables() d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name] d_vars_cp = [var for var in t_vars if 'DISCRIMINATOR_CP/' in var.name] c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] g_vars = [var for var in t_vars if 'GENERATOR/' in var.name] # Ops to roll back the variable values of discriminator_cp # Will be executed each time before the unrolling steps with tf.variable_scope('assign'): roll_back = [] for var, var_cp in zip(d_vars, d_vars_cp): roll_back.append(tf.assign(var_cp, var)) d_optim = ops.optimizer(opts, 'd').minimize(d_loss, var_list=d_vars) d_optim_cp = ops.optimizer(opts, 'd').minimize( d_loss_cp, var_list=d_vars_cp) c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) g_optim = ops.optimizer(opts, 'g').minimize(g_loss, var_list=g_vars) # writer = tf.summary.FileWriter(opts['work_dir']+'/tensorboard', self._session.graph) self._real_points_ph = real_points_ph self._fake_points_ph = fake_points_ph self._noise_ph = noise_ph self._G = G self._roll_back = roll_back self._d_loss = d_loss self._d_loss_cp = d_loss_cp self._g_loss = g_loss self._c_loss = c_loss self._c_training = c_training self._g_optim = g_optim self._d_optim = d_optim self._d_optim_cp = d_optim_cp self._c_optim = c_optim logging.debug("Building Graph Done.") def _train_internal(self, opts): """Train a GAN model. """ batches_num = self._data.num_points / opts['batch_size'] train_size = self._data.num_points counter = 0 logging.debug('Training GAN') for _epoch in xrange(opts["gan_epoch_num"]): for _idx in TQDM(opts, xrange(batches_num), desc='Epoch %2d/%2d' %\ (_epoch+1, opts["gan_epoch_num"])): data_ids = np.random.choice(train_size, opts['batch_size'], replace=False, p=self._data_weights) batch_images = self._data.data[data_ids].astype(np.float) batch_noise = utils.generate_noise(opts, opts['batch_size']) # Update discriminator parameters for _iter in xrange(opts['d_steps']): _ = self._session.run( self._d_optim, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise}) # Roll back discriminator_cp's variables self._session.run(self._roll_back) # Unrolling steps for _iter in xrange(opts['unrolling_steps']): self._session.run( self._d_optim_cp, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise}) # Update generator parameters for _iter in xrange(opts['g_steps']): _ = self._session.run( self._g_optim, feed_dict={self._noise_ph: batch_noise}) counter += 1 if opts['verbose'] and counter % opts['plot_every'] == 0: metrics = Metrics() points_to_plot = self._run_batch( opts, self._G, self._noise_ph, self._noise_for_plots[0:320]) data_ids = np.random.choice(train_size, 320, replace=False, p=self._data_weights) metrics.make_plots( opts, counter, self._data.data[data_ids], points_to_plot, prefix='sample_e%04d_mb%05d_' % (_epoch, _idx)) class ImageGan(Gan): """A simple GAN implementation, suitable for pictures. """ def __init__(self, opts, data, weights): # One more placeholder for batch norm self._is_training_ph = None Gan.__init__(self, opts, data, weights) def generator(self, opts, noise, is_training, reuse=False): """Generator function, suitable for simple picture experiments. Args: noise: [num_points, dim] array, where dim is dimensionality of the latent noise space. is_training: bool, defines whether to use batch_norm in the train or test mode. Returns: [num_points, dim1, dim2, dim3] array, where the first coordinate indexes the points, which all are of the shape (dim1, dim2, dim3). """ output_shape = self._data.data_shape # (dim1, dim2, dim3) # Computing the number of noise vectors on-the-go dim1 = tf.shape(noise)[0] num_filters = opts['g_num_filters'] with tf.variable_scope("GENERATOR", reuse=reuse): height = output_shape[0] / 4 width = output_shape[1] / 4 h0 = ops.linear(opts, noise, num_filters * height * width, scope='h0_lin') h0 = tf.reshape(h0, [-1, height, width, num_filters]) h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1') # h0 = tf.nn.relu(h0) h0 = ops.lrelu(h0) _out_shape = [dim1, height * 2, width * 2, num_filters / 2] # for 28 x 28 does 7 x 7 --> 14 x 14 h1 = ops.deconv2d(opts, h0, _out_shape, scope='h1_deconv') h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2') # h1 = tf.nn.relu(h1) h1 = ops.lrelu(h1) _out_shape = [dim1, height * 4, width * 4, num_filters / 4] # for 28 x 28 does 14 x 14 --> 28 x 28 h2 = ops.deconv2d(opts, h1, _out_shape, scope='h2_deconv') h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3') # h2 = tf.nn.relu(h2) h2 = ops.lrelu(h2) _out_shape = [dim1] + list(output_shape) # data_shape[0] x data_shape[1] x ? -> data_shape h3 = ops.deconv2d(opts, h2, _out_shape, d_h=1, d_w=1, scope='h3_deconv') h3 = ops.batch_norm(opts, h3, is_training, reuse, scope='bn_layer4') if opts['input_normalize_sym']: return tf.nn.tanh(h3) else: return tf.nn.sigmoid(h3) def discriminator(self, opts, input_, is_training, prefix='DISCRIMINATOR', reuse=False): """Discriminator function, suitable for simple toy experiments. """ num_filters = opts['d_num_filters'] with tf.variable_scope(prefix, reuse=reuse): h0 = ops.conv2d(opts, input_, num_filters, scope='h0_conv') h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1') h0 = ops.lrelu(h0) h1 = ops.conv2d(opts, h0, num_filters * 2, scope='h1_conv') h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2') h1 = ops.lrelu(h1) h2 = ops.conv2d(opts, h1, num_filters * 4, scope='h2_conv') h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3') h2 = ops.lrelu(h2) h3 = ops.linear(opts, h2, 1, scope='h3_lin') return h3 def _build_model_internal(self, opts): """Build the Graph corresponding to GAN implementation. """ data_shape = self._data.data_shape # Placeholders real_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') fake_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='fake_points_ph') noise_ph = tf.placeholder( tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph') is_training_ph = tf.placeholder(tf.bool, name='is_train_ph') # Operations G = self.generator(opts, noise_ph, is_training_ph) # We use conv2d_transpose in the generator, which results in the # output tensor of undefined shapes. However, we statically know # the shape of the generator output, which is [-1, dim1, dim2, dim3] # where (dim1, dim2, dim3) is given by self._data.data_shape G.set_shape([None] + list(self._data.data_shape)) d_logits_real = self.discriminator(opts, real_points_ph, is_training_ph) d_logits_fake = self.discriminator(opts, G, is_training_ph, reuse=True) c_logits_real = self.discriminator( opts, real_points_ph, is_training_ph, prefix='CLASSIFIER') c_logits_fake = self.discriminator( opts, fake_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True) c_training = tf.nn.sigmoid( self.discriminator(opts, real_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True)) d_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real, labels=tf.ones_like(d_logits_real))) d_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.zeros_like(d_logits_fake))) d_loss = d_loss_real + d_loss_fake g_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.ones_like(d_logits_fake))) c_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_real, labels=tf.ones_like(c_logits_real))) c_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_fake, labels=tf.zeros_like(c_logits_fake))) c_loss = c_loss_real + c_loss_fake t_vars = tf.trainable_variables() d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name] g_vars = [var for var in t_vars if 'GENERATOR/' in var.name] d_optim = ops.optimizer(opts, 'd').minimize(d_loss, var_list=d_vars) g_optim = ops.optimizer(opts, 'g').minimize(g_loss, var_list=g_vars) # d_optim_op = ops.optimizer(opts, 'd') # g_optim_op = ops.optimizer(opts, 'g') # def debug_grads(grad, var): # _grad = tf.Print( # grad, # grads_and_vars, # [tf.global_norm([grad])], # 'Global grad norm of %s: ' % var.name) # return _grad, var # d_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # d_optim_op.compute_gradients(d_loss, var_list=d_vars)] # g_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # g_optim_op.compute_gradients(g_loss, var_list=g_vars)] # d_optim = d_optim_op.apply_gradients(d_grads_and_vars) # g_optim = g_optim_op.apply_gradients(g_grads_and_vars) c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) self._real_points_ph = real_points_ph self._fake_points_ph = fake_points_ph self._noise_ph = noise_ph self._is_training_ph = is_training_ph self._G = G self._d_loss = d_loss self._g_loss = g_loss self._c_loss = c_loss self._c_training = c_training self._g_optim = g_optim self._d_optim = d_optim self._c_optim = c_optim logging.debug("Building Graph Done.") def _train_internal(self, opts): """Train a GAN model. """ batches_num = self._data.num_points / opts['batch_size'] train_size = self._data.num_points counter = 0 logging.debug('Training GAN') for _epoch in xrange(opts["gan_epoch_num"]): for _idx in xrange(batches_num): # logging.debug('Step %d of %d' % (_idx, batches_num ) ) data_ids = np.random.choice(train_size, opts['batch_size'], replace=False, p=self._data_weights) batch_images = self._data.data[data_ids].astype(np.float) batch_noise = utils.generate_noise(opts, opts['batch_size']) # Update discriminator parameters for _iter in xrange(opts['d_steps']): _ = self._session.run( self._d_optim, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise, self._is_training_ph: True}) # Update generator parameters for _iter in xrange(opts['g_steps']): _ = self._session.run( self._g_optim, feed_dict={self._noise_ph: batch_noise, self._is_training_ph: True}) counter += 1 if opts['verbose'] and counter % opts['plot_every'] == 0: logging.debug( 'Epoch: %d/%d, batch:%d/%d' % \ (_epoch+1, opts['gan_epoch_num'], _idx+1, batches_num)) metrics = Metrics() points_to_plot = self._run_batch( opts, self._G, self._noise_ph, self._noise_for_plots[0:320], self._is_training_ph, False) metrics.make_plots( opts, counter, None, points_to_plot, prefix='sample_e%04d_mb%05d_' % (_epoch, _idx)) if opts['early_stop'] > 0 and counter > opts['early_stop']: break def _sample_internal(self, opts, num): """Sample from the trained GAN model. """ noise = utils.generate_noise(opts, num) sample = self._run_batch( opts, self._G, self._noise_ph, noise, self._is_training_ph, False) # sample = self._session.run( # self._G, feed_dict={self._noise_ph: noise}) return sample def _train_mixture_discriminator_internal(self, opts, fake_images): """Train a classifier separating true data from points in fake_images. """ batches_num = self._data.num_points / opts['batch_size'] logging.debug('Training a mixture discriminator') logging.debug('Using %d real points and %d fake ones' %\ (self._data.num_points, len(fake_images))) for epoch in xrange(opts["mixture_c_epoch_num"]): for idx in xrange(batches_num): ids = np.random.choice(len(fake_images), opts['batch_size'], replace=False) batch_fake_images = fake_images[ids] ids = np.random.choice(self._data.num_points, opts['batch_size'], replace=False) batch_real_images = self._data.data[ids] _ = self._session.run( self._c_optim, feed_dict={self._real_points_ph: batch_real_images, self._fake_points_ph: batch_fake_images, self._is_training_ph: True}) # Evaluating trained classifier on real points res = self._run_batch( opts, self._c_training, self._real_points_ph, self._data.data, self._is_training_ph, False) # Evaluating trained classifier on fake points res_fake = self._run_batch( opts, self._c_training, self._real_points_ph, fake_images, self._is_training_ph, False) return res, res_fake class MNISTLabelGan(ImageGan): """Architecture for MNIST from "Improved techniques for training GANs" """ def generator(self, opts, noise, is_training, reuse=False): with tf.variable_scope("GENERATOR", reuse=reuse): h0 = ops.linear(opts, noise, 100, scope='h0_lin') h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1', scale=False) h0 = tf.nn.softplus(h0) h1 = ops.linear(opts, h0, 100, scope='h1_lin') h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2', scale=False) h1 = tf.nn.softplus(h1) h2 = ops.linear(opts, h1, 28 * 28, scope='h2_lin') # h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3') h2 = tf.reshape(h2, [-1, 28, 28, 1]) if opts['input_normalize_sym']: return tf.nn.tanh(h2) else: return tf.nn.sigmoid(h2) def discriminator(self, opts, input_, is_training, prefix='DISCRIMINATOR', reuse=False): shape = tf.shape(input_) num = shape[0] with tf.variable_scope(prefix, reuse=reuse): h0 = input_ h0 = tf.add(h0, tf.random_normal(shape, stddev=0.3)) h0 = ops.linear(opts, h0, 1000, scope='h0_linear') # h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1') h0 = tf.nn.relu(h0) h1 = tf.add(h0, tf.random_normal([num, 1000], stddev=0.5)) h1 = ops.linear(opts, h1, 500, scope='h1_linear') # h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2') h1 = tf.nn.relu(h1) h2 = tf.add(h1, tf.random_normal([num, 500], stddev=0.5)) h2 = ops.linear(opts, h2, 250, scope='h2_linear') # h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3') h2 = tf.nn.relu(h2) h3 = tf.add(h2, tf.random_normal([num, 250], stddev=0.5)) h3 = ops.linear(opts, h3, 250, scope='h3_linear') # h3 = ops.batch_norm(opts, h3, is_training, reuse, scope='bn_layer4') h3 = tf.nn.relu(h3) h4 = tf.add(h3, tf.random_normal([num, 250], stddev=0.5)) h4 = ops.linear(opts, h4, 250, scope='h4_linear') # h4 = ops.batch_norm(opts, h4, is_training, reuse, scope='bn_layer5') h4 = tf.nn.relu(h4) h5 = ops.linear(opts, h4, 10, scope='h5_linear') return h5, h3 def _build_model_internal(self, opts): """Build the Graph corresponding to GAN implementation. """ data_shape = self._data.data_shape # Placeholders real_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') real_points_unl_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') fake_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='fake_points_ph') noise_ph = tf.placeholder( tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph') is_training_ph = tf.placeholder(tf.bool, name='is_train_ph') dropout_rate_ph = tf.placeholder(tf.float32) # labels_ph = tf.placeholder(tf.int8, [None, 10]) labels_ph = tf.placeholder(tf.int64, [None]) lr_ph = tf.placeholder(tf.float32) # Operations G = self.generator(opts, noise_ph, is_training_ph) # We use conv2d_transpose in the generator, which results in the # output tensor of undefined shapes. However, we statically know # the shape of the generator output, which is [-1, dim1, dim2, dim3] # where (dim1, dim2, dim3) is given by self._data.data_shape # G.set_shape([None] + list(self._data.data_shape)) # Here we follow a proposal of "Improved techniques for training # GANs" paper, Section 5 d_logits_real, _ = self.discriminator(opts, real_points_ph, is_training_ph) d_logits_real_unl, d_features_real_unl = self.discriminator( opts, real_points_unl_ph, is_training_ph, reuse=True) d_logits_fake, d_features_fake = self.discriminator( opts, G, is_training_ph, reuse=True) d_loss_labelled = tf.reduce_mean( tf.nn.sparse_softmax_cross_entropy_with_logits( logits=d_logits_real, labels=labels_ph)) correct_predictions = tf.equal( tf.argmax(d_logits_real, axis=1), # tf.argmax(labels_ph, axis=1)) labels_ph) d_accuracy = tf.reduce_mean(tf.cast(correct_predictions, tf.float32)) # 0 / 1 labels: # Z_real = ops.log_sum_exp(d_logits_real_unl) # Z_fake = ops.log_sum_exp(d_logits_fake) # D_real = Z_real - tf.nn.softplus(ops.log_sum_exp(d_logits_real_unl)) # D_real = tf.Print(D_real, [D_real], 'Res:') # D_fake = -tf.nn.softplus(ops.log_sum_exp(d_logits_fake)) # D_fake = tf.Print(D_fake, [D_fake]) # d_loss_unl = - tf.reduce_mean(D_real) - tf.reduce_mean(D_fake) # Label smoothing Z_real = ops.log_sum_exp(d_logits_real_unl) Z_fake = ops.log_sum_exp(d_logits_fake) cross_entropy_real_0 = Z_real - tf.nn.softplus( ops.log_sum_exp(d_logits_real_unl)) # cross_entropy_real_0 = tf.Print(cross_entropy_real_0, # [tf.exp(cross_entropy_real_0)], # 'D(X):') cross_entropy_real = 0.65 * cross_entropy_real_0 + 0.35 * ( -tf.nn.softplus(ops.log_sum_exp(d_logits_real_unl))) cross_entropy_fake_0 = -tf.nn.softplus( ops.log_sum_exp(d_logits_fake)) # cross_entropy_fake_0 = tf.Print(cross_entropy_fake_0, # [tf.exp(cross_entropy_fake_0)], # '1-D(G(Z)):') cross_entropy_fake = 1. * cross_entropy_fake_0 + 0. * ( Z_fake - tf.nn.softplus(ops.log_sum_exp(d_logits_fake))) d_loss_unl = - tf.reduce_mean(cross_entropy_fake) \ - tf.reduce_mean(cross_entropy_real) d_loss = d_loss_labelled + 0.5 * d_loss_unl # Log trick: # g_loss = -(ops.log_sum_exp(d_logits_fake) + cross_entropy_fake_0) # No log trick: # g_loss = tf.reduce_mean(cross_entropy_fake_0) # Feature matching f_mean_fake = tf.reduce_mean(d_features_fake, axis=0) f_mean_real = tf.reduce_mean(d_features_real_unl, axis=0) g_loss = tf.reduce_mean(tf.square(f_mean_fake - f_mean_real)) c_logits_real, _ = self.discriminator( opts, real_points_ph, is_training_ph, prefix='CLASSIFIER') c_logits_fake, _ = self.discriminator( opts, fake_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True) c_training_logits, _ = self.discriminator( opts, real_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True) CZ_real = ops.log_sum_exp(c_logits_real) CD_real = CZ_real - tf.nn.softplus(ops.log_sum_exp(c_logits_real)) CD_fake = -tf.nn.softplus(ops.log_sum_exp(c_logits_fake)) c_loss = - tf.reduce_mean(CD_real) - tf.reduce_mean(CD_fake) c_training = tf.exp(CD_real) # d_optim_op = ops.optimizer(opts, 'd') # g_optim_op = ops.optimizer(opts, 'g') # def debug_grads(grad, var): # _grad = tf.Print( # grad, # grads_and_vars, # [tf.global_norm([grad])], # 'Global grad norm of %s: ' % var.name) # return _grad, var t_vars = tf.trainable_variables() d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name] g_vars = [var for var in t_vars if 'GENERATOR/' in var.name] # d_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # d_optim_op.compute_gradients(d_loss, var_list=d_vars)] # g_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # g_optim_op.compute_gradients(g_loss, var_list=g_vars)] # d_optim = d_optim_op.apply_gradients(d_grads_and_vars) # g_optim = g_optim_op.apply_gradients(g_grads_and_vars) d_optim = tf.train.AdamOptimizer(lr_ph, beta1=opts["opt_beta1"]) g_optim = tf.train.AdamOptimizer(lr_ph, beta1=opts["opt_beta1"]) # g_optim = tf.train.GradientDescentOptimizer(lr_ph) d_optim = d_optim.minimize(d_loss, var_list=d_vars) g_optim = g_optim.minimize(g_loss, var_list=g_vars) c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) self._real_points_ph = real_points_ph self._fake_points_ph = fake_points_ph self._noise_ph = noise_ph self._real_points_unl_ph = real_points_unl_ph self._is_training_ph = is_training_ph self._dropout_rate_ph = dropout_rate_ph self._G = G self._d_loss = d_loss self._g_loss = g_loss self._c_loss = c_loss self._c_training = c_training self._g_optim = g_optim self._d_optim = d_optim self._c_optim = c_optim self._labels_ph = labels_ph self._d_accuracy = d_accuracy self._g_loss = g_loss self._lr_ph = lr_ph logging.debug("Building Graph Done.") def _train_internal(self, opts): """Train a GAN model. """ train_data = self._data.data[:60000] train_labels = self._data.labels[:60000] train_weights = self._data_weights[:60000] train_weights = train_weights / np.sum(train_weights) test_data = self._data.data[60000:] test_labels = self._data.labels[60000:] batches_num = len(train_data) / opts['batch_size'] train_size = len(train_data) counter = 0 logging.debug('Training GAN') lr_g = opts['opt_g_learning_rate'] lr_d = opts['opt_d_learning_rate'] accuracy = 0. for _epoch in xrange(opts["gan_epoch_num"]): for _idx in xrange(batches_num): # logging.debug('Step %d of %d' % (_idx, batches_num ) ) data_ids = np.random.choice(train_size, opts['batch_size'], replace=False, p=train_weights) data_ids_unl = np.random.choice(train_size, opts['batch_size'], replace=False, p=train_weights) batch_images = train_data[data_ids].astype(np.float) batch_images_unl = train_data[data_ids_unl].astype(np.float) batch_noise = utils.generate_noise(opts, opts['batch_size']) # Update discriminator parameters # labels_oh = utils.one_hot(self._data.labels[data_ids]) labels_oh = train_labels[data_ids] lr = lr_d * min(1., 1. - ((0. + _epoch) / opts['gan_epoch_num'])) for _iter in xrange(opts['d_steps']): _ = self._session.run( self._d_optim, feed_dict={self._real_points_ph: batch_images, self._real_points_unl_ph: batch_images_unl, self._is_training_ph: True, self._lr_ph: lr, self._noise_ph: batch_noise, self._labels_ph: labels_oh}) # Update generator parameters lr = lr_g * min(1., 1. - ((0. + _epoch) / opts['gan_epoch_num'])) for _iter in xrange(opts['g_steps']): _ = self._session.run( self._g_optim, feed_dict={self._noise_ph: batch_noise, self._is_training_ph: True, self._lr_ph: lr, self._real_points_unl_ph: batch_images_unl}) counter += 1 if opts['verbose'] and counter % opts['plot_every'] == 0: accuracy = self._d_accuracy.eval( feed_dict={self._real_points_ph: test_data, self._is_training_ph: False, # self._labels_ph: utils.one_hot(self._data.labels[:1000])}) self._labels_ph: test_labels}) g_loss = self._g_loss.eval( feed_dict={self._noise_ph: batch_noise, self._is_training_ph: False, self._real_points_unl_ph: batch_images_unl}) logging.debug( 'Epoch:%3d/%d, batch:%4d/%d, lr_g=%.4f, D accuracy in telling digits:%f, G feature matching loss:%f' % \ (_epoch+1, opts['gan_epoch_num'], _idx+1, batches_num, lr, accuracy, g_loss)) metrics = Metrics() points_to_plot = self._run_batch( opts, self._G, self._noise_ph, self._noise_for_plots[0:320], self._is_training_ph, False) metrics.make_plots( opts, counter, None, points_to_plot, prefix='sample_e%04d_mb%05d_' % (_epoch, _idx)) if opts['early_stop'] > 0 and counter > opts['early_stop']: break def _train_mixture_discriminator_internal(self, opts, fake_images): """Train a classifier separating true data from points in fake_images. """ batches_num = self._data.num_points / opts['batch_size'] logging.debug('Training a mixture discriminator') logging.debug('Using %d real points and %d fake ones' %\ (self._data.num_points, len(fake_images))) for epoch in xrange(opts["mixture_c_epoch_num"]): for idx in xrange(batches_num): ids = np.random.choice(len(fake_images), opts['batch_size'], replace=False) batch_fake_images = fake_images[ids] ids = np.random.choice(self._data.num_points, opts['batch_size'], replace=False) batch_real_images = self._data.data[ids] _ = self._session.run( self._c_optim, feed_dict={self._real_points_ph: batch_real_images, self._fake_points_ph: batch_fake_images, self._is_training_ph: True}) # Evaluating trained classifier on real points res = self._run_batch( opts, self._c_training, self._real_points_ph, self._data.data, self._is_training_ph, False) # Evaluating trained classifier on fake points res_fake = self._run_batch( opts, self._c_training, self._real_points_ph, fake_images, self._is_training_ph, False) return res, res_fake class BigImageGan(ImageGan): """A bit more flexible generator, compared to ImageGan. """ def generator(self, opts, noise, is_training, reuse=False): """Generator function, suitable for bigger simple pictures. Args: noise: [num_points, dim] array, where dim is dimensionality of the latent noise space. is_training: bool, defines whether to use batch_norm in the train or test mode. Returns: [num_points, dim1, dim2, dim3] array, where the first coordinate indexes the points, which all are of the shape (dim1, dim2, dim3). """ output_shape = self._data.data_shape # (dim1, dim2, dim3) # Computing the number of noise vectors on-the-go dim1 = tf.shape(noise)[0] num_filters = opts['g_num_filters'] with tf.variable_scope("GENERATOR", reuse=reuse): height = output_shape[0] / 16 width = output_shape[1] / 16 h0 = ops.linear(opts, noise, num_filters * height * width, scope='h0_lin') h0 = tf.reshape(h0, [-1, height, width, num_filters]) h0 = ops.batch_norm(opts, h0, is_training, reuse, scope='bn_layer1') h0 = tf.nn.relu(h0) _out_shape = [dim1, height * 2, width * 2, num_filters / 2] # for 128 x 128 does 8 x 8 --> 16 x 16 h1 = ops.deconv2d(opts, h0, _out_shape, scope='h1_deconv') h1 = ops.batch_norm(opts, h1, is_training, reuse, scope='bn_layer2') h1 = tf.nn.relu(h1) _out_shape = [dim1, height * 4, width * 4, num_filters / 4] # for 128 x 128 does 16 x 16 --> 32 x 32 h2 = ops.deconv2d(opts, h1, _out_shape, scope='h2_deconv') h2 = ops.batch_norm(opts, h2, is_training, reuse, scope='bn_layer3') h2 = tf.nn.relu(h2) _out_shape = [dim1, height * 8, width * 8, num_filters / 8] # for 128 x 128 does 32 x 32 --> 64 x 64 h3 = ops.deconv2d(opts, h2, _out_shape, scope='h3_deconv') h3 = ops.batch_norm(opts, h3, is_training, reuse, scope='bn_layer4') h3 = tf.nn.relu(h3) _out_shape = [dim1, height * 16, width * 16, num_filters / 16] # for 128 x 128 does 64 x 64 --> 128 x 128 h4 = ops.deconv2d(opts, h3, _out_shape, scope='h4_deconv') h4 = ops.batch_norm(opts, h4, is_training, reuse, scope='bn_layer5') h4 = tf.nn.relu(h4) _out_shape = [dim1] + list(output_shape) # data_shape[0] x data_shape[1] x ? -> data_shape h5 = ops.deconv2d(opts, h4, _out_shape, d_h=1, d_w=1, scope='h5_deconv') h5 = ops.batch_norm(opts, h5, is_training, reuse, scope='bn_layer6') if opts['input_normalize_sym']: return tf.nn.tanh(h5) else: return tf.nn.sigmoid(h5) class ImageUnrolledGan(ImageGan): """A simple GAN implementation, suitable for pictures. """ def __init__(self, opts, data, weights): # Losses of the copied discriminator network self._d_loss_cp = None self._d_optim_cp = None # Rolling back ops (assign variable values fo true # to copied discriminator network) self._roll_back = None ImageGan.__init__(self, opts, data, weights) def _build_model_internal(self, opts): """Build the Graph corresponding to GAN implementation. """ data_shape = self._data.data_shape # Placeholders real_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='real_points_ph') fake_points_ph = tf.placeholder( tf.float32, [None] + list(data_shape), name='fake_points_ph') noise_ph = tf.placeholder( tf.float32, [None] + [opts['latent_space_dim']], name='noise_ph') is_training_ph = tf.placeholder(tf.bool, name='is_train_ph') # Operations G = self.generator(opts, noise_ph, is_training_ph) # We use conv2d_transpose in the generator, which results in the # output tensor of undefined shapes. However, we statically know # the shape of the generator output, which is [-1, dim1, dim2, dim3] # where (dim1, dim2, dim3) is given by self._data.data_shape G.set_shape([None] + list(self._data.data_shape)) d_logits_real = self.discriminator(opts, real_points_ph, is_training_ph) d_logits_fake = self.discriminator(opts, G, is_training_ph, reuse=True) # Disccriminator copy for the unrolling steps d_logits_real_cp = self.discriminator( opts, real_points_ph, is_training_ph, prefix='DISCRIMINATOR_CP') d_logits_fake_cp = self.discriminator( opts, G, is_training_ph, prefix='DISCRIMINATOR_CP', reuse=True) c_logits_real = self.discriminator( opts, real_points_ph, is_training_ph, prefix='CLASSIFIER') c_logits_fake = self.discriminator( opts, fake_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True) c_training = tf.nn.sigmoid( self.discriminator(opts, real_points_ph, is_training_ph, prefix='CLASSIFIER', reuse=True)) d_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real, labels=tf.ones_like(d_logits_real))) d_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake, labels=tf.zeros_like(d_logits_fake))) d_loss = d_loss_real + d_loss_fake d_loss_real_cp = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_real_cp, labels=tf.ones_like(d_logits_real_cp))) d_loss_fake_cp = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake_cp, labels=tf.zeros_like(d_logits_fake_cp))) d_loss_cp = d_loss_real_cp + d_loss_fake_cp if opts['objective'] == 'JS': g_loss = - d_loss_cp elif opts['objective'] == 'JS_modified': g_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=d_logits_fake_cp, labels=tf.ones_like(d_logits_fake_cp))) else: assert False, 'No objective %r implemented' % opts['objective'] c_loss_real = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_real, labels=tf.ones_like(c_logits_real))) c_loss_fake = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits( logits=c_logits_fake, labels=tf.zeros_like(c_logits_fake))) c_loss = c_loss_real + c_loss_fake t_vars = tf.trainable_variables() d_vars = [var for var in t_vars if 'DISCRIMINATOR/' in var.name] d_vars_cp = [var for var in t_vars if 'DISCRIMINATOR_CP/' in var.name] c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] g_vars = [var for var in t_vars if 'GENERATOR/' in var.name] # Ops to roll back the variable values of discriminator_cp # Will be executed each time before the unrolling steps with tf.variable_scope('assign'): roll_back = [] for var, var_cp in zip(d_vars, d_vars_cp): roll_back.append(tf.assign(var_cp, var)) d_optim = ops.optimizer(opts, 'd').minimize(d_loss, var_list=d_vars) d_optim_cp = ops.optimizer(opts, 'd').minimize( d_loss_cp, var_list=d_vars_cp) c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) g_optim = ops.optimizer(opts, 'g').minimize(g_loss, var_list=g_vars) # writer = tf.summary.FileWriter(opts['work_dir']+'/tensorboard', self._session.graph) # d_optim_op = ops.optimizer(opts, 'd') # g_optim_op = ops.optimizer(opts, 'g') # def debug_grads(grad, var): # _grad = tf.Print( # grad, # grads_and_vars, # [tf.global_norm([grad])], # 'Global grad norm of %s: ' % var.name) # return _grad, var # d_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # d_optim_op.compute_gradients(d_loss, var_list=d_vars)] # g_grads_and_vars = [debug_grads(grad, var) for (grad, var) in \ # g_optim_op.compute_gradients(g_loss, var_list=g_vars)] # d_optim = d_optim_op.apply_gradients(d_grads_and_vars) # g_optim = g_optim_op.apply_gradients(g_grads_and_vars) c_vars = [var for var in t_vars if 'CLASSIFIER/' in var.name] c_optim = ops.optimizer(opts).minimize(c_loss, var_list=c_vars) self._real_points_ph = real_points_ph self._fake_points_ph = fake_points_ph self._noise_ph = noise_ph self._is_training_ph = is_training_ph self._G = G self._roll_back = roll_back self._d_loss = d_loss self._d_loss_cp = d_loss_cp self._g_loss = g_loss self._c_loss = c_loss self._c_training = c_training self._g_optim = g_optim self._d_optim = d_optim self._d_optim_cp = d_optim_cp self._c_optim = c_optim logging.debug("Building Graph Done.") def _train_internal(self, opts): """Train a GAN model. """ batches_num = self._data.num_points / opts['batch_size'] train_size = self._data.num_points counter = 0 logging.debug('Training GAN') for _epoch in xrange(opts["gan_epoch_num"]): for _idx in TQDM(opts, xrange(batches_num), desc='Epoch %2d/%2d' %\ (_epoch + 1, opts["gan_epoch_num"])): # logging.debug('Step %d of %d' % (_idx, batches_num ) ) data_ids = np.random.choice(train_size, opts['batch_size'], replace=False, p=self._data_weights) batch_images = self._data.data[data_ids].astype(np.float) batch_noise = utils.generate_noise(opts, opts['batch_size']) # Update discriminator parameters for _iter in xrange(opts['d_steps']): _ = self._session.run( self._d_optim, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise, self._is_training_ph: True}) # Roll back discriminator_cp's variables self._session.run(self._roll_back) # Unrolling steps for _iter in xrange(opts['unrolling_steps']): self._session.run( self._d_optim_cp, feed_dict={self._real_points_ph: batch_images, self._noise_ph: batch_noise, self._is_training_ph: True}) # Update generator parameters for _iter in xrange(opts['g_steps']): _ = self._session.run( self._g_optim, feed_dict={self._noise_ph: batch_noise, self._is_training_ph: True}) counter += 1 if opts['verbose'] and counter % opts['plot_every'] == 0: logging.debug( 'Epoch: %d/%d, batch:%d/%d' % \ (_epoch+1, opts['gan_epoch_num'], _idx+1, batches_num)) metrics = Metrics() points_to_plot = self._run_batch( opts, self._G, self._noise_ph, self._noise_for_plots[0:320], self._is_training_ph, False) metrics.make_plots( opts, counter, None, points_to_plot, prefix='sample_e%04d_mb%05d_' % (_epoch, _idx)) if opts['early_stop'] > 0 and counter > opts['early_stop']: break