Function main

tensorflow/contrib/eager/python/examples/gan/mnist.py:265–308

(_)

def main(_):
  (device, data_format) = ('/gpu:0', 'channels_first')
  if FLAGS.no_gpu or tf.contrib.eager.num_gpus() <= 0:
    (device, data_format) = ('/cpu:0', 'channels_last')
  print('Using device %s, and data format %s.' % (device, data_format))

  # Load the datasets
  data = input_data.read_data_sets(FLAGS.data_dir)
  dataset = (
      tf.data.Dataset.from_tensor_slices(data.train.images).shuffle(60000)
      .batch(FLAGS.batch_size))

  # Create the models and optimizers.
  model_objects = {
      'generator': Generator(data_format),
      'discriminator': Discriminator(data_format),
      'generator_optimizer': tf.compat.v1.train.AdamOptimizer(FLAGS.lr),
      'discriminator_optimizer': tf.compat.v1.train.AdamOptimizer(FLAGS.lr),
      'step_counter': tf.train.get_or_create_global_step(),
  }

  # Prepare summary writer and checkpoint info
  summary_writer = tf.contrib.summary.create_summary_file_writer(
      FLAGS.output_dir, flush_millis=1000)
  checkpoint_prefix = os.path.join(FLAGS.checkpoint_dir, 'ckpt')
  latest_cpkt = tf.train.latest_checkpoint(FLAGS.checkpoint_dir)
  if latest_cpkt:
    print('Using latest checkpoint at ' + latest_cpkt)
  checkpoint = tf.train.Checkpoint(**model_objects)
  # Restore variables on creation if a checkpoint exists.
  checkpoint.restore(latest_cpkt)

  with tf.device(device):
    for _ in range(100):
      start = time.time()
      with summary_writer.as_default():
        train_one_epoch(dataset=dataset, log_interval=FLAGS.log_interval,
                        noise_dim=FLAGS.noise, **model_objects)
      end = time.time()
      checkpoint.save(checkpoint_prefix)
      print('\nTrain time for epoch #%d (step %d): %f' %
            (checkpoint.save_counter.numpy(),
             checkpoint.step_counter.numpy(),
             end - start))


if __name__ == '__main__':