retinanet_benchmark.py 13.1 KB
Newer Older
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
28
29
# 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 RetinaNet benchmarks and accuracy tests."""

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

# pylint: disable=g-bad-import-order
import json
import time

from absl import flags
from absl.testing import flagsaver
import tensorflow as tf
# pylint: enable=g-bad-import-order

30
from official.benchmark import benchmark_wrappers
Jose Baiocchi's avatar
Jose Baiocchi committed
31
32
33
from official.benchmark import perfzero_benchmark
from official.utils.flags import core as flags_core
from official.utils.misc import keras_utils
34
from official.vision.detection import main as detection
Jose Baiocchi's avatar
Jose Baiocchi committed
35
from official.vision.detection.configs import base_config
36
37
38
39
40
41
42
43
44
45
46

FLAGS = flags.FLAGS

# pylint: disable=line-too-long
COCO_TRAIN_DATA = 'gs://tf-perfzero-data/coco/train*'
COCO_EVAL_DATA = 'gs://tf-perfzero-data/coco/val*'
COCO_EVAL_JSON = 'gs://tf-perfzero-data/coco/instances_val2017.json'
RESNET_CHECKPOINT_PATH = 'gs://cloud-tpu-checkpoints/retinanet/resnet50-checkpoint-2018-02-07'
# pylint: enable=line-too-long


Allen Wang's avatar
Allen Wang committed
47
class BenchmarkBase(perfzero_benchmark.PerfZeroBenchmark):
Jose Baiocchi's avatar
Jose Baiocchi committed
48
49
50
  """Base class to hold methods common to test classes."""

  def __init__(self, **kwargs):
Allen Wang's avatar
Allen Wang committed
51
    super(BenchmarkBase, self).__init__(**kwargs)
Jose Baiocchi's avatar
Jose Baiocchi committed
52
53
54
55
    self.timer_callback = None

  def _report_benchmark(self, stats, start_time_sec, wall_time_sec, min_ap,
                        max_ap, warmup):
56
57
58
59
    """Report benchmark results by writing to local protobuf file.

    Args:
      stats: dict returned from Detection models with known entries.
Jose Baiocchi's avatar
Jose Baiocchi committed
60
61
      start_time_sec: the start of the benchmark execution in seconds
      wall_time_sec: the duration of the benchmark execution in seconds
62
63
64
65
      min_ap: Minimum detection AP constraint to verify correctness of the
        model.
      max_ap: Maximum detection AP accuracy constraint to verify correctness of
        the model.
Jose Baiocchi's avatar
Jose Baiocchi committed
66
      warmup: Number of time log entries to ignore when computing examples/sec.
67
68
69
70
71
72
73
    """
    metrics = [{
        'name': 'total_loss',
        'value': stats['total_loss'],
    }]
    if self.timer_callback:
      metrics.append({
74
          'name': 'exp_per_second',
Jose Baiocchi's avatar
Jose Baiocchi committed
75
76
77
78
79
          'value': self.timer_callback.get_examples_per_sec(warmup)
      })
      metrics.append({
          'name': 'startup_time',
          'value': self.timer_callback.get_startup_time(start_time_sec)
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
      })
    else:
      metrics.append({
          'name': 'exp_per_second',
          'value': 0.0,
      })

    if 'eval_metrics' in stats:
      metrics.append({
          'name': 'AP',
          'value': stats['AP'],
          'min_value': min_ap,
          'max_value': max_ap,
      })
    flags_str = flags_core.get_nondefault_flags_as_str()
    self.report_benchmark(
        iters=stats['total_steps'],
        wall_time=wall_time_sec,
        metrics=metrics,
        extras={'flags': flags_str})


Allen Wang's avatar
Allen Wang committed
102
class DetectionBenchmarkBase(BenchmarkBase):
103
104
  """Base class to hold methods common to test classes in the module."""

Jose Baiocchi's avatar
Jose Baiocchi committed
105
  def __init__(self, **kwargs):
106
107
108
109
    self.train_data_path = COCO_TRAIN_DATA
    self.eval_data_path = COCO_EVAL_DATA
    self.eval_json_path = COCO_EVAL_JSON
    self.resnet_checkpoint_path = RESNET_CHECKPOINT_PATH
Allen Wang's avatar
Allen Wang committed
110
    super(DetectionBenchmarkBase, self).__init__(**kwargs)
111
112
113

  def _run_detection_main(self):
    """Starts detection job."""
Yeqing Li's avatar
Yeqing Li committed
114
    if self.timer_callback:
Jose Baiocchi's avatar
Jose Baiocchi committed
115
      FLAGS.log_steps = 0  # prevent detection.run from adding the same callback
Yeqing Li's avatar
Yeqing Li committed
116
117
118
      return detection.run(callbacks=[self.timer_callback])
    else:
      return detection.run()
119
120


Allen Wang's avatar
Allen Wang committed
121
class DetectionAccuracy(DetectionBenchmarkBase):
122
123
124
125
126
127
128
  """Accuracy test for RetinaNet model.

  Tests RetinaNet detection task model accuracy. The naming
  convention of below test cases follow
  `benchmark_(number of gpus)_gpu_(dataset type)` format.
  """

129
  def __init__(self, model, per_gpu_batch_size=8, **kwargs):
Allen Wang's avatar
Allen Wang committed
130
    self.model = model
131
    self.per_gpu_batch_size = per_gpu_batch_size
Allen Wang's avatar
Allen Wang committed
132
133
    super(DetectionAccuracy, self).__init__(**kwargs)

134
  @benchmark_wrappers.enable_runtime_flags
Jose Baiocchi's avatar
Jose Baiocchi committed
135
136
137
138
139
140
  def _run_and_report_benchmark(self,
                                params,
                                min_ap=0.325,
                                max_ap=0.35,
                                do_eval=True,
                                warmup=1):
Allen Wang's avatar
Allen Wang committed
141
    """Starts Detection accuracy benchmark test."""
Jose Baiocchi's avatar
Jose Baiocchi committed
142
143
    FLAGS.params_override = json.dumps(params)
    # Need timer callback to measure performance
Hongkun Yu's avatar
Hongkun Yu committed
144
    self.timer_callback = keras_utils.TimeHistory(
Jose Baiocchi's avatar
Jose Baiocchi committed
145
146
147
        batch_size=params['train']['batch_size'],
        log_steps=FLAGS.log_steps,
    )
148
149
150
151
152
153

    start_time_sec = time.time()
    FLAGS.mode = 'train'
    summary, _ = self._run_detection_main()
    wall_time_sec = time.time() - start_time_sec

Jose Baiocchi's avatar
Jose Baiocchi committed
154
155
156
157
    if do_eval:
      FLAGS.mode = 'eval'
      eval_metrics = self._run_detection_main()
      summary.update(eval_metrics)
158

Jose Baiocchi's avatar
Jose Baiocchi committed
159
160
161
    summary['total_steps'] = params['train']['total_steps']
    self._report_benchmark(summary, start_time_sec, wall_time_sec, min_ap,
                           max_ap, warmup)
162
163

  def _setup(self):
Allen Wang's avatar
Allen Wang committed
164
165
    super(DetectionAccuracy, self)._setup()
    FLAGS.model = self.model
166

Jose Baiocchi's avatar
Jose Baiocchi committed
167
168
  def _params(self):
    return {
Pengchong Jin's avatar
Pengchong Jin committed
169
170
171
        'architecture': {
            'use_bfloat16': True,
        },
172
173
174
175
176
        'train': {
            'batch_size': 64,
            'iterations_per_loop': 100,
            'total_steps': 22500,
            'train_file_pattern': self.train_data_path,
Yeqing Li's avatar
Yeqing Li committed
177
178
179
180
            'checkpoint': {
                'path': self.resnet_checkpoint_path,
                'prefix': 'resnet50/'
            },
Jose Baiocchi's avatar
Jose Baiocchi committed
181
182
            # Speed up ResNet training when loading from the checkpoint.
            'frozen_variable_prefix': base_config.RESNET_FROZEN_VAR_PREFIX,
183
184
185
186
187
188
189
190
191
192
193
194
195
        },
        'eval': {
            'batch_size': 8,
            'eval_samples': 5000,
            'val_json_file': self.eval_json_path,
            'eval_file_pattern': self.eval_data_path,
        },
    }

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
    """Run RetinaNet model accuracy test with 8 GPUs."""
    self._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
196
197
    params = self._params()
    FLAGS.num_gpus = 8
198
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_coco')
Jose Baiocchi's avatar
Jose Baiocchi committed
199
200
    FLAGS.strategy_type = 'mirrored'
    self._run_and_report_benchmark(params)
201
202


