Fix/log ex per sec (tensorflow#4360)

karmel · web-flow · commit d626b9084570 · 2018-05-25T09:41:24.000-07:00
* Using BenchmarkLogger

* Using BenchmarkLogger

* Fixing tests

* Linting fixes.

* Adding comments

* Moving mock logger

* Moving mock logger

* Glinting

* Responding to CR

* Reverting assertEmpty
diff --git a/official/utils/logs/hooks.py b/official/utils/logs/hooks.py
@@ -20,7 +20,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import tensorflow as tf
+import tensorflow as tf  # pylint: disable=g-bad-import-order
+
+from official.utils.logs import logger
 
 
 class ExamplesPerSecondHook(tf.train.SessionRunHook):
@@ -36,7 +38,8 @@ def __init__(self,
                batch_size,
                every_n_steps=None,
                every_n_secs=None,
-               warm_steps=0):
+               warm_steps=0,
+               metric_logger=None):
     """Initializer for ExamplesPerSecondHook.
 
     Args:
@@ -48,15 +51,20 @@ def __init__(self,
       warm_steps: The number of steps to be skipped before logging and running
         average calculation. warm_steps steps refers to global steps across all
         workers, not on each worker
+      metric_logger: instance of `BenchmarkLogger`, the benchmark logger that
+          hook should use to write the log. If None, BaseBenchmarkLogger will
+          be used.
 
     Raises:
       ValueError: if neither `every_n_steps` or `every_n_secs` is set, or
       both are set.
     """
 
     if (every_n_steps is None) == (every_n_secs is None):
-      raise ValueError('exactly one of every_n_steps'
-                       ' and every_n_secs should be provided.')
+      raise ValueError("exactly one of every_n_steps"
+                       " and every_n_secs should be provided.")
+
+    self._logger = metric_logger or logger.BaseBenchmarkLogger()
 
     self._timer = tf.train.SecondOrStepTimer(
         every_steps=every_n_steps, every_secs=every_n_secs)
@@ -71,7 +79,7 @@ def begin(self):
     self._global_step_tensor = tf.train.get_global_step()
     if self._global_step_tensor is None:
       raise RuntimeError(
-          'Global step should be created to use StepCounterHook.')
+          "Global step should be created to use StepCounterHook.")
 
   def before_run(self, run_context):  # pylint: disable=unused-argument
     """Called before each call to run().
@@ -109,7 +117,11 @@ def after_run(self, run_context, run_values):  # pylint: disable=unused-argument
         # and training time per batch
         current_examples_per_sec = self._batch_size * (
             elapsed_steps / elapsed_time)
-        # Current examples/sec followed by average examples/sec
-        tf.logging.info('Batch [%g]:  current exp/sec = %g, average exp/sec = '
-                        '%g', self._total_steps, current_examples_per_sec,
-                        average_examples_per_sec)
+
+        self._logger.log_metric(
+            "average_examples_per_sec", average_examples_per_sec,
+            global_step=global_step)
+
+        self._logger.log_metric(
+            "current_examples_per_sec", current_examples_per_sec,
+            global_step=global_step)
diff --git a/official/utils/logs/hooks_helper.py b/official/utils/logs/hooks_helper.py
@@ -119,9 +119,9 @@ def get_examples_per_second_hook(every_n_steps=100,
     Returns a ProfilerHook that writes out timelines that can be loaded into
     profiling tools like chrome://tracing.
   """
