bert_squad_benchmark.py 14.4 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."""

David Chen's avatar
David Chen committed
55
56
57
58
  def __init__(self, output_dir=None, tpu=None):
    super(BertSquadBenchmarkBase, self).__init__(output_dir=output_dir)
    self.tpu = tpu

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

66
67
68
69
  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'))
70

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

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

83
  def _get_distribution_strategy(self, use_ds=True):
84
    """Gets the distribution strategy."""
David Chen's avatar
David Chen committed
85
86
87
88
89
90
91
    if self.tpu:
      return distribution_utils.get_distribution_strategy(
          distribution_strategy='tpu', tpu_address=self.tpu)
    else:
      return distribution_utils.get_distribution_strategy(
          distribution_strategy='mirrored' if use_ds else 'off',
          num_gpus=self.num_gpus)
92

davidmochen's avatar
davidmochen committed
93
  @flagsaver.flagsaver
94
  def _train_squad(self, use_ds=True, run_eagerly=False):
95
    """Runs BERT SQuAD training."""
David Chen's avatar
David Chen committed
96
    assert tf.version.VERSION.startswith('2.')
97
    input_meta_data = self._read_input_meta_data_from_file()
98
    strategy = self._get_distribution_strategy(use_ds)
davidmochen's avatar
davidmochen committed
99
100
101
102

    run_squad.train_squad(
        strategy=strategy,
        input_meta_data=input_meta_data,
103
        run_eagerly=run_eagerly,
davidmochen's avatar
davidmochen committed
104
        custom_callbacks=[self.timer_callback])
105
106

  @flagsaver.flagsaver
107
  def _evaluate_squad(self, use_ds=True):
108
    """Runs BERT SQuAD evaluation."""
David Chen's avatar
David Chen committed
109
    assert tf.version.VERSION.startswith('2.')
110
    input_meta_data = self._read_input_meta_data_from_file()
111
    strategy = self._get_distribution_strategy(use_ds)
112

113
    run_squad.predict_squad(strategy=strategy, input_meta_data=input_meta_data)
114
115
116
117
118

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

    eval_metrics = squad_evaluate_v1_1.evaluate(dataset, predictions)
119
120
    # Use F1 score as reported evaluation metric.
    self.eval_metrics = eval_metrics['f1']
davidmochen's avatar
davidmochen committed
121
122


123
class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
davidmochen's avatar
davidmochen committed
124
125
126
127
  """Short benchmark performance tests for BERT SQuAD model.

  Tests BERT SQuAD performance in different GPU configurations.
  The naming convention of below test cases follow
David Chen's avatar
David Chen committed
128
129
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
davidmochen's avatar
davidmochen committed
130
131
  """

David Chen's avatar
David Chen committed
132
133
  def __init__(self, output_dir=TMP_DIR, tpu=None, **kwargs):
    super(BertSquadBenchmarkReal, self).__init__(output_dir=output_dir, tpu=tpu)
davidmochen's avatar
davidmochen committed
134
135

  def _setup(self):
136
137
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadBenchmarkReal, self)._setup()
davidmochen's avatar
davidmochen committed
138
139
140
    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
141
    FLAGS.input_meta_data_path = SQUAD_MEDIUM_INPUT_META_DATA_PATH
davidmochen's avatar
davidmochen committed
142
143
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
144
    FLAGS.steps_per_loop = 1
davidmochen's avatar
davidmochen committed
145

146
  @benchmark_wrappers.enable_runtime_flags
147
148
149
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
150
    """Runs the benchmark and reports various metrics."""
151
    start_time_sec = time.time()
152
    self._train_squad(use_ds=use_ds, run_eagerly=run_eagerly)
153
154
155
    wall_time_sec = time.time() - start_time_sec

    summary = self._read_training_summary_from_file()
David Chen's avatar
David Chen committed
156
    summary['start_time_sec'] = start_time_sec
157
158
159
160
161
162

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

  def benchmark_1_gpu(self):
165
    """Tests BERT SQuAD model performance with 1 GPU."""
davidmochen's avatar
davidmochen committed
166
167
168
169

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

172
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
173

174
175
176
177
178
179
  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')
180
181
    # XLA runs out of memory when running with batch size 4.
    FLAGS.train_batch_size = 3
182
    FLAGS.enable_xla = True
183

184
    self._run_and_report_benchmark()
185
186
187
188
189
190
191

  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')
192
    FLAGS.train_batch_size = 3
193
194
195
196
197
198
199
200
201
202

    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')
203
    FLAGS.train_batch_size = 3
204
205
206

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

davidmochen's avatar
davidmochen committed
207
  def benchmark_2_gpu(self):
208
    """Tests BERT SQuAD model performance with 2 GPUs."""
davidmochen's avatar
davidmochen committed
209
210
211
212

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

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

  def benchmark_4_gpu(self):
218
    """Tests BERT SQuAD model performance with 4 GPUs."""
davidmochen's avatar
davidmochen committed
219
220
221
222

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

225
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
226
227

  def benchmark_8_gpu(self):
228
229
230
231
232
    """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')
233
    FLAGS.train_batch_size = 24
234

235
    self._run_and_report_benchmark()
236

237
238
239
240
241
242
243
244
245
246
247
248
  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()

249
250
251
252
253
254
255
256
257
258
259
260
261
  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()

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
289
290
291
292
293
294
295
296
297
  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()

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
324
325
326
327
328
329
330
331
332
  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()
333

David Chen's avatar
David Chen committed
334
335
336
337
338
339
340
341
342
  def benchmark_2x2_tpu(self):
    """Tests BERT SQuAD model performance with 2x2 TPU."""

    self._setup()
    FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu')
    FLAGS.train_batch_size = 48

    self._run_and_report_benchmark()

343
344
345
346
347

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

  Tests BERT SQuAD accuracy. The naming convention of below test cases follow
David Chen's avatar
David Chen committed
348
349
  `benchmark_(number of gpus)_gpu` format for GPUs and
  `benchmark_(topology)_tpu` format for TPUs.
350
351
  """

David Chen's avatar
David Chen committed
352
353
  def __init__(self, output_dir=None, tpu=None, **kwargs):
    super(BertSquadAccuracy, self).__init__(output_dir=output_dir, tpu=tpu)
354
355
356
357
358
359
360
361
362
363
364

  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
365
    FLAGS.steps_per_loop = 1
366

367
  @benchmark_wrappers.enable_runtime_flags
368
369
370
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
371
    """Runs the benchmark and reports various metrics."""
372
    start_time_sec = time.time()
373
    self._train_squad(use_ds=use_ds, run_eagerly=run_eagerly)
374
375
376
377
378
379
380
381
382
    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,
383
        min_accuracy=0.900,
384
        max_accuracy=0.920)
385

386
387
388
389
390
391
392
393
394
395
  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)

396
397
  def benchmark_8_gpu(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""
davidmochen's avatar
davidmochen committed
398
399
400
401

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

404
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
405

406
407
408
409
410
411
412
413
414
415
416
417
  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()

418
419
420
421
422
423
424
  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
425
    FLAGS.enable_xla = True
426

427
    self._run_and_report_benchmark()
428

David Chen's avatar
David Chen committed
429
430
431
432
433
434
435
436
437
  def benchmark_2x2_tpu(self):
    """Tests BERT SQuAD model accuracy with 2x2 TPU."""

    self._setup()
    FLAGS.model_dir = self._get_model_dir('benchmark_2x2_tpu')
    FLAGS.train_batch_size = 48

    self._run_and_report_benchmark()

davidmochen's avatar
davidmochen committed
438
439
440

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