Allen Wang's avatar
Allen Wang committed
203
204
class DetectionBenchmarkReal(DetectionAccuracy):
  """Short benchmark performance tests for a detection model.
205

Allen Wang's avatar
Allen Wang committed
206
  Tests detection performance in different accelerator configurations.
207
208
209
210
  The naming convention of below test cases follow
  `benchmark_(number of gpus)_gpu` format.
  """

Jose Baiocchi's avatar
Jose Baiocchi committed
211
  def _setup(self):
Allen Wang's avatar
Allen Wang committed
212
    super(DetectionBenchmarkReal, self)._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
213
214
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
215
216
217

  @flagsaver.flagsaver
  def benchmark_8_gpu_coco(self):
Allen Wang's avatar
Allen Wang committed
218
    """Run detection model accuracy test with 8 GPUs."""
219
    self._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
220
    params = self._params()
Pengchong Jin's avatar
Pengchong Jin committed
221
    params['architecture']['use_bfloat16'] = False
David Chen's avatar
David Chen committed
222
    params['train']['total_steps'] = 1875  # One epoch.
223
    params['train']['batch_size'] = 8 * self.per_gpu_batch_size
Yeqing Li's avatar
Yeqing Li committed
224
225
226
227
228
229
230
    # The iterations_per_loop must be one, otherwise the number of examples per
    # second would be wrong. Currently only support calling callback per batch
    # when each loop only runs on one batch, i.e. host loop for one step. The
    # performance of this situation might be lower than the case of
    # iterations_per_loop > 1.
    # Related bug: b/135933080
    params['train']['iterations_per_loop'] = 1
231
    params['eval']['eval_samples'] = 8
Jose Baiocchi's avatar
Jose Baiocchi committed
232
    FLAGS.num_gpus = 8
233
    FLAGS.model_dir = self._get_model_dir('real_benchmark_8_gpu_coco')
Jose Baiocchi's avatar
Jose Baiocchi committed
234
235
    FLAGS.strategy_type = 'mirrored'
    self._run_and_report_benchmark(params)
236

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
237
238
  @flagsaver.flagsaver
  def benchmark_1_gpu_coco(self):
Allen Wang's avatar
Allen Wang committed
239
    """Run detection model accuracy test with 1 GPU."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
240
    self._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
241
    params = self._params()
Pengchong Jin's avatar
Pengchong Jin committed
242
    params['architecture']['use_bfloat16'] = False
243
    params['train']['batch_size'] = 1 * self.per_gpu_batch_size
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
244
245
246
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
Jose Baiocchi's avatar
Jose Baiocchi committed
247
    FLAGS.num_gpus = 1
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
248
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
Yeqing Li's avatar
Yeqing Li committed
249
    FLAGS.strategy_type = 'one_device'
Jose Baiocchi's avatar
Jose Baiocchi committed
250
    self._run_and_report_benchmark(params)
251

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
252
253
  @flagsaver.flagsaver
  def benchmark_xla_1_gpu_coco(self):
Allen Wang's avatar
Allen Wang committed
254
    """Run detection model accuracy test with 1 GPU and XLA enabled."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
255
    self._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
256
    params = self._params()
Pengchong Jin's avatar
Pengchong Jin committed
257
    params['architecture']['use_bfloat16'] = False
258
    params['train']['batch_size'] = 1 * self.per_gpu_batch_size
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
259
260
261
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
Jose Baiocchi's avatar
Jose Baiocchi committed
262
263
    FLAGS.num_gpus = 1
    FLAGS.model_dir = self._get_model_dir('real_benchmark_xla_1_gpu_coco')
Yeqing Li's avatar
Yeqing Li committed
264
    FLAGS.strategy_type = 'one_device'
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
265
    FLAGS.enable_xla = True
Jose Baiocchi's avatar
Jose Baiocchi committed
266
267
268
269
    self._run_and_report_benchmark(params)

  @flagsaver.flagsaver
  def benchmark_2x2_tpu_coco(self):
Allen Wang's avatar
Allen Wang committed
270
    """Run detection model accuracy test with 4 TPUs."""
Jose Baiocchi's avatar
Jose Baiocchi committed
271
272
273
274
275
276
277
278
279
    self._setup()
    params = self._params()
    params['train']['batch_size'] = 64
    params['train']['total_steps'] = 1875  # One epoch.
    params['train']['iterations_per_loop'] = 500
    FLAGS.model_dir = self._get_model_dir('real_benchmark_2x2_tpu_coco')
    FLAGS.strategy_type = 'tpu'
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