-  return hooks.ExamplesPerSecondHook(every_n_steps=every_n_steps,
-                                     batch_size=batch_size,
-                                     warm_steps=warm_steps)
+  return hooks.ExamplesPerSecondHook(
+      batch_size=batch_size, every_n_steps=every_n_steps,
+      warm_steps=warm_steps, metric_logger=logger.get_benchmark_logger())
 
 
 def get_logging_metric_hook(tensors_to_log=None,
diff --git a/official/utils/logs/hooks_test.py b/official/utils/logs/hooks_test.py
@@ -25,77 +25,74 @@
 from tensorflow.python.training import monitored_session  # pylint: disable=g-bad-import-order
 
 from official.utils.logs import hooks
+from official.utils.testing import mock_lib
 
 
-tf.logging.set_verbosity(tf.logging.ERROR)
+tf.logging.set_verbosity(tf.logging.DEBUG)
 
 
 class ExamplesPerSecondHookTest(tf.test.TestCase):
   """Tests for the ExamplesPerSecondHook."""
 
   def setUp(self):
     """Mock out logging calls to verify if correct info is being monitored."""
-    self._actual_log = tf.logging.info
-    self.logged_message = None
-
-    def mock_log(*args, **kwargs):
-      self.logged_message = args
-      self._actual_log(*args, **kwargs)
-
-    tf.logging.info = mock_log
+    self._logger = mock_lib.MockBenchmarkLogger()
 
     self.graph = tf.Graph()
     with self.graph.as_default():
       self.global_step = tf.train.get_or_create_global_step()
       self.train_op = tf.assign_add(self.global_step, 1)
 
-  def tearDown(self):
-    tf.logging.info = self._actual_log
-
   def test_raise_in_both_secs_and_steps(self):
     with self.assertRaises(ValueError):
       hooks.ExamplesPerSecondHook(
           batch_size=256,
           every_n_steps=10,
-          every_n_secs=20)
+          every_n_secs=20,
+          metric_logger=self._logger)
 
   def test_raise_in_none_secs_and_steps(self):
     with self.assertRaises(ValueError):
       hooks.ExamplesPerSecondHook(
           batch_size=256,
           every_n_steps=None,
-          every_n_secs=None)
+          every_n_secs=None,
+          metric_logger=self._logger)
 
   def _validate_log_every_n_steps(self, sess, every_n_steps, warm_steps):
     hook = hooks.ExamplesPerSecondHook(
         batch_size=256,
         every_n_steps=every_n_steps,
-        warm_steps=warm_steps)
+        warm_steps=warm_steps,
+        metric_logger=self._logger)
     hook.begin()
     mon_sess = monitored_session._HookedSession(sess, [hook])  # pylint: disable=protected-access
     sess.run(tf.global_variables_initializer())
 
-    self.logged_message = ''
     for _ in range(every_n_steps):
       mon_sess.run(self.train_op)
-      self.assertEqual(str(self.logged_message).find('exp/sec'), -1)
+      # Nothing should be in the list yet
+      self.assertFalse(self._logger.logged_metric)
 
     mon_sess.run(self.train_op)
     global_step_val = sess.run(self.global_step)
-    # assertNotRegexpMatches is not supported by python 3.1 and later
+
     if global_step_val > warm_steps:
-      self.assertRegexpMatches(str(self.logged_message), 'exp/sec')
+      self._assert_metrics()
     else:
-      self.assertEqual(str(self.logged_message).find('exp/sec'), -1)
+      # Nothing should be in the list yet
+      self.assertFalse(self._logger.logged_metric)
 
     # Add additional run to verify proper reset when called multiple times.
-    self.logged_message = ''
+    prev_log_len = len(self._logger.logged_metric)
     mon_sess.run(self.train_op)
     global_step_val = sess.run(self.global_step)
     if every_n_steps == 1 and global_step_val > warm_steps:
-      self.assertRegexpMatches(str(self.logged_message), 'exp/sec')
+      # Each time, we log two additional metrics. Did exactly 2 get added?
+      self.assertEqual(len(self._logger.logged_metric), prev_log_len + 2)
     else:
-      self.assertEqual(str(self.logged_message).find('exp/sec'), -1)
+      # No change in the size of the metric list.
+      self.assertEqual(len(self._logger.logged_metric), prev_log_len)
 
     hook.end(sess)
 
@@ -119,19 +116,19 @@ def _validate_log_every_n_secs(self, sess, every_n_secs):
     hook = hooks.ExamplesPerSecondHook(
         batch_size=256,
         every_n_steps=None,
-        every_n_secs=every_n_secs)
+        every_n_secs=every_n_secs,
+        metric_logger=self._logger)
     hook.begin()
     mon_sess = monitored_session._HookedSession(sess, [hook])  # pylint: disable=protected-access
     sess.run(tf.global_variables_initializer())
 
