text: # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Utility functions or classes shared between BERT benchmarks."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import time
  
  # pylint: disable=g-bad-import-order
  
  import numpy as np
  from absl import flags
  import tensorflow as tf
  # pylint: enable=g-bad-import-order
  
  from official.utils.flags import core as flags_core
  from official.benchmark.perfzero_benchmark import PerfZeroBenchmark
  
  FLAGS = flags.FLAGS
  
  
  class BenchmarkTimerCallback(tf.keras.callbacks.Callback):
    """Callback that records time it takes to run each batch."""
  
    def __init__(self, num_batches_to_skip=10):
      super(BenchmarkTimerCallback, self).__init__()
      self.batch_start_times = {}
      self.batch_stop_times = {}
  
    def on_batch_begin(self, batch, logs=None):
      self.batch_start_times[batch] = time.time()
  
    def on_batch_end(self, batch, logs=None):
      # If there are multiple steps_per_loop, the end batch index will not be the
      # same as the starting index. Use the last starting index instead.
      if batch not in self.batch_start_times:
        batch = max(self.batch_start_times.keys())
  
      self.batch_stop_times[batch] = time.time()
  
    def get_examples_per_sec(self, batch_size, num_batches_to_skip=1):
      batch_durations = []
      for batch in self.batch_start_times:
        if batch in self.batch_stop_times and batch >= num_batches_to_skip:
          batch_durations.append(self.batch_stop_times[batch] -
                                 self.batch_start_times[batch])
      return batch_size / np.mean(batch_durations)
  
    def get_startup_time(self, program_start_time):
      return self.batch_start_times[0] - program_start_time
  
  
  class BertBenchmarkBase(PerfZeroBenchmark):
    """Base class to hold methods common to test classes."""
    local_flags = None
  
    def __init__(self, output_dir=None, tpu=None, **kwargs):
      super(BertBenchmarkBase, self).__init__(
          output_dir=output_dir, tpu=tpu, **kwargs)
      self.num_gpus = 8
      self.timer_callback = None
  
    def _setup(self):
      """Sets up and resets flags before each test."""
      super(BertBenchmarkBase, self)._setup()
      self.timer_callback = BenchmarkTimerCallback()
  
    def _report_benchmark(self, stats, wall_time_sec, min_accuracy, max_accuracy):
      """Report benchmark results by writing to local protobuf file.
  
      Args:
        stats: dict returned from BERT models with known entries.
        wall_time_sec: the during of the benchmark execution in seconds
        min_accuracy: Minimum classification accuracy constraint to verify
          correctness of the model.
        max_accuracy: Maximum classification accuracy constraint to verify
          correctness of the model.
      """
      metrics = [{
          'name': 'training_loss',
          'value': stats['train_loss'],
      }]
      if self.timer_callback:
        metrics.append({
            'name':
                'exp_per_second',
            'value':
                self.timer_callback.get_examples_per_sec(FLAGS.train_batch_size *
                                                         FLAGS.steps_per_loop)
        })
      else:
        metrics.append({
            'name': 'exp_per_second',
            'value': 0.0,
        })
      if self.timer_callback and 'start_time_sec' in stats:
        metrics.append({
            'name': 'startup_time',
            'value': self.timer_callback.get_startup_time(stats['start_time_sec'])
        })
  
      if 'eval_metrics' in stats:
        metrics.append({
            'name': 'eval_accuracy',
            'value': stats['eval_metrics'],
            'min_value': min_accuracy,
            'max_value': max_accuracy,
        })
      flags_str = flags_core.get_nondefault_flags_as_str()
      self.report_benchmark(
          iters=stats['total_training_steps'],
          wall_time=wall_time_sec,
          metrics=metrics,
          extras={'flags': flags_str})
  