Allen Wang's avatar
Allen Wang committed
280
281
  @flagsaver.flagsaver
  def benchmark_4x4_tpu_coco(self):
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
282
    """Run detection model accuracy test with 4x4 TPU."""
Allen Wang's avatar
Allen Wang committed
283
284
285
    self._setup()
    params = self._params()
    params['train']['batch_size'] = 256
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
286
    params['train']['total_steps'] = 10 * 469  # 10 epochs.
Allen Wang's avatar
Allen Wang committed
287
288
289
290
291
292
293
    params['train']['iterations_per_loop'] = 500
    FLAGS.model_dir = self._get_model_dir('real_benchmark_4x4_tpu_coco')
    FLAGS.strategy_type = 'tpu'
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

  @flagsaver.flagsaver
  def benchmark_2x2_tpu_coco_mlir(self):
Allen Wang's avatar
Allen Wang committed
294
    """Run detection model accuracy test with 4 TPUs."""
Allen Wang's avatar
Allen Wang committed
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
    self._setup()
    params = self._params()
    params['train']['batch_size'] = 64
    params['train']['total_steps'] = 1875  # One epoch.
    params['train']['iterations_per_loop'] = 500
    FLAGS.model_dir = self._get_model_dir('real_benchmark_2x2_tpu_coco_mlir')
    FLAGS.strategy_type = 'tpu'
    tf.config.experimental.enable_mlir_bridge()
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

  @flagsaver.flagsaver
  def benchmark_4x4_tpu_coco_mlir(self):
    """Run RetinaNet model accuracy test with 4 TPUs."""
    self._setup()
    params = self._params()
    params['train']['batch_size'] = 256
    params['train']['total_steps'] = 469  # One epoch.
    params['train']['iterations_per_loop'] = 500
    FLAGS.model_dir = self._get_model_dir('real_benchmark_4x4_tpu_coco_mlir')
    FLAGS.strategy_type = 'tpu'
    tf.config.experimental.enable_mlir_bridge()
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

318
319
  @flagsaver.flagsaver
  def benchmark_2x2_tpu_spinenet_coco(self):
Allen Wang's avatar
Allen Wang committed
320
    """Run detection model with SpineNet backbone accuracy test with 4 TPUs."""
321
322
323
324
325
326
327
328
329
330
331
332
333
334
    self._setup()
    params = self._params()
    params['architecture']['backbone'] = 'spinenet'
    params['architecture']['multilevel_features'] = 'identity'
    params['architecture']['use_bfloat16'] = False
    params['train']['batch_size'] = 64
    params['train']['total_steps'] = 1875  # One epoch.
    params['train']['iterations_per_loop'] = 500
    params['train']['checkpoint']['path'] = ''
    FLAGS.model_dir = self._get_model_dir(
        'real_benchmark_2x2_tpu_spinenet_coco')
    FLAGS.strategy_type = 'tpu'
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
335

Allen Wang's avatar
Allen Wang committed
336
337
338
339
class RetinanetBenchmarkReal(DetectionBenchmarkReal):
  """Short benchmark performance tests for Retinanet model."""

  def __init__(self, **kwargs):
340
341
342
    super(RetinanetBenchmarkReal, self).__init__(model='retinanet',
                                                 per_gpu_batch_size=8,
                                                 **kwargs)
Allen Wang's avatar
Allen Wang committed
343
344
345
346
347
348


class MaskRCNNBenchmarkReal(DetectionBenchmarkReal):
  """Short benchmark performance tests for Mask RCNN model."""

  def __init__(self, **kwargs):
349
350
351
    super(MaskRCNNBenchmarkReal, self).__init__(model='mask_rcnn',
                                                per_gpu_batch_size=4,
                                                **kwargs)
Allen Wang's avatar
Allen Wang committed
352
353
354
355
356
357


class ShapeMaskBenchmarkReal(DetectionBenchmarkReal):
  """Short benchmark performance tests for ShapeMask model."""

  def __init__(self, **kwargs):
358
    super(ShapeMaskBenchmarkReal, self).__init__(model='shapemask',
Hongkun Yu's avatar
Hongkun Yu committed
359
                                                 per_gpu_batch_size=4,
360
                                                 **kwargs)
Allen Wang's avatar
Allen Wang committed
361
362


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