Skip to content

GitLab

Commit 46245913 authored by Hongkun Yu's avatar Hongkun Yu Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 314466805
parent 20897493
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
official/benchmark/bert_benchmark.py
View file @ 46245913
...@@ -56,7 +56,7 @@ class BertClassifyBenchmarkBase(benchmark_utils.BertBenchmarkBase): ...@@ -56,7 +56,7 @@ class BertClassifyBenchmarkBase(benchmark_utils.BertBenchmarkBase):
super(BertClassifyBenchmarkBase, self).__init__(output_dir, tpu=tpu) super(BertClassifyBenchmarkBase, self).__init__(output_dir, tpu=tpu)
self.num_epochs = None self.num_epochs = None
self.num_steps_per_epoch = None self.num_steps_per_epoch = None
FLAGS.steps_per_loop = 50 FLAGS.steps_per_loop = 1
@flagsaver.flagsaver @flagsaver.flagsaver
def _run_bert_classifier(self, callbacks=None, use_ds=True): def _run_bert_classifier(self, callbacks=None, use_ds=True):
...@@ -269,6 +269,7 @@ class BertClassifyBenchmarkReal(BertClassifyBenchmarkBase): ...@@ -269,6 +269,7 @@ class BertClassifyBenchmarkReal(BertClassifyBenchmarkBase):
"""Test BERT model performance with 2x2 TPU.""" """Test BERT model performance with 2x2 TPU."""
self._setup() self._setup()
FLAGS.steps_per_loop = 50
FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu_mrpc') FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu_mrpc')
FLAGS.train_data_path = self.train_data_path FLAGS.train_data_path = self.train_data_path
FLAGS.eval_data_path = self.eval_data_path FLAGS.eval_data_path = self.eval_data_path
...@@ -352,6 +353,7 @@ class BertClassifyAccuracy(BertClassifyBenchmarkBase): ...@@ -352,6 +353,7 @@ class BertClassifyAccuracy(BertClassifyBenchmarkBase):
def benchmark_2x2_tpu_mrpc(self): def benchmark_2x2_tpu_mrpc(self):
"""Run BERT model accuracy test on 2x2 TPU.""" """Run BERT model accuracy test on 2x2 TPU."""
self._setup() self._setup()
FLAGS.steps_per_loop = 50
FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu_mrpc') FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu_mrpc')
summary_path = os.path.join(FLAGS.model_dir, summary_path = os.path.join(FLAGS.model_dir,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment