retinanet.py 14.7 KB
Newer Older
Yeqing Li's avatar
Yeqing Li committed
1
# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# 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.
Yeqing Li's avatar
Yeqing Li committed
14
15

# Lint as: python3
Abdullah Rashwan's avatar
Abdullah Rashwan committed
16
17
"""RetinaNet configuration definition."""

18
import dataclasses
Abdullah Rashwan's avatar
Abdullah Rashwan committed
19
import os
Xianzhi Du's avatar
Xianzhi Du committed
20
from typing import List, Optional, Union
Abdullah Rashwan's avatar
Abdullah Rashwan committed
21

22
from official.core import config_definitions as cfg
Abdullah Rashwan's avatar
Abdullah Rashwan committed
23
24
25
26
27
from official.core import exp_factory
from official.modeling import hyperparams
from official.modeling import optimization
from official.vision.beta.configs import common
from official.vision.beta.configs import decoders
Abdullah Rashwan's avatar
Abdullah Rashwan committed
28
from official.vision.beta.configs import backbones
Abdullah Rashwan's avatar
Abdullah Rashwan committed
29
30
31


# pylint: disable=missing-class-docstring
32
# Keep for backward compatibility.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
33
@dataclasses.dataclass
34
35
class TfExampleDecoder(common.TfExampleDecoder):
  """A simple TF Example decoder config."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
36
37


38
# Keep for backward compatibility.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
39
@dataclasses.dataclass
40
41
class TfExampleDecoderLabelMap(common.TfExampleDecoderLabelMap):
  """TF Example decoder with label map config."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
42
43


