bert_squad_benchmark.py 13.5 KB
Newer Older
davidmochen's avatar
davidmochen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# 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.
# ==============================================================================
"""Executes BERT SQuAD benchmarks and accuracy tests."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import json
import os
import time

# pylint: disable=g-bad-import-order
from absl import flags
from absl.testing import flagsaver
28
import tensorflow as tf
davidmochen's avatar
davidmochen committed
29
30
# pylint: enable=g-bad-import-order

31
32
from official.benchmark import bert_benchmark_utils as benchmark_utils
from official.benchmark import squad_evaluate_v1_1
33
from official.nlp.bert import run_squad
davidmochen's avatar
davidmochen committed
34
from official.utils.misc import distribution_utils
35
36
from official.utils.testing import benchmark_wrappers

davidmochen's avatar
davidmochen committed
37
38

# pylint: disable=line-too-long
David Chen's avatar
David Chen committed
39
PRETRAINED_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16/bert_model.ckpt'
davidmochen's avatar
davidmochen committed
40
41
SQUAD_TRAIN_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_train.tf_record'
SQUAD_PREDICT_FILE = 'gs://tf-perfzero-data/bert/squad/dev-v1.1.json'
David Chen's avatar
David Chen committed
42
SQUAD_VOCAB_FILE = 'gs://tf-perfzero-data/bert/squad/vocab.txt'
David Chen's avatar
David Chen committed
43
SQUAD_MEDIUM_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_medium_meta_data'
44
SQUAD_FULL_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_full_meta_data'
David Chen's avatar
David Chen committed
45
MODEL_CONFIG_FILE_PATH = 'gs://cloud-tpu-checkpoints/bert/keras_bert/uncased_L-24_H-1024_A-16/bert_config.json'
davidmochen's avatar
davidmochen committed
46
47
# pylint: enable=line-too-long

David Chen's avatar
David Chen committed
48
TMP_DIR = os.getenv('TMPDIR')
davidmochen's avatar
davidmochen committed
49
50
51
52
53
54
FLAGS = flags.FLAGS


class BertSquadBenchmarkBase(benchmark_utils.BertBenchmarkBase):
  """Base class to hold methods common to test classes in the module."""

55
56
  def _read_training_summary_from_file(self):
    """Reads the training summary from a file."""
57
58
    summary_path = os.path.join(FLAGS.model_dir,
                                'summaries/training_summary.txt')
59
60
    with tf.io.gfile.GFile(summary_path, 'rb') as reader:
      return json.loads(reader.read().decode('utf-8'))
61

62
63
64
65
  def _read_input_meta_data_from_file(self):
    """Reads the input metadata from a file."""
    with tf.io.gfile.GFile(FLAGS.input_meta_data_path, 'rb') as reader:
      return json.loads(reader.read().decode('utf-8'))
66

67
68
  def _read_predictions_dataset_from_file(self):
    """Reads the predictions dataset from a file."""
69
70
    with tf.io.gfile.GFile(SQUAD_PREDICT_FILE, 'r') as reader:
      dataset_json = json.load(reader)
71
      return dataset_json['data']
72

73
74
75
  def _read_predictions_from_file(self):
    """Reads the predictions from a file."""
    predictions_file = os.path.join(FLAGS.model_dir, 'predictions.json')
76
    with tf.io.gfile.GFile(predictions_file, 'r') as reader:
77
      return json.load(reader)
78

79
  def _get_distribution_strategy(self, use_ds=True):
80
81
    """Gets the distribution strategy."""
    return distribution_utils.get_distribution_strategy(
82
83
        distribution_strategy='mirrored' if use_ds else 'off',
        num_gpus=self.num_gpus)
84

davidmochen's avatar
davidmochen committed
85
  @flagsaver.flagsaver
86
  def _train_squad(self, use_ds=True, run_eagerly=False):
87
    """Runs BERT SQuAD training."""
David Chen's avatar
David Chen committed
88
    assert tf.version.VERSION.startswith('2.')
89
    input_meta_data = self._read_input_meta_data_from_file()
90
    strategy = self._get_distribution_strategy(use_ds)
davidmochen's avatar
davidmochen committed
91
92
93
94

    run_squad.train_squad(
        strategy=strategy,
        input_meta_data=input_meta_data,
95
        run_eagerly=run_eagerly,
davidmochen's avatar
davidmochen committed
96
        custom_callbacks=[self.timer_callback])
97
98

  @flagsaver.flagsaver
99
  def _evaluate_squad(self, use_ds=True):
100
    """Runs BERT SQuAD evaluation."""
David Chen's avatar
David Chen committed
101
    assert tf.version.VERSION.startswith('2.')
102
    input_meta_data = self._read_input_meta_data_from_file()
103
    strategy = self._get_distribution_strategy(use_ds)
104

105
    run_squad.predict_squad(strategy=strategy, input_meta_data=input_meta_data)
106
107
108
109
110

    dataset = self._read_predictions_dataset_from_file()
    predictions = self._read_predictions_from_file()

    eval_metrics = squad_evaluate_v1_1.evaluate(dataset, predictions)
111
112
    # Use F1 score as reported evaluation metric.
    self.eval_metrics = eval_metrics['f1']
davidmochen's avatar
davidmochen committed
113
114


115
class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
davidmochen's avatar
davidmochen committed
116
117
118
119
120
121
122
  """Short benchmark performance tests for BERT SQuAD model.

  Tests BERT SQuAD performance in different GPU configurations.
  The naming convention of below test cases follow
  `benchmark_(number of gpus)_gpu` format.
  """

David Chen's avatar
David Chen committed
123
  def __init__(self, output_dir=TMP_DIR, **kwargs):
124
    super(BertSquadBenchmarkReal, self).__init__(output_dir=output_dir)
davidmochen's avatar
davidmochen committed
125
126

  def _setup(self):
127
128
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadBenchmarkReal, self)._setup()
davidmochen's avatar
davidmochen committed
129
130
131
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
David Chen's avatar
David Chen committed
132
    FLAGS.input_meta_data_path = SQUAD_MEDIUM_INPUT_META_DATA_PATH
davidmochen's avatar
davidmochen committed
133
134
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
135
    FLAGS.steps_per_loop = 1
davidmochen's avatar
davidmochen committed
136

137
  @benchmark_wrappers.enable_runtime_flags
138
139
140
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
141
    """Runs the benchmark and reports various metrics."""
142
    start_time_sec = time.time()
143
    self._train_squad(use_ds=use_ds, run_eagerly=run_eagerly)
144
145
146
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
David Chen's avatar
David Chen committed
147
    summary['start_time_sec'] = start_time_sec
148
149
150
151
152
153

    super(BertSquadBenchmarkReal, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
        min_accuracy=0,
        max_accuracy=1)
davidmochen's avatar
davidmochen committed
154
155

  def benchmark_1_gpu(self):
156
    """Tests BERT SQuAD model performance with 1 GPU."""
davidmochen's avatar
davidmochen committed
157
158
159
160

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad')
161
    FLAGS.train_batch_size = 3
davidmochen's avatar
davidmochen committed
162

163
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
164

165
166
167
168
169
170
  def benchmark_1_gpu_xla(self):
    """Tests BERT SQuAD model performance with 1 GPU with XLA."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_xla_squad')
