Function train_and_export

tensorflow/examples/tf2_showcase/mnist.py:184–234

Run MNIST training and eval loop in eager mode. Args: flags_obj: An object containing parsed flag values.

(flags_obj)

def train_and_export(flags_obj):
  """Run MNIST training and eval loop in eager mode.

  Args:
    flags_obj: An object containing parsed flag values.
  """
  # Load the datasets
  train_ds, test_ds = mnist_datasets()
  train_ds = train_ds.shuffle(60000).batch(flags_obj.batch_size)
  test_ds = test_ds.batch(flags_obj.batch_size)

  # Create the model and optimizer
  model = create_model()
  optimizer = tf.train.MomentumOptimizer(
      flags_obj.learning_rate, flags_obj.momentum)

  # See summaries with `tensorboard --logdir=<model_dir>`
  train_dir = os.path.join(flags_obj.model_dir, 'summaries', 'train')
  test_dir = os.path.join(flags_obj.model_dir, 'summaries', 'eval')
  summary_writer = tf.contrib.summary.create_file_writer(
      train_dir, flush_millis=10000)
  test_summary_writer = tf.contrib.summary.create_file_writer(
      test_dir, flush_millis=10000, name='test')

  # Create and restore checkpoint (if one exists on the path)
  checkpoint_dir = os.path.join(flags_obj.model_dir, 'checkpoints')
  checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
  step_counter = tf.train.get_or_create_global_step()
  checkpoint = tf.train.Checkpoint(
      model=model, optimizer=optimizer, step_counter=step_counter)
  # Restore variables on creation if a checkpoint exists.
  checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))

  # Train and evaluate for a set number of epochs.
  for _ in range(flags_obj.train_epochs):
    start = time.time()
    with summary_writer.as_default():
      train(model, optimizer, train_ds, step_counter,
            flags_obj.log_interval, num_steps=1)
    end = time.time()
    print('\nTrain time for epoch #%d (%d total steps): %f' %
          (checkpoint.save_counter.numpy() + 1,
           step_counter.numpy(),
           end - start))
    with test_summary_writer.as_default():
      test(model, test_ds)
    checkpoint.save(checkpoint_prefix)

  # TODO(brianklee): Enable this functionality after @allenl implements this.
  # export_path = os.path.join(flags_obj.model_dir, 'export')
  # tf.saved_model.save(export_path, model)


def import_and_eval(flags_obj):