44
# Keep for backward compatibility.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
45
@dataclasses.dataclass
46
47
class DataDecoder(common.DataDecoder):
  """Data decoder config."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
48
49
50
51
52
53
54
55
56
57
58
59


@dataclasses.dataclass
class Parser(hyperparams.Config):
  num_channels: int = 3
  match_threshold: float = 0.5
  unmatched_threshold: float = 0.5
  aug_rand_hflip: bool = False
  aug_scale_min: float = 1.0
  aug_scale_max: float = 1.0
  skip_crowd_during_training: bool = True
  max_num_instances: int = 100
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
60
61
62
63
64
65
  # Can choose AutoAugment and RandAugment.
  # TODO(b/205346436) Support RandAugment.
  aug_type: Optional[common.Augmentation] = None

  # Keep for backward compatibility. Not used.
  aug_policy: Optional[str] = None
Abdullah Rashwan's avatar
Abdullah Rashwan committed
66
67
68
69
70
71
72
73
74


@dataclasses.dataclass
class DataConfig(cfg.DataConfig):
  """Input config for training."""
  input_path: str = ''
  global_batch_size: int = 0
  is_training: bool = False
  dtype: str = 'bfloat16'
75
  decoder: common.DataDecoder = common.DataDecoder()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
76
77
  parser: Parser = Parser()
  shuffle_buffer_size: int = 10000
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
78
  file_type: str = 'tfrecord'
Abdullah Rashwan's avatar
Abdullah Rashwan committed
79
80
81
82
83
84
85
86
87
88
89
90


@dataclasses.dataclass
class Anchor(hyperparams.Config):
  num_scales: int = 3
  aspect_ratios: List[float] = dataclasses.field(
      default_factory=lambda: [0.5, 1.0, 2.0])
  anchor_size: float = 4.0


@dataclasses.dataclass
class Losses(hyperparams.Config):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
91
  loss_weight: float = 1.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
92
93
94
95
96
97
98
  focal_loss_alpha: float = 0.25
  focal_loss_gamma: float = 1.5
  huber_loss_delta: float = 0.1
  box_loss_weight: int = 50
  l2_weight_decay: float = 0.0


A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
99
100
101
102
103
104
105
@dataclasses.dataclass
class AttributeHead(hyperparams.Config):
  name: str = ''
  type: str = 'regression'
  size: int = 1


Abdullah Rashwan's avatar
Abdullah Rashwan committed
106
107
108
109
110
@dataclasses.dataclass
class RetinaNetHead(hyperparams.Config):
  num_convs: int = 4
  num_filters: int = 256
  use_separable_conv: bool = False
Xianzhi Du's avatar
Xianzhi Du committed
111
  attribute_heads: List[AttributeHead] = dataclasses.field(default_factory=list)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
112
113
114
115


@dataclasses.dataclass
class DetectionGenerator(hyperparams.Config):
Fan Yang's avatar
Fan Yang committed
116
  apply_nms: bool = True
Abdullah Rashwan's avatar
Abdullah Rashwan committed
117
118
119
120
  pre_nms_top_k: int = 5000
  pre_nms_score_threshold: float = 0.05
  nms_iou_threshold: float = 0.5
  max_num_detections: int = 100
Xianzhi Du's avatar
Xianzhi Du committed
121
  nms_version: str = 'v2'  # `v2`, `v1`, `batched`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
122
  use_cpu_nms: bool = False
Xianzhi Du's avatar
Xianzhi Du committed
123
  soft_nms_sigma: Optional[float] = None  # Only works when nms_version='v1'.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141


@dataclasses.dataclass
class RetinaNet(hyperparams.Config):
  num_classes: int = 0
  input_size: List[int] = dataclasses.field(default_factory=list)
  min_level: int = 3
  max_level: int = 7
  anchor: Anchor = Anchor()
  backbone: backbones.Backbone = backbones.Backbone(
      type='resnet', resnet=backbones.ResNet())
  decoder: decoders.Decoder = decoders.Decoder(
      type='fpn', fpn=decoders.FPN())
  head: RetinaNetHead = RetinaNetHead()
  detection_generator: DetectionGenerator = DetectionGenerator()
  norm_activation: common.NormActivation = common.NormActivation()


142
143
144
145
146
147
148
@dataclasses.dataclass
class ExportConfig(hyperparams.Config):
  output_normalized_coordinates: bool = False
  cast_num_detections_to_float: bool = False
  cast_detection_classes_to_float: bool = False


Abdullah Rashwan's avatar
Abdullah Rashwan committed
149
150
151
152
153
154
155
@dataclasses.dataclass
class RetinaNetTask(cfg.TaskConfig):
  model: RetinaNet = RetinaNet()
  train_data: DataConfig = DataConfig(is_training=True)
  validation_data: DataConfig = DataConfig(is_training=False)
  losses: Losses = Losses()
  init_checkpoint: Optional[str] = None
Xianzhi Du's avatar
Xianzhi Du committed
156
157
  init_checkpoint_modules: Union[
      str, List[str]] = 'all'  # all, backbone, and/or decoder
Zhenyu Tan's avatar
Zhenyu Tan committed
158
  annotation_file: Optional[str] = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
159
  per_category_metrics: bool = False
160
  export_config: ExportConfig = ExportConfig()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174


@exp_factory.register_config_factory('retinanet')
def retinanet() -> cfg.ExperimentConfig:
  """RetinaNet general config."""
  return cfg.ExperimentConfig(
      task=RetinaNetTask(),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
      ])


COCO_INPUT_PATH_BASE = 'coco'
175
COCO_TRAIN_EXAMPLES = 118287
Abdullah Rashwan's avatar
Abdullah Rashwan committed
176
177
178
179
180
181
182
183
COCO_VAL_EXAMPLES = 5000


@exp_factory.register_config_factory('retinanet_resnetfpn_coco')
def retinanet_resnetfpn_coco() -> cfg.ExperimentConfig:
  """COCO object detection with RetinaNet."""
  train_batch_size = 256
  eval_batch_size = 8
184
  steps_per_epoch = COCO_TRAIN_EXAMPLES // train_batch_size
Abdullah Rashwan's avatar
Abdullah Rashwan committed
185
186
187
188
189
190

  config = cfg.ExperimentConfig(
      runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
      task=RetinaNetTask(
          init_checkpoint='gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080',
          init_checkpoint_modules='backbone',
Zhenyu Tan's avatar
Zhenyu Tan committed
191
192
          annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                       'instances_val2017.json'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
193
194
195
          model=RetinaNet(
              num_classes=91,
              input_size=[640, 640, 3],
Xianzhi Du's avatar
Xianzhi Du committed
196
              norm_activation=common.NormActivation(use_sync_bn=False),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
197
198
199
200
201
202
203
204
              min_level=3,
              max_level=7),
          losses=Losses(l2_weight_decay=1e-4),
          train_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
              is_training=True,
              global_batch_size=train_batch_size,
              parser=Parser(
Xianzhi Du's avatar
Xianzhi Du committed
205
                  aug_rand_hflip=True, aug_scale_min=0.8, aug_scale_max=1.2)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
          validation_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
              is_training=False,
              global_batch_size=eval_batch_size)),
      trainer=cfg.TrainerConfig(
          train_steps=72 * steps_per_epoch,
          validation_steps=COCO_VAL_EXAMPLES // eval_batch_size,
          validation_interval=steps_per_epoch,
          steps_per_loop=steps_per_epoch,
          summary_interval=steps_per_epoch,
          checkpoint_interval=steps_per_epoch,
          optimizer_config=optimization.OptimizationConfig({
              'optimizer': {
                  'type': 'sgd',
                  'sgd': {
                      'momentum': 0.9
                  }
              },
              'learning_rate': {
                  'type': 'stepwise',
                  'stepwise': {
                      'boundaries': [
                          57 * steps_per_epoch, 67 * steps_per_epoch
                      ],
                      'values': [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
231
232
233
                          0.32 * train_batch_size / 256.0,
                          0.032 * train_batch_size / 256.0,
                          0.0032 * train_batch_size / 256.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
                      ],
                  }
              },
              'warmup': {
                  'type': 'linear',
                  'linear': {
                      'warmup_steps': 500,
                      'warmup_learning_rate': 0.0067
                  }
              }
          })),
      restrictions=[
          'task.train_data.is_training != None',
          'task.validation_data.is_training != None'
      ])

  return config


@exp_factory.register_config_factory('retinanet_spinenet_coco')
def retinanet_spinenet_coco() -> cfg.ExperimentConfig:
  """COCO object detection with RetinaNet using SpineNet backbone."""
  train_batch_size = 256
  eval_batch_size = 8
258
  steps_per_epoch = COCO_TRAIN_EXAMPLES // train_batch_size
Abdullah Rashwan's avatar
Abdullah Rashwan committed
259
260
261
262
263
  input_size = 640

  config = cfg.ExperimentConfig(
      runtime=cfg.RuntimeConfig(mixed_precision_dtype='float32'),
      task=RetinaNetTask(
Zhenyu Tan's avatar
Zhenyu Tan committed
264
265
          annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                       'instances_val2017.json'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
266
267
268
          model=RetinaNet(
              backbone=backbones.Backbone(
                  type='spinenet',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
269
                  spinenet=backbones.SpineNet(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
270
271
272
273
                      model_id='49',
                      stochastic_depth_drop_rate=0.2,
                      min_level=3,
                      max_level=7)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
274
275
276
              decoder=decoders.Decoder(
                  type='identity', identity=decoders.Identity()),
              anchor=Anchor(anchor_size=3),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
277
278
              norm_activation=common.NormActivation(
                  use_sync_bn=True, activation='swish'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
279
280
281
282
283
284
285
286
287
288
              num_classes=91,
              input_size=[input_size, input_size, 3],
              min_level=3,
              max_level=7),
          losses=Losses(l2_weight_decay=4e-5),
          train_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
              is_training=True,
              global_batch_size=train_batch_size,
              parser=Parser(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
289
                  aug_rand_hflip=True, aug_scale_min=0.1, aug_scale_max=2.0)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
290
291
292
293
294
          validation_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
              is_training=False,
              global_batch_size=eval_batch_size)),
      trainer=cfg.TrainerConfig(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
295
          train_steps=500 * steps_per_epoch,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
          validation_steps=COCO_VAL_EXAMPLES // eval_batch_size,
          validation_interval=steps_per_epoch,
          steps_per_loop=steps_per_epoch,
          summary_interval=steps_per_epoch,
          checkpoint_interval=steps_per_epoch,
          optimizer_config=optimization.OptimizationConfig({
              'optimizer': {
                  'type': 'sgd',
                  'sgd': {
                      'momentum': 0.9
                  }
              },
              'learning_rate': {
                  'type': 'stepwise',
                  'stepwise': {
                      'boundaries': [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
312
                          475 * steps_per_epoch, 490 * steps_per_epoch
Abdullah Rashwan's avatar
Abdullah Rashwan committed
313
314
                      ],
                      'values': [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
315
316
317
                          0.32 * train_batch_size / 256.0,
                          0.032 * train_batch_size / 256.0,
                          0.0032 * train_batch_size / 256.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
318
319
320
321
322
323
324
325
326
327
328
329
330
                      ],
                  }
              },
              'warmup': {
                  'type': 'linear',
                  'linear': {
                      'warmup_steps': 2000,
                      'warmup_learning_rate': 0.0067
                  }
              }
          })),
      restrictions=[
          'task.train_data.is_training != None',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
331
          'task.validation_data.is_training != None',
Xianzhi Du's avatar
Xianzhi Du committed
332
333
          'task.model.min_level == task.model.backbone.spinenet.min_level',
          'task.model.max_level == task.model.backbone.spinenet.max_level',
Abdullah Rashwan's avatar
Abdullah Rashwan committed
334
335
336
      ])

  return config
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
337
338


Xianzhi Du's avatar
Xianzhi Du committed
339
@exp_factory.register_config_factory('retinanet_mobile_coco')
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
340
def retinanet_spinenet_mobile_coco() -> cfg.ExperimentConfig:
Xianzhi Du's avatar
Xianzhi Du committed
341
  """COCO object detection with mobile RetinaNet."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