171
172
    # XLA runs out of memory when running with batch size 4.
    FLAGS.train_batch_size = 3
173
    FLAGS.enable_xla = True
174

175
    self._run_and_report_benchmark()
176
177
178
179
180
181
182

  def benchmark_1_gpu_no_dist_strat(self):
    """Tests BERT SQuAD model performance with 1 GPU without DS."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_no_dist_strat_squad')
183
    FLAGS.train_batch_size = 3
184
185
186
187
188
189
190
191
192
193

    self._run_and_report_benchmark(use_ds=False)

  def benchmark_1_gpu_eager_no_dist_strat(self):
    """Tests BERT SQuAD model performance with 1 GPU with eager execution."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir(
        'benchmark_1_gpu_eager_no_dist_strat_squad')
194
    FLAGS.train_batch_size = 3
195
196
197

    self._run_and_report_benchmark(use_ds=False, run_eagerly=True)

davidmochen's avatar
davidmochen committed
198
  def benchmark_2_gpu(self):
199
    """Tests BERT SQuAD model performance with 2 GPUs."""
davidmochen's avatar
davidmochen committed
200
201
202
203

    self._setup()
    self.num_gpus = 2
    FLAGS.model_dir = self._get_model_dir('benchmark_2_gpu_squad')
204
    FLAGS.train_batch_size = 6
davidmochen's avatar
davidmochen committed
205

206
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
207
208

  def benchmark_4_gpu(self):
209
    """Tests BERT SQuAD model performance with 4 GPUs."""
davidmochen's avatar
davidmochen committed
210
211
212
213

    self._setup()
    self.num_gpus = 4
    FLAGS.model_dir = self._get_model_dir('benchmark_4_gpu_squad')
214
    FLAGS.train_batch_size = 12
davidmochen's avatar
davidmochen committed
215

216
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
217
218

  def benchmark_8_gpu(self):
