Add end-to-end tests for Estimator. (#6023)

* Add end-to-end tests for Estimator.

* comment from keras to estimator.

* remove trailing new line

official/resnet/cifar10_main.py

@@ -250,6 +250,9 @@ def run_cifar(flags_obj):
 
   Args:
     flags_obj: An object containing parsed flag values.
+
+  Returns:
+    Dictionary of results. Including final accuracy.
   """
   if flags_obj.image_bytes_as_serving_input:
     tf.logging.fatal('--image_bytes_as_serving_input cannot be set to True '
@@ -259,10 +262,12 @@ def run_cifar(flags_obj):
   input_function = (flags_obj.use_synthetic_data and
                     get_synth_input_fn(flags_core.get_tf_dtype(flags_obj)) or
                     input_fn)
-  resnet_run_loop.resnet_main(
+  result = resnet_run_loop.resnet_main(
       flags_obj, cifar10_model_fn, input_function, DATASET_NAME,
       shape=[HEIGHT, WIDTH, NUM_CHANNELS])
 
+  return result
+
 
 def main(_):
   with logger.benchmark_context(flags.FLAGS):

official/resnet/estimator_cifar_benchmark.py

+"""Executes Estimator benchmarks and accuracy tests."""
+from __future__ import print_function
+
+import os
+
+from absl import flags
+from absl.testing import flagsaver
+import tensorflow as tf  # pylint: disable=g-bad-import-order
+
+from official.resnet import cifar10_main as cifar_main
+
+DATA_DIR = '/data/cifar10_data/'
+
+
+class EstimatorCifar10BenchmarkTests(object):
+  """Benchmarks and accuracy tests for Estimator ResNet56."""
+
+  local_flags = None
+
+  def __init__(self, output_dir=None):
+    self.oss_report_object = None
+    self.output_dir = output_dir
+
+  def resnet56_1_gpu(self):
+    """Test layers model with Estimator and distribution strategies."""
+    self._setup()
+    flags.FLAGS.num_gpus = 1
+    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.batch_size = 128
+    flags.FLAGS.train_epochs = 182
+    flags.FLAGS.model_dir = self._get_model_dir('resnet56_1_gpu')
+    flags.FLAGS.resnet_size = 56
+    flags.FLAGS.dtype = 'fp32'
+    stats = cifar_main.run_cifar(flags.FLAGS)
+    self._fill_report_object(stats)
+
+  def resnet56_fp16_1_gpu(self):
+    """Test layers FP16 model with Estimator and distribution strategies."""
+    self._setup()
+    flags.FLAGS.num_gpus = 1
+    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.batch_size = 128
+    flags.FLAGS.train_epochs = 182
+    flags.FLAGS.model_dir = self._get_model_dir('resnet56_fp16_1_gpu')
+    flags.FLAGS.resnet_size = 56
+    flags.FLAGS.dtype = 'fp16'
+    stats = cifar_main.run_cifar(flags.FLAGS)
+    self._fill_report_object(stats)
+
+  def resnet56_2_gpu(self):
+    """Test layers model with Estimator and dist_strat. 2 GPUs."""
+    self._setup()
+    flags.FLAGS.num_gpus = 1
+    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.batch_size = 128
+    flags.FLAGS.train_epochs = 182
+    flags.FLAGS.model_dir = self._get_model_dir('resnet56_2_gpu')
+    flags.FLAGS.resnet_size = 56
+    flags.FLAGS.dtype = 'fp32'
+    stats = cifar_main.run_cifar(flags.FLAGS)
+    self._fill_report_object(stats)
+
+  def resnet56_fp16_2_gpu(self):
+    """Test layers FP16 model with Estimator and dist_strat. 2 GPUs."""
+    self._setup()
+    flags.FLAGS.num_gpus = 2
+    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.batch_size = 128
+    flags.FLAGS.train_epochs = 182
+    flags.FLAGS.model_dir = self._get_model_dir('resnet56_fp16_2_gpu')
+    flags.FLAGS.resnet_size = 56
+    flags.FLAGS.dtype = 'fp16'
+    stats = cifar_main.run_cifar(flags.FLAGS)
+    self._fill_report_object(stats)
+
+  def _fill_report_object(self, stats):
+    # Also "available global_step"
+    if self.oss_report_object:
+      self.oss_report_object.top_1 = stats['accuracy'].item()
+      self.oss_report_object.top_5 = stats['accuracy_top_5'].item()
+    else:
+      raise ValueError('oss_report_object has not been set.')
+
+  def _get_model_dir(self, folder_name):
+    return os.path.join(self.output_dir, folder_name)
+
+  def _setup(self):
+    tf.logging.set_verbosity(tf.logging.DEBUG)
+    if EstimatorCifar10BenchmarkTests.local_flags is None:
+      cifar_main.define_cifar_flags()
+      # Loads flags to get defaults to then override.
+      flags.FLAGS(['foo'])
+      saved_flag_values = flagsaver.save_flag_values()
+      EstimatorCifar10BenchmarkTests.local_flags = saved_flag_values
+      return
+    flagsaver.restore_flag_values(EstimatorCifar10BenchmarkTests.local_flags)

official/resnet/resnet_run_loop.py

@@ -452,6 +452,9 @@ def resnet_main(
       used for logging purpose.
     shape: list of ints representing the shape of the images used for training.
       This is only used if flags_obj.export_dir is passed.
+
+  Returns:
+    Dict of results of the run.
   """
 
   model_helpers.apply_clean(flags.FLAGS)
@@ -588,7 +591,7 @@ def resnet_main(
           shape, batch_size=flags_obj.batch_size, dtype=export_dtype)
     classifier.export_savedmodel(flags_obj.export_dir, input_receiver_fn,
                                  strip_default_attrs=True)
-
+  return eval_results
 
 def define_resnet_flags(resnet_size_choices=None):
   """Add flags and validators for ResNet."""