bert_squad_benchmark.py 13.2 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
35
36
from official.utils.misc import distribution_utils

# pylint: disable=line-too-long
37
PRETRAINED_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/bert/tf_20/uncased_L-24_H-1024_A-16/bert_model.ckpt'
davidmochen's avatar
davidmochen committed
38
39
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'
SQUAD_VOCAB_FILE = 'gs://tf-perfzero-data/bert/squad/vocab.txt'
SQUAD_SMALL_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_small_meta_data'
42
SQUAD_FULL_INPUT_META_DATA_PATH = 'gs://tf-perfzero-data/bert/squad/squad_full_meta_data'
davidmochen's avatar
davidmochen committed
43
44
45
MODEL_CONFIG_FILE_PATH = 'gs://cloud-tpu-checkpoints/bert/tf_20/uncased_L-24_H-1024_A-16/bert_config'
# pylint: enable=line-too-long

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


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

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

59
60
61
62
  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'))
63

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

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

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

davidmochen's avatar
davidmochen committed
82
  @flagsaver.flagsaver
83
  def _train_squad(self, use_ds=True, run_eagerly=False):
84
    """Runs BERT SQuAD training."""
85
    tf.enable_v2_behavior()
86
    input_meta_data = self._read_input_meta_data_from_file()
87
    strategy = self._get_distribution_strategy(use_ds)
davidmochen's avatar
davidmochen committed
88
89
90
91

    run_squad.train_squad(
        strategy=strategy,
        input_meta_data=input_meta_data,
92
        run_eagerly=run_eagerly,
davidmochen's avatar
davidmochen committed
93
        custom_callbacks=[self.timer_callback])
94
95

  @flagsaver.flagsaver
96
  def _evaluate_squad(self, use_ds=True):
97
    """Runs BERT SQuAD evaluation."""
98
    tf.enable_v2_behavior()
99
    input_meta_data = self._read_input_meta_data_from_file()
100
    strategy = self._get_distribution_strategy(use_ds)
101

102
    run_squad.predict_squad(strategy=strategy, input_meta_data=input_meta_data)
103
104
105
106
107

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

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


112
class BertSquadBenchmarkReal(BertSquadBenchmarkBase):
davidmochen's avatar
davidmochen committed
113
114
115
116
117
118
119
  """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
120
  def __init__(self, output_dir=TMP_DIR, **kwargs):
121
    super(BertSquadBenchmarkReal, self).__init__(output_dir=output_dir)
davidmochen's avatar
davidmochen committed
122
123

  def _setup(self):
124
125
    """Sets up the benchmark and SQuAD flags."""
    super(BertSquadBenchmarkReal, self)._setup()
davidmochen's avatar
davidmochen committed
126
127
128
129
130
131
    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_SMALL_INPUT_META_DATA_PATH
    FLAGS.bert_config_file = MODEL_CONFIG_FILE_PATH
    FLAGS.num_train_epochs = 1
132
    FLAGS.steps_per_loop = 1
davidmochen's avatar
davidmochen committed
133

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

    summary = self._read_training_summary_from_file()

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

  def benchmark_1_gpu(self):
151
    """Tests BERT SQuAD model performance with 1 GPU."""
davidmochen's avatar
davidmochen committed
152
153
154
155

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

158
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
159

160
161
162
163
164
165
  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')
166
167
    # XLA runs out of memory when running with batch size 4.
    FLAGS.train_batch_size = 3
168
    FLAGS.enable_xla = True
169

170
    self._run_and_report_benchmark()
171
172
173
174
175
176
177

  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')
178
    FLAGS.train_batch_size = 3
179
180
181
182
183
184
185
186
187
188

    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')
189
    FLAGS.train_batch_size = 3
190
191
192

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

davidmochen's avatar
davidmochen committed
193
  def benchmark_2_gpu(self):
194
    """Tests BERT SQuAD model performance with 2 GPUs."""
davidmochen's avatar
davidmochen committed
195
196
197
198

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

201
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
202
203

  def benchmark_4_gpu(self):
204
    """Tests BERT SQuAD model performance with 4 GPUs."""
davidmochen's avatar
davidmochen committed
205
206
207
208

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

211
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
212
213

  def benchmark_8_gpu(self):
214
215
216
217
218
    """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')
219
    FLAGS.train_batch_size = 24
220

221
    self._run_and_report_benchmark()
222

223
224
225
226
227
228
229
230
231
232
233
234
  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()

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

248
249
250
251
252
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
  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()

284
285
286
287
288
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
  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()
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340


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
341
    FLAGS.steps_per_loop = 1
342

343
344
345
  def _run_and_report_benchmark(self,
                                use_ds=True,
                                run_eagerly=False):
346
    """Runs the benchmark and reports various metrics."""
347
    start_time_sec = time.time()
348
    self._train_squad(use_ds=use_ds, run_eagerly=run_eagerly)
349
350
351
352
353
354
355
356
357
    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,
358
        min_accuracy=0.900,
359
        max_accuracy=0.920)
360

361
362
363
364
365
366
367
368
369
370
  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)

371
372
  def benchmark_8_gpu(self):
    """Tests BERT SQuAD model accuracy with 8 GPUs."""
davidmochen's avatar
davidmochen committed
373
374
375
376

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

379
    self._run_and_report_benchmark()
davidmochen's avatar
davidmochen committed
380

381
382
383
384
385
386
387
388
389
390
391
392
  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()

393
394
395
396
397
398
399
  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
400
    FLAGS.enable_xla = True
401

402
    self._run_and_report_benchmark()
403

davidmochen's avatar
davidmochen committed
404
405
406

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