219
220
221
222
223
    """Tests BERT SQuAD model performance with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad')
224
    FLAGS.train_batch_size = 24
225

226
    self._run_and_report_benchmark()
227

228
229
230
231
232
233
234
235
236
237
238
239
  def benchmark_1_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 1 GPU and FP16."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_fp16')
    FLAGS.train_batch_size = 4
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

240
241
242
243
244
245
246
247
248
249
250
251
252
  def benchmark_1_gpu_xla_fp16(self):
    """Tests BERT SQuAD model performance with 1 GPU with XLA and FP16."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_xla_squad_fp16')
    FLAGS.train_batch_size = 4
    FLAGS.enable_xla = True
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
  def benchmark_2_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 2 GPUs and FP16."""

    self._setup()
    self.num_gpus = 2
    FLAGS.model_dir = self._get_model_dir('benchmark_2_gpu_squad_fp16')
    FLAGS.train_batch_size = 8
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

  def benchmark_4_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 4 GPUs and FP16."""

    self._setup()
    self.num_gpus = 4
    FLAGS.model_dir = self._get_model_dir('benchmark_4_gpu_squad_fp16')
    FLAGS.train_batch_size = 16
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

  def benchmark_8_gpu_fp16(self):
    """Tests BERT SQuAD model performance with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_fp16')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
  def benchmark_1_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_amp_squad')
    FLAGS.train_batch_size = 4
    FLAGS.dtype = 'fp16'
    FLAGS.fp16_implementation = 'graph_rewrite'

    self._run_and_report_benchmark()

  def benchmark_4_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 4
    FLAGS.model_dir = self._get_model_dir('benchmark_4_gpu_amp_squad')
    FLAGS.train_batch_size = 16
    FLAGS.dtype = 'fp16'
    FLAGS.fp16_implementation = 'graph_rewrite'

    self._run_and_report_benchmark()

  def benchmark_8_gpu_amp(self):
    """Tests BERT SQuAD model performance with 1 GPU with automatic mixed precision."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_amp_squad')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.fp16_implementation = 'graph_rewrite'

    self._run_and_report_benchmark()
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345


class BertSquadAccuracy(BertSquadBenchmarkBase):
  """Short accuracy test for BERT SQuAD model.

  Tests BERT SQuAD accuracy. The naming convention of below test cases follow
  `benchmark_(number of gpus)_gpu` format.
  """

  def __init__(self, output_dir=None, **kwargs):
    super(BertSquadAccuracy, self).__init__(output_dir=output_dir)

  def _setup(self):
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadAccuracy, self)._setup()
    FLAGS.train_data_path = SQUAD_TRAIN_DATA_PATH
    FLAGS.predict_file = SQUAD_PREDICT_FILE
    FLAGS.vocab_file = SQUAD_VOCAB_FILE
    FLAGS.input_meta_data_path = SQUAD_FULL_INPUT_META_DATA_PATH
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.init_checkpoint = PRETRAINED_CHECKPOINT_PATH
    FLAGS.num_train_epochs = 2
346
    FLAGS.steps_per_loop = 1
347

348
  @benchmark_wrappers.enable_runtime_flags
349
350
351
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
352
    """Runs the benchmark and reports various metrics."""
353
    start_time_sec = time.time()
354
    self._train_squad(use_ds=use_ds, run_eagerly=run_eagerly)
355
356
357
358
359
360
361
362
363
    self._evaluate_squad()
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
    summary['eval_metrics'] = self.eval_metrics

    super(BertSquadAccuracy, self)._report_benchmark(
        stats=summary,
        wall_time_sec=wall_time_sec,
364
        min_accuracy=0.900,
365
        max_accuracy=0.920)
366

367
368
369
370
371
372
373
374
375
376
  def benchmark_1_gpu_eager(self):
    """Tests BERT SQuAD model accuracy with 1 GPU with eager execution."""

    self._setup()
    self.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu_squad_eager')
    FLAGS.train_batch_size = 4

    self._run_and_report_benchmark(use_ds=False, run_eagerly=True)

377
378
  def benchmark_8_gpu(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""
davidmochen's avatar
davidmochen committed
379
380
381
382

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad')
383
    FLAGS.train_batch_size = 24
davidmochen's avatar
davidmochen committed
384

385
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
386

387
388
389
390
391
392
393
394
395
396
397
398
  def benchmark_8_gpu_fp16(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs and FP16."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_fp16')
    FLAGS.train_batch_size = 32
    FLAGS.dtype = 'fp16'
    FLAGS.loss_scale = 'dynamic'

    self._run_and_report_benchmark()

399
400
401
402
403
404
405
  def benchmark_8_gpu_xla(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""

    self._setup()
    self.num_gpus = 8
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_squad_xla')
    FLAGS.train_batch_size = 32
406
    FLAGS.enable_xla = True
407

408
    self._run_and_report_benchmark()
409

davidmochen's avatar
davidmochen committed
410
411
412

if __name__ == '__main__':
  tf.test.main()