Add eager for keras benchmark (tensorflow#4825)

yhliang2018 · web-flow · commit 2689c9aee73e · 2018-07-20T09:47:27.000-07:00
* Add more arguments

* Add eager mode

* Add notes for eager mode

* Address the comments

* Fix argument typos

* Add warning for eager and multi-gpu

* Fix typo

* Fix notes

* Fix pylint
diff --git a/official/keras_application_models/benchmark_main.py b/official/keras_application_models/benchmark_main.py
@@ -54,6 +54,11 @@ def run_keras_model_benchmark(_):
     raise AssertionError("The --model command line argument should "
                          "be a key in the `MODELS` dictionary.")
 
+  # Check if eager execution is enabled
+  if FLAGS.eager:
+    tf.logging.info("Eager execution is enabled...")
+    tf.enable_eager_execution()
+
   # Load the model
   tf.logging.info("Benchmark on {} model...".format(FLAGS.model))
   keras_model = MODELS[FLAGS.model]
@@ -64,30 +69,35 @@ def run_keras_model_benchmark(_):
   if FLAGS.use_synthetic_data:
     tf.logging.info("Using synthetic dataset...")
     dataset_name += "_Synthetic"
-    train_num_images = FLAGS.batch_size
-    val_num_images = FLAGS.batch_size
     train_dataset = dataset.generate_synthetic_input_dataset(
-        FLAGS.model, train_num_images)
+        FLAGS.model, FLAGS.batch_size)
     val_dataset = dataset.generate_synthetic_input_dataset(
-        FLAGS.model, val_num_images)
+        FLAGS.model, FLAGS.batch_size)
   else:
     raise ValueError("Only synthetic dataset is supported!")
 
   # If run with multiple GPUs
+  # If eager execution is enabled, only one GPU is utilized even if multiple
+  # GPUs are provided.
   num_gpus = flags_core.get_num_gpus(FLAGS)
-  if num_gpus > 0:
+  if num_gpus > 1:
+    if FLAGS.eager:
+      tf.logging.warning(
+          "{} GPUs are provided, but only one GPU is utilized as "
+          "eager execution is enabled.".format(num_gpus))
     model = tf.keras.utils.multi_gpu_model(model, gpus=num_gpus)
 
-  # Configure the model
   model.compile(loss="categorical_crossentropy",
-                optimizer="sgd",
+                optimizer=tf.train.AdamOptimizer(),
                 metrics=["accuracy"])
 
   # Create benchmark logger for benchmark logging
   run_params = {
       "batch_size": FLAGS.batch_size,
       "synthetic_data": FLAGS.use_synthetic_data,
-      "train_epochs": FLAGS.train_epochs
+      "train_epochs": FLAGS.train_epochs,
+      "num_train_images": FLAGS.num_images,
+      "num_eval_images": FLAGS.num_images,
   }
 
   benchmark_logger = logger.get_benchmark_logger()
@@ -108,8 +118,8 @@ def run_keras_model_benchmark(_):
       epochs=FLAGS.train_epochs,
       callbacks=callbacks,
       validation_data=val_dataset,
-      steps_per_epoch=int(np.ceil(train_num_images / FLAGS.batch_size)),
-      validation_steps=int(np.ceil(val_num_images / FLAGS.batch_size))
+      steps_per_epoch=int(np.ceil(FLAGS.num_images / FLAGS.batch_size)),
+      validation_steps=int(np.ceil(FLAGS.num_images / FLAGS.batch_size))
   )
 
   tf.logging.info("Logging the evaluation results...")
@@ -118,7 +128,7 @@ def run_keras_model_benchmark(_):
         "accuracy": history.history["val_acc"][epoch],
         "loss": history.history["val_loss"][epoch],
         tf.GraphKeys.GLOBAL_STEP: (epoch + 1) * np.ceil(
-            train_num_images/FLAGS.batch_size)
+            FLAGS.num_images/FLAGS.batch_size)
     }
     benchmark_logger.log_evaluation_result(eval_results)
 
@@ -146,6 +156,18 @@ def define_keras_benchmark_flags():
       help=flags_core.help_wrap(
           "Model to be benchmarked."))
 
+  flags.DEFINE_integer(
+      name="num_images", default=1000,
+      help=flags_core.help_wrap(
+          "The number of synthetic images for training and evaluation. The "
+          "default value is 1000."))
+
+  flags.DEFINE_boolean(
+      name="eager", default=False, help=flags_core.help_wrap(
+          "To enable eager execution. Note that if eager execution is enabled, "
+          "only one GPU is utilized even if multiple GPUs are provided and "
+          "multi_gpu_model is used."))
+
   flags.DEFINE_list(
       name="callbacks",
       default=["ExamplesPerSecondCallback", "LoggingMetricCallback"],
@@ -159,6 +181,7 @@ def main(_):
   with logger.benchmark_context(FLAGS):
     run_keras_model_benchmark(FLAGS)
 
+
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
   define_keras_benchmark_flags()
diff --git a/official/keras_application_models/dataset.py b/official/keras_application_models/dataset.py
@@ -19,6 +19,8 @@
 
 import tensorflow as tf
 
+from official.utils.misc import model_helpers  # pylint: disable=g-bad-import-order
+
 # Default values for dataset.
 _NUM_CHANNELS = 3
 _NUM_CLASSES = 1000
@@ -34,12 +36,14 @@ def _get_default_image_size(model):
   return image_size
 
 
-def generate_synthetic_input_dataset(model, num_imgs):
+def generate_synthetic_input_dataset(model, batch_size):
   """Generate synthetic dataset."""
   image_size = _get_default_image_size(model)
-  input_shape = (num_imgs,) + image_size + (_NUM_CHANNELS,)
-
-  images = tf.zeros(input_shape, dtype=tf.float32)
-  labels = tf.zeros((num_imgs, _NUM_CLASSES), dtype=tf.float32)
-
-  return tf.data.Dataset.from_tensors((images, labels)).repeat()
+  image_shape = (batch_size,) + image_size + (_NUM_CHANNELS,)
+  label_shape = (batch_size, _NUM_CLASSES)
+
+  return model_helpers.generate_synthetic_data(
+      input_shape=tf.TensorShape(image_shape),
+      input_dtype=tf.float32,
+      label_shape=tf.TensorShape(label_shape),
+      label_dtype=tf.float32)
