Add NHNet quality benchmark.

PiperOrigin-RevId: 312099777

official/benchmark/nhnet_benchmark.py

+# Lint as: python3
+# Copyright 2020 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.
+# ==============================================================================
+"""Executes benchmark testing for bert pretraining."""
+# pylint: disable=line-too-long
+from __future__ import print_function
+
+import time
+from typing import Optional
+
+from absl import flags
+import tensorflow as tf
+
+from official.benchmark import benchmark_wrappers
+from official.benchmark import owner_utils
+from official.benchmark import perfzero_benchmark
+from official.nlp.nhnet import trainer
+from official.utils.flags import core as flags_core
+
+MIN_LOSS = 0.35
+MAX_LOSS = 0.45
+NHNET_DATA = 'nhnet_training_data_files'
+
+FLAGS = flags.FLAGS
+
+
+class NHNetBenchmark(perfzero_benchmark.PerfZeroBenchmark):
+  """Base benchmark class for NHNet."""
+
+  def __init__(self, output_dir=None, default_flags=None, tpu=None):
+    self.default_flags = default_flags or {}
+    flag_methods = trainer.define_flags()
+    super(NHNetBenchmark, self).__init__(
+        output_dir=output_dir,
+        default_flags=default_flags,
+        flag_methods=flag_methods,
+        tpu=tpu)
+
+  def _report_benchmark(self,
+                        stats,
+                        wall_time_sec,
+                        max_value=None,
+                        min_value=None):
+    """Report benchmark results by writing to local protobuf file.
+
+    Args:
+      stats: dict returned from keras models with known entries.
+      wall_time_sec: the during of the benchmark execution in seconds
+      max_value: highest passing level.
+      min_value: lowest passing level.
+    """
+
+    metrics = []
+    metrics.append({
+        'name': 'training_loss',
+        'value': stats['training_loss'],
+        'min_value': min_value,
+        'max_value': max_value
+    })
+    # These metrics are placeholders to avoid PerfZero failure.
+    metrics.append({
+        'name': 'exp_per_second',
+        'value': 0.0,
+    })
+    metrics.append({
+        'name': 'startup_time',
+        'value': 9999.,
+    })
+    flags_str = flags_core.get_nondefault_flags_as_str()
+    self.report_benchmark(
+        iters=-1,
+        wall_time=wall_time_sec,
+        metrics=metrics,
+        extras={'flags': flags_str})
+
+
+class NHNetAccuracyBenchmark(NHNetBenchmark):
+  """Benchmark accuracy tests for NHNet."""
+
+  def __init__(self,
+               output_dir: Optional[str] = None,
+               tpu: Optional[str] = None,
+               **kwargs):
+    default_flags = dict(
+        mode='train',
+        train_file_pattern=NHNET_DATA,
+        train_batch_size=1024,
+        model_type='nhnet',
+        len_title=15,
+        len_passage=200,
+        num_encoder_layers=12,
+        num_decoder_layers=12,
+        num_nhnet_articles=5,
+        steps_per_loop=1000,
+        params_override='init_from_bert2bert=false')
+    super(NHNetAccuracyBenchmark, self).__init__(
+        output_dir=output_dir, default_flags=default_flags, tpu=tpu, **kwargs)
+
+  @benchmark_wrappers.enable_runtime_flags
+  def _run_and_report_benchmark(self, max_value=MAX_LOSS, min_value=MIN_LOSS):
+    """Runs and reports the benchmark given the provided configuration."""
+    start_time_sec = time.time()
+    stats = trainer.run()
+    wall_time_sec = time.time() - start_time_sec
+    self._report_benchmark(
+        stats, wall_time_sec, max_value=max_value, min_value=min_value)
+
+  @owner_utils.Owner('tf-model-garden')
+  def benchmark_accuracy_4x4_tpu_f32_50k_steps(self):
+    """Test bert pretraining with 4x4 TPU for 50k steps."""
+    # This is used for accuracy test.
+    self._setup()
+    FLAGS.train_steps = 50000
+    FLAGS.checkpoint_interval = FLAGS.train_steps
+    FLAGS.distribution_strategy = 'tpu'
+    FLAGS.model_dir = self._get_model_dir(
+        'benchmark_accuracy_4x4_tpu_bf32_50k_steps')
+    self._run_and_report_benchmark()
+
+  @owner_utils.Owner('tf-model-garden')
+  def benchmark_accuracy_4x4_tpu_f32_1k_steps(self):
+    """Test bert pretraining with 4x4 TPU for 1k steps."""
+    self._setup()
+    FLAGS.train_steps = 1000
+    FLAGS.checkpoint_interval = FLAGS.train_steps
+    FLAGS.distribution_strategy = 'tpu'
+    FLAGS.model_dir = self._get_model_dir(
+        'benchmark_accuracy_4x4_tpu_bf32_1k_steps')
+    self._run_and_report_benchmark()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/nlp/nhnet/trainer.py