342
343
344
345
346
347
348
349
350
351
352
353
354
355
  train_batch_size = 256
  eval_batch_size = 8
  steps_per_epoch = COCO_TRAIN_EXAMPLES // train_batch_size
  input_size = 384

  config = cfg.ExperimentConfig(
      runtime=cfg.RuntimeConfig(mixed_precision_dtype='float32'),
      task=RetinaNetTask(
          annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                       'instances_val2017.json'),
          model=RetinaNet(
              backbone=backbones.Backbone(
                  type='spinenet_mobile',
                  spinenet_mobile=backbones.SpineNetMobile(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
356
357
358
                      model_id='49',
                      stochastic_depth_drop_rate=0.2,
                      min_level=3,
359
360
                      max_level=7,
                      use_keras_upsampling_2d=False)),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
361
362
              decoder=decoders.Decoder(
                  type='identity', identity=decoders.Identity()),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
363
              head=RetinaNetHead(num_filters=48, use_separable_conv=True),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
364
365
366
367
368
369
370
              anchor=Anchor(anchor_size=3),
              norm_activation=common.NormActivation(
                  use_sync_bn=True, activation='swish'),
              num_classes=91,
              input_size=[input_size, input_size, 3],
              min_level=3,
              max_level=7),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
371
          losses=Losses(l2_weight_decay=3e-5),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
          train_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'train*'),
              is_training=True,
              global_batch_size=train_batch_size,
              parser=Parser(
                  aug_rand_hflip=True, aug_scale_min=0.1, aug_scale_max=2.0)),
          validation_data=DataConfig(
              input_path=os.path.join(COCO_INPUT_PATH_BASE, 'val*'),
              is_training=False,
              global_batch_size=eval_batch_size)),
      trainer=cfg.TrainerConfig(
          train_steps=600 * steps_per_epoch,
          validation_steps=COCO_VAL_EXAMPLES // eval_batch_size,
          validation_interval=steps_per_epoch,
          steps_per_loop=steps_per_epoch,
          summary_interval=steps_per_epoch,
          checkpoint_interval=steps_per_epoch,
          optimizer_config=optimization.OptimizationConfig({
              'optimizer': {
                  'type': 'sgd',
                  'sgd': {
                      'momentum': 0.9
                  }
              },
              'learning_rate': {
                  'type': 'stepwise',
                  'stepwise': {
                      'boundaries': [
                          575 * steps_per_epoch, 590 * steps_per_epoch
                      ],
                      'values': [
                          0.32 * train_batch_size / 256.0,
                          0.032 * train_batch_size / 256.0,
                          0.0032 * train_batch_size / 256.0
                      ],
                  }
              },
              'warmup': {
                  'type': 'linear',
                  'linear': {
                      'warmup_steps': 2000,
                      'warmup_learning_rate': 0.0067
                  }
              }
          })),
      restrictions=[
          'task.train_data.is_training != None',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
419
          'task.validation_data.is_training != None',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
420
421
422
      ])

  return config