-    self.logged_message = ''
     mon_sess.run(self.train_op)
-    self.assertEqual(str(self.logged_message).find('exp/sec'), -1)
+    # Nothing should be in the list yet
+    self.assertFalse(self._logger.logged_metric)
     time.sleep(every_n_secs)
 
-    self.logged_message = ''
     mon_sess.run(self.train_op)
-    self.assertRegexpMatches(str(self.logged_message), 'exp/sec')
+    self._assert_metrics()
 
     hook.end(sess)
 
@@ -143,6 +140,11 @@ def test_examples_per_sec_every_5_secs(self):
     with self.graph.as_default(), tf.Session() as sess:
       self._validate_log_every_n_secs(sess, 5)
 
+  def _assert_metrics(self):
+    metrics = self._logger.logged_metric
+    self.assertEqual(metrics[-2]["name"], "average_examples_per_sec")
+    self.assertEqual(metrics[-1]["name"], "current_examples_per_sec")
+
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   tf.test.main()
diff --git a/official/utils/logs/logger.py b/official/utils/logs/logger.py
@@ -47,20 +47,20 @@
 
 
 def config_benchmark_logger(flag_obj=None):
-  """Config the global benchmark logger"""
+  """Config the global benchmark logger."""
   _logger_lock.acquire()
   try:
     global _benchmark_logger
     if not flag_obj:
       flag_obj = FLAGS
 
-    if (not hasattr(flag_obj, 'benchmark_logger_type') or
-        flag_obj.benchmark_logger_type == 'BaseBenchmarkLogger'):
+    if (not hasattr(flag_obj, "benchmark_logger_type") or
+        flag_obj.benchmark_logger_type == "BaseBenchmarkLogger"):
       _benchmark_logger = BaseBenchmarkLogger()
-    elif flag_obj.benchmark_logger_type == 'BenchmarkFileLogger':
+    elif flag_obj.benchmark_logger_type == "BenchmarkFileLogger":
       _benchmark_logger = BenchmarkFileLogger(flag_obj.benchmark_log_dir)
-    elif flag_obj.benchmark_logger_type == 'BenchmarkBigQueryLogger':
-      from official.benchmark import benchmark_uploader as bu # pylint: disable=g-import-not-at-top
+    elif flag_obj.benchmark_logger_type == "BenchmarkBigQueryLogger":
+      from official.benchmark import benchmark_uploader as bu  # pylint: disable=g-import-not-at-top
       bq_uploader = bu.BigQueryUploader(gcp_project=flag_obj.gcp_project)
       _benchmark_logger = BenchmarkBigQueryLogger(
           bigquery_uploader=bq_uploader,
@@ -69,8 +69,8 @@ def config_benchmark_logger(flag_obj=None):
           bigquery_metric_table=flag_obj.bigquery_metric_table,
           run_id=str(uuid.uuid4()))
     else:
-      raise ValueError('Unrecognized benchmark_logger_type: %s',
-                       flag_obj.benchmark_logger_type)
+      raise ValueError("Unrecognized benchmark_logger_type: %s"
+                       % flag_obj.benchmark_logger_type)
 
   finally:
     _logger_lock.release()
@@ -247,6 +247,7 @@ def log_run_info(self, model_name, dataset_name, run_params):
          self._run_id,
          run_info))
 
+
 def _gather_run_info(model_name, dataset_name, run_params):
   """Collect the benchmark run information for the local environment."""
   run_info = {
@@ -303,6 +304,7 @@ def process_param(name, value):
     run_info["run_parameters"] = [
         process_param(k, v) for k, v in sorted(run_params.items())]
 
+
 def _collect_tensorflow_environment_variables(run_info):
   run_info["tensorflow_environment_variables"] = [
       {"name": k, "value": v}
diff --git a/official/utils/logs/metric_hook_test.py b/official/utils/logs/metric_hook_test.py
diff --git a/official/utils/testing/mock_lib.py b/official/utils/testing/mock_lib.py