@@ -162,12 +162,16 @@ def train(params, strategy, dataset=None):
   # Trains the model.
   steps_per_epoch = min(FLAGS.train_steps, FLAGS.checkpoint_interval)
   epochs = FLAGS.train_steps // steps_per_epoch
-  trainer.fit(
+  history = trainer.fit(
       x=dataset,
       steps_per_epoch=steps_per_epoch,
       epochs=epochs,
       callbacks=[summary_callback, checkpoint_callback],
       verbose=2)
+  train_hist = history.history
+  # Gets final loss from training.
+  stats = dict(training_loss=float(train_hist["training_loss"][-1]))
+  return stats
 
 
 def run():
@@ -197,7 +201,7 @@ def run():
       is_strict=False)
   stats = {}
   if "train" in FLAGS.mode:
-    train(params, strategy)
+    stats = train(params, strategy)
   if "eval" in FLAGS.mode:
     timeout = 0 if FLAGS.mode == "train_and_eval" else 3000
     # Uses padded decoding for TPU. Always uses cache.

official/nlp/nhnet/trainer_test.py

@@ -48,6 +48,7 @@ def get_trivial_data(config) -> tf.data.Dataset:
   batch_size, num_docs = 2, len(config.passage_list),
   len_passage = config.len_passage
   len_title = config.len_title
+
   def generate_data(_) -> tf.data.Dataset:
     fake_ids = tf.zeros((num_docs, len_passage), dtype=tf.int32)
     title = tf.zeros((len_title), dtype=tf.int32)
@@ -59,8 +60,8 @@ def get_trivial_data(config) -> tf.data.Dataset:
 
   dataset = tf.data.Dataset.range(1)
   dataset = dataset.repeat()
-  dataset = dataset.map(generate_data,
-                        num_parallel_calls=tf.data.experimental.AUTOTUNE)
+  dataset = dataset.map(
+      generate_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
   dataset = dataset.prefetch(buffer_size=1).batch(batch_size)
   return dataset
 
@@ -91,7 +92,9 @@ class TrainerTest(tf.test.TestCase, parameterized.TestCase):
     FLAGS.checkpoint_interval = 5
     FLAGS.model_dir = self.get_temp_dir()
     FLAGS.model_type = "nhnet"
-    trainer.train(self._config, distribution, get_trivial_data(self._config))
+    stats = trainer.train(self._config, distribution,
+                          get_trivial_data(self._config))
+    self.assertIn("training_loss", stats)
     self.assertLen(
         tf.io.gfile.glob(os.path.join(FLAGS.model_dir, "ckpt*.index")), 2)
 

official/recommendation/data_pipeline.py

@@ -37,7 +37,7 @@ from absl import logging
 
 from official.recommendation import constants as rconst
 from official.recommendation import movielens
-from official.recommendation import popen_helper
+from official.recommendation import popen_helper_internal as popen_helper
 from official.recommendation import stat_utils
 from tensorflow.python.tpu.datasets import StreamingFilesDataset
 

official/recommendation/data_test.py

@@ -30,7 +30,7 @@ import tensorflow as tf
 from official.recommendation import constants as rconst
 from official.recommendation import data_preprocessing
 from official.recommendation import movielens
-from official.recommendation import popen_helper
+from official.recommendation import popen_helper_internal as popen_helper
 from official.utils.misc import keras_utils