# Copyright 2019 DeepMind Technologies Limited and Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Training script.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import os from absl import app from absl import flags from absl import logging import tensorflow.compat.v1 as tf from cs_gan import file_utils from cs_gan import gan from cs_gan import image_metrics from cs_gan import utils flags.DEFINE_integer( 'num_training_iterations', 1200000, 'Number of training iterations.') flags.DEFINE_string( 'ode_mode', 'rk4', 'Integration method.') flags.DEFINE_integer( 'batch_size', 64, 'Training batch size.') flags.DEFINE_float( 'grad_reg_weight', 0.02, 'Step size for latent optimisation.') flags.DEFINE_string( 'opt_name', 'gd', 'Name of the optimiser (gd|adam).') flags.DEFINE_bool( 'schedule_lr', True, 'The method to project z.') flags.DEFINE_bool( 'reg_first_grad_only', True, 'Whether only to regularise the first grad.') flags.DEFINE_integer( 'num_latents', 128, 'The number of latents') flags.DEFINE_integer( 'summary_every_step', 1000, 'The interval at which to log debug ops.') flags.DEFINE_integer( 'image_metrics_every_step', 1000, 'The interval at which to log (expensive) image metrics.') flags.DEFINE_integer( 'export_every', 10, 'The interval at which to export samples.') # Use 50k to reproduce scores from the paper. Default to 10k here to avoid the # runtime error caused by too large graph with 50k samples on some machines. flags.DEFINE_integer( 'num_eval_samples', 10000, 'The number of samples used to evaluate FID/IS.') flags.DEFINE_string( 'dataset', 'cifar', 'The dataset used for learning (cifar|mnist).') flags.DEFINE_string( 'output_dir', '/tmp/ode_gan/gan', 'Location where to save output files.') flags.DEFINE_float('disc_lr', 4e-2, 'Discriminator Learning rate.') flags.DEFINE_float('gen_lr', 4e-2, 'Generator Learning rate.') flags.DEFINE_bool( 'run_real_data_metrics', False, 'Whether or not to run image metrics on real data.') flags.DEFINE_bool( 'run_sample_metrics', True, 'Whether or not to run image metrics on samples.') FLAGS = flags.FLAGS # Log info level (for Hooks). tf.logging.set_verbosity(tf.logging.INFO) def _copy_vars(v_list): """Copy variables in v_list.""" t_list = [] for v in v_list: t_list.append(tf.identity(v)) return t_list def _restore_vars(v_list, t_list): """Restore variables in v_list from t_list.""" ops = [] for v, t in zip(v_list, t_list): ops.append(v.assign(t)) return ops def _scale_vars(s, v_list): """Scale all variables in v_list by s.""" return [s * v for v in v_list] def _acc_grads(g_sum, g_w, g): """Accumulate gradients in g, weighted by g_w.""" return [g_sum_i + g_w * g_i for g_sum_i, g_i in zip(g_sum, g)] def _compute_reg_grads(gen_grads, disc_vars): """Compute gradients norm (this is an upper-bpund of the full-batch norm).""" gen_norm = tf.accumulate_n([tf.reduce_sum(u * u) for u in gen_grads]) disc_reg_grads = tf.gradients(gen_norm, disc_vars) return disc_reg_grads def run_model(prior, images, model, disc_reg_weight): """Run the model with new data and samples. Args: prior: the noise source as the generator input. images: images sampled from dataset. model: a GAN model defined in gan.py. disc_reg_weight: regularisation weight for discrmininator gradients. Returns: debug_ops: statistics from the model, see gan.py for more detials. disc_grads: discriminator gradients. gen_grads: generator gradients. """ generator_inputs = prior.sample(FLAGS.batch_size) model_output = model.connect(images, generator_inputs) optimization_components = model_output.optimization_components disc_grads = tf.gradients( optimization_components['disc'].loss, optimization_components['disc'].vars) gen_grads = tf.gradients( optimization_components['gen'].loss, optimization_components['gen'].vars) if disc_reg_weight > 0.0: reg_grads = _compute_reg_grads(gen_grads, optimization_components['disc'].vars) disc_grads = _acc_grads(disc_grads, disc_reg_weight, reg_grads) debug_ops = model_output.debug_ops return debug_ops, disc_grads, gen_grads def update_model(model, disc_grads, gen_grads, disc_opt, gen_opt, global_step, update_scale): """Update model with gradients.""" disc_vars, gen_vars = model.get_variables() with tf.control_dependencies(gen_grads + disc_grads): disc_update_op = disc_opt.apply_gradients( zip(_scale_vars(update_scale, disc_grads), disc_vars)) gen_update_op = gen_opt.apply_gradients( zip(_scale_vars(update_scale, gen_grads), gen_vars), global_step=global_step) update_op = tf.group([disc_update_op, gen_update_op]) return update_op def main(argv): del argv utils.make_output_dir(FLAGS.output_dir) data_processor = utils.DataProcessor() # Compute the batch-size multiplier if FLAGS.ode_mode == 'rk2': batch_mul = 2 elif FLAGS.ode_mode == 'rk4': batch_mul = 4 else: batch_mul = 1 images = utils.get_train_dataset(data_processor, FLAGS.dataset, int(FLAGS.batch_size * batch_mul)) image_splits = tf.split(images, batch_mul) logging.info('Generator learning rate: %d', FLAGS.gen_lr) logging.info('Discriminator learning rate: %d', FLAGS.disc_lr) global_step = tf.train.get_or_create_global_step() # Construct optimizers. if FLAGS.opt_name == 'adam': disc_opt = tf.train.AdamOptimizer(FLAGS.disc_lr, beta1=0.5, beta2=0.999) gen_opt = tf.train.AdamOptimizer(FLAGS.gen_lr, beta1=0.5, beta2=0.999) elif FLAGS.opt_name == 'gd': if FLAGS.schedule_lr: gd_disc_lr = tf.train.piecewise_constant( global_step, values=[FLAGS.disc_lr / 4., FLAGS.disc_lr, FLAGS.disc_lr / 2.], boundaries=[500, 400000]) gd_gen_lr = tf.train.piecewise_constant( global_step, values=[FLAGS.gen_lr / 4., FLAGS.gen_lr, FLAGS.gen_lr / 2.], boundaries=[500, 400000]) else: gd_disc_lr = FLAGS.disc_lr gd_gen_lr = FLAGS.gen_lr disc_opt = tf.train.GradientDescentOptimizer(gd_disc_lr) gen_opt = tf.train.GradientDescentOptimizer(gd_gen_lr) else: raise ValueError('Unknown ODE mode!') # Create the networks and models. generator = utils.get_generator(FLAGS.dataset) metric_net = utils.get_metric_net(FLAGS.dataset, use_sn=False) model = gan.GAN(metric_net, generator) prior = utils.make_prior(FLAGS.num_latents) # Setup ODE parameters. if FLAGS.ode_mode == 'rk2': ode_grad_weights = [0.5, 0.5] step_scale = [1.0] elif FLAGS.ode_mode == 'rk4': ode_grad_weights = [1. / 6., 1. / 3., 1. / 3., 1. / 6.] step_scale = [0.5, 0.5, 1.] elif FLAGS.ode_mode == 'euler': # Euler update ode_grad_weights = [1.0] step_scale = [] else: raise ValueError('Unknown ODE mode!') # Extra steps for RK updates. num_extra_steps = len(step_scale) if FLAGS.reg_first_grad_only: first_reg_weight = FLAGS.grad_reg_weight / ode_grad_weights[0] other_reg_weight = 0.0 else: first_reg_weight = FLAGS.grad_reg_weight other_reg_weight = FLAGS.grad_reg_weight debug_ops, disc_grads, gen_grads = run_model(prior, image_splits[0], model, first_reg_weight) disc_vars, gen_vars = model.get_variables() final_disc_grads = _scale_vars(ode_grad_weights[0], disc_grads) final_gen_grads = _scale_vars(ode_grad_weights[0], gen_grads) restore_ops = [] # Preparing for further RK steps. if num_extra_steps > 0: # copy the variables before they are changed by update_op saved_disc_vars = _copy_vars(disc_vars) saved_gen_vars = _copy_vars(gen_vars) # Enter RK loop. with tf.control_dependencies(saved_disc_vars + saved_gen_vars): step_deps = [] for i_step in range(num_extra_steps): with tf.control_dependencies(step_deps): # Compute gradient steps for intermediate updates. update_op = update_model( model, disc_grads, gen_grads, disc_opt, gen_opt, None, step_scale[i_step]) with tf.control_dependencies([update_op]): _, disc_grads, gen_grads = run_model( prior, image_splits[i_step + 1], model, other_reg_weight) # Accumlate gradients for final update. final_disc_grads = _acc_grads(final_disc_grads, ode_grad_weights[i_step + 1], disc_grads) final_gen_grads = _acc_grads(final_gen_grads, ode_grad_weights[i_step + 1], gen_grads) # Make new restore_op for each step. restore_ops = [] restore_ops += _restore_vars(disc_vars, saved_disc_vars) restore_ops += _restore_vars(gen_vars, saved_gen_vars) step_deps = restore_ops with tf.control_dependencies(restore_ops): update_op = update_model( model, final_disc_grads, final_gen_grads, disc_opt, gen_opt, global_step, 1.0) samples = generator(prior.sample(FLAGS.batch_size), is_training=False) # Get data needed to compute FID. We also compute metrics on # real data as a sanity check and as a reference point. eval_real_data = utils.get_real_data_for_eval(FLAGS.num_eval_samples, FLAGS.dataset, split='train') def sample_fn(x): return utils.optimise_and_sample(x, module=model, data=None, is_training=False)[0] if FLAGS.run_sample_metrics: sample_metrics = image_metrics.get_image_metrics_for_samples( eval_real_data, sample_fn, prior, data_processor, num_eval_samples=FLAGS.num_eval_samples) else: sample_metrics = {} if FLAGS.run_real_data_metrics: data_metrics = image_metrics.get_image_metrics( eval_real_data, eval_real_data) else: data_metrics = {} sample_exporter = file_utils.FileExporter( os.path.join(FLAGS.output_dir, 'samples')) # Hooks. debug_ops['it'] = global_step # Abort training on Nans. nan_disc_hook = tf.train.NanTensorHook(debug_ops['disc_loss']) nan_gen_hook = tf.train.NanTensorHook(debug_ops['gen_loss']) # Step counter. step_conter_hook = tf.train.StepCounterHook() checkpoint_saver_hook = tf.train.CheckpointSaverHook( checkpoint_dir=utils.get_ckpt_dir(FLAGS.output_dir), save_secs=10 * 60) loss_summary_saver_hook = tf.train.SummarySaverHook( save_steps=FLAGS.summary_every_step, output_dir=os.path.join(FLAGS.output_dir, 'summaries'), summary_op=utils.get_summaries(debug_ops)) metrics_summary_saver_hook = tf.train.SummarySaverHook( save_steps=FLAGS.image_metrics_every_step, output_dir=os.path.join(FLAGS.output_dir, 'summaries'), summary_op=utils.get_summaries(sample_metrics)) hooks = [checkpoint_saver_hook, metrics_summary_saver_hook, nan_disc_hook, nan_gen_hook, step_conter_hook, loss_summary_saver_hook] # Start training. with tf.train.MonitoredSession(hooks=hooks) as sess: logging.info('starting training') for key, value in sess.run(data_metrics).items(): logging.info('%s: %d', key, value) for i in range(FLAGS.num_training_iterations): sess.run(update_op) if i % FLAGS.export_every == 0: samples_np, data_np = sess.run([samples, image_splits[0]]) # Create an object which gets data and does the processing. data_np = data_processor.postprocess(data_np) samples_np = data_processor.postprocess(samples_np) sample_exporter.save(samples_np, 'samples') sample_exporter.save(data_np, 'data') if __name__ == '__main__': tf.enable_resource_variables() app.run(main)