retinanet_benchmark.py 11.7 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


Jose Baiocchi's avatar
Jose Baiocchi committed
47
48
49
50
51
52
53
54
55
class DetectionBenchmarkBase(perfzero_benchmark.PerfZeroBenchmark):
  """Base class to hold methods common to test classes."""

  def __init__(self, **kwargs):
    super(DetectionBenchmarkBase, self).__init__(**kwargs)
    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
102
103
104
      })
    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})


class RetinanetBenchmarkBase(DetectionBenchmarkBase):
  """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
Jose Baiocchi's avatar
Jose Baiocchi committed
110
    super(RetinanetBenchmarkBase, 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
121
122
123
124
125
126
127
128


class RetinanetAccuracy(RetinanetBenchmarkBase):
  """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
  @benchmark_wrappers.enable_runtime_flags
Jose Baiocchi's avatar
Jose Baiocchi committed
130
131
132
133
134
135
  def _run_and_report_benchmark(self,
                                params,
                                min_ap=0.325,
                                max_ap=0.35,
                                do_eval=True,
                                warmup=1):
136
    """Starts RetinaNet accuracy benchmark test."""
Jose Baiocchi's avatar
Jose Baiocchi committed
137
138
    FLAGS.params_override = json.dumps(params)
    # Need timer callback to measure performance
Hongkun Yu's avatar
Hongkun Yu committed
139
    self.timer_callback = keras_utils.TimeHistory(
Jose Baiocchi's avatar
Jose Baiocchi committed
140
141
142
        batch_size=params['train']['batch_size'],
        log_steps=FLAGS.log_steps,
    )
143
144
145
146
147
148

    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
149
150
151
152
    if do_eval:
      FLAGS.mode = 'eval'
      eval_metrics = self._run_detection_main()
      summary.update(eval_metrics)
153

Jose Baiocchi's avatar
Jose Baiocchi committed
154
155
156
    summary['total_steps'] = params['train']['total_steps']
    self._report_benchmark(summary, start_time_sec, wall_time_sec, min_ap,
                           max_ap, warmup)
157
158
159
160
161

  def _setup(self):
    super(RetinanetAccuracy, self)._setup()
    FLAGS.model = 'retinanet'

Jose Baiocchi's avatar
Jose Baiocchi committed
162
163
  def _params(self):
    return {
Pengchong Jin's avatar
Pengchong Jin committed
164
165
166
        'architecture': {
            'use_bfloat16': True,
        },
167
168
169
170
171
        '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
172
173
174
175
            'checkpoint': {
                'path': self.resnet_checkpoint_path,
                'prefix': 'resnet50/'
            },
Jose Baiocchi's avatar
Jose Baiocchi committed
176
177
            # Speed up ResNet training when loading from the checkpoint.
            'frozen_variable_prefix': base_config.RESNET_FROZEN_VAR_PREFIX,
178
179
180
181
182
183
184
185
186
187
188
189
190
        },
        '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
191
192
    params = self._params()
    FLAGS.num_gpus = 8
193
    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_coco')
Jose Baiocchi's avatar
Jose Baiocchi committed
194
195
    FLAGS.strategy_type = 'mirrored'
    self._run_and_report_benchmark(params)
196
197
198
199
200
201
202
203
204
205


class RetinanetBenchmarkReal(RetinanetAccuracy):
  """Short benchmark performance tests for RetinaNet model.

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

Jose Baiocchi's avatar
Jose Baiocchi committed
206
207
208
209
  def _setup(self):
    super(RetinanetBenchmarkReal, self)._setup()
    # Use negative value to avoid saving checkpoints.
    FLAGS.save_checkpoint_freq = -1
210
211
212
213
214

  @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
215
    params = self._params()
Pengchong Jin's avatar
Pengchong Jin committed
216
    params['architecture']['use_bfloat16'] = False
David Chen's avatar
David Chen committed
217
    params['train']['total_steps'] = 1875  # One epoch.
Yeqing Li's avatar
Yeqing Li committed
218
219
220
221
222
223
224
    # 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
225
    params['eval']['eval_samples'] = 8
Jose Baiocchi's avatar
Jose Baiocchi committed
226
    FLAGS.num_gpus = 8
227
    FLAGS.model_dir = self._get_model_dir('real_benchmark_8_gpu_coco')
Jose Baiocchi's avatar
Jose Baiocchi committed
228
229
    FLAGS.strategy_type = 'mirrored'
    self._run_and_report_benchmark(params)
230

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
231
232
233
234
  @flagsaver.flagsaver
  def benchmark_1_gpu_coco(self):
    """Run RetinaNet model accuracy test with 1 GPU."""
    self._setup()
Jose Baiocchi's avatar
Jose Baiocchi committed
235
    params = self._params()
Pengchong Jin's avatar
Pengchong Jin committed
236
    params['architecture']['use_bfloat16'] = False
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
237
238
239
240
    params['train']['batch_size'] = 8
    params['train']['total_steps'] = 200
    params['train']['iterations_per_loop'] = 1
    params['eval']['eval_samples'] = 8
Jose Baiocchi's avatar
Jose Baiocchi committed
241
    FLAGS.num_gpus = 1
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
242
    FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
Yeqing Li's avatar
Yeqing Li committed
243
    FLAGS.strategy_type = 'one_device'
Jose Baiocchi's avatar
Jose Baiocchi committed
244
    self._run_and_report_benchmark(params)
245

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

  @flagsaver.flagsaver
  def benchmark_2x2_tpu_coco(self):
    """Run RetinaNet model accuracy test with 4 TPUs."""
    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
274
275
276
277
278
279
280
281
282
283
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
  @flagsaver.flagsaver
  def benchmark_4x4_tpu_coco(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')
    FLAGS.strategy_type = 'tpu'
    self._run_and_report_benchmark(params, do_eval=False, warmup=0)

  @flagsaver.flagsaver
  def benchmark_2x2_tpu_coco_mlir(self):
    """Run RetinaNet model accuracy test with 4 TPUs."""
    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)

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
  @flagsaver.flagsaver
  def benchmark_2x2_tpu_spinenet_coco(self):
    """Run SpineNet with RetinaNet model accuracy test with 4 TPUs."""
    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
329

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