retinanet.py 14.5 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


@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
Fan Yang's avatar
Fan Yang committed
58
  aug_policy: Optional[str] = None
Abdullah Rashwan's avatar
Abdullah Rashwan committed
59
60
61
62
63
64
65
66
67
68
69
  skip_crowd_during_training: bool = True
  max_num_instances: int = 100


@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'
70
  decoder: common.DataDecoder = common.DataDecoder()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
71
72
  parser: Parser = Parser()
  shuffle_buffer_size: int = 10000
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
73
  file_type: str = 'tfrecord'
Abdullah Rashwan's avatar
Abdullah Rashwan committed
74
75
76
77
78
79
80
81
82
83
84
85


@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
86
  loss_weight: float = 1.0
Abdullah Rashwan's avatar
Abdullah Rashwan committed
87
88
89
90
91
92
93
  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
94
95
96
97
98
99
100
@dataclasses.dataclass
class AttributeHead(hyperparams.Config):
  name: str = ''
  type: str = 'regression'
  size: int = 1


Abdullah Rashwan's avatar
Abdullah Rashwan committed
101
102
103
104
105
@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
106
  attribute_heads: List[AttributeHead] = dataclasses.field(default_factory=list)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
107
108
109
110


@dataclasses.dataclass
class DetectionGenerator(hyperparams.Config):
Fan Yang's avatar
Fan Yang committed
111
  apply_nms: bool = True
Abdullah Rashwan's avatar
Abdullah Rashwan committed
112
113
114
115
  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
116
  nms_version: str = 'v2'  # `v2`, `v1`, `batched`.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
117
  use_cpu_nms: bool = False
Xianzhi Du's avatar
Xianzhi Du committed
118
  soft_nms_sigma: Optional[float] = None  # Only works when nms_version='v1'.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136


@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()


137
138
139
140
141
142
143
@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
144
145
146
147
148
149
150
@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
151
152
  init_checkpoint_modules: Union[
      str, List[str]] = 'all'  # all, backbone, and/or decoder
Zhenyu Tan's avatar
Zhenyu Tan committed
153
  annotation_file: Optional[str] = None
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
154
  per_category_metrics: bool = False
155
  export_config: ExportConfig = ExportConfig()
Abdullah Rashwan's avatar
Abdullah Rashwan committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169


@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'
170
COCO_TRAIN_EXAMPLES = 118287
Abdullah Rashwan's avatar
Abdullah Rashwan committed
171
172
173
174
175
176
177
178
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
179
  steps_per_epoch = COCO_TRAIN_EXAMPLES // train_batch_size
Abdullah Rashwan's avatar
Abdullah Rashwan committed
180
181
182
183
184
185

  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
186
187
          annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                       'instances_val2017.json'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
188
189
190
          model=RetinaNet(
              num_classes=91,
              input_size=[640, 640, 3],
Xianzhi Du's avatar
Xianzhi Du committed
191
              norm_activation=common.NormActivation(use_sync_bn=False),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
192
193
194
195
196
197
198
199
              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
200
                  aug_rand_hflip=True, aug_scale_min=0.8, aug_scale_max=1.2)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
          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
226
227
228
                          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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
                      ],
                  }
              },
              '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
253
  steps_per_epoch = COCO_TRAIN_EXAMPLES // train_batch_size
Abdullah Rashwan's avatar
Abdullah Rashwan committed
254
255
256
257
258
  input_size = 640

  config = cfg.ExperimentConfig(
      runtime=cfg.RuntimeConfig(mixed_precision_dtype='float32'),
      task=RetinaNetTask(
Zhenyu Tan's avatar
Zhenyu Tan committed
259
260
          annotation_file=os.path.join(COCO_INPUT_PATH_BASE,
                                       'instances_val2017.json'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
261
262
263
          model=RetinaNet(
              backbone=backbones.Backbone(
                  type='spinenet',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
264
                  spinenet=backbones.SpineNet(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
265
266
267
268
                      model_id='49',
                      stochastic_depth_drop_rate=0.2,
                      min_level=3,
                      max_level=7)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
269
270
271
              decoder=decoders.Decoder(
                  type='identity', identity=decoders.Identity()),
              anchor=Anchor(anchor_size=3),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
272
273
              norm_activation=common.NormActivation(
                  use_sync_bn=True, activation='swish'),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
274
275
276
277
278
279
280
281
282
283
              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
284
                  aug_rand_hflip=True, aug_scale_min=0.1, aug_scale_max=2.0)),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
285
286
287
288
289
          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
290
          train_steps=500 * steps_per_epoch,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
          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
307
                          475 * steps_per_epoch, 490 * steps_per_epoch
Abdullah Rashwan's avatar
Abdullah Rashwan committed
308
309
                      ],
                      'values': [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
310
311
312
                          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
313
314
315
316
317
318
319
320
321
322
323
324
325
                      ],
                  }
              },
              '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
326
          'task.validation_data.is_training != None',
Xianzhi Du's avatar
Xianzhi Du committed
327
328
          '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
329
330
331
      ])

  return config
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
332
333


Xianzhi Du's avatar
Xianzhi Du committed
334
@exp_factory.register_config_factory('retinanet_mobile_coco')
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
335
def retinanet_spinenet_mobile_coco() -> cfg.ExperimentConfig:
Xianzhi Du's avatar
Xianzhi Du committed
336
  """COCO object detection with mobile RetinaNet."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
  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
351
352
353
                      model_id='49',
                      stochastic_depth_drop_rate=0.2,
                      min_level=3,
354
355
                      max_level=7,
                      use_keras_upsampling_2d=False)),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
356
357
              decoder=decoders.Decoder(
                  type='identity', identity=decoders.Identity()),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
358
              head=RetinaNetHead(num_filters=48, use_separable_conv=True),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
359
360
361
362
363
364
365
              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
366
          losses=Losses(l2_weight_decay=3e-5),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
367
368
369
370
371
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
          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
414
          'task.validation_data.is_training != None',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
415
416
417
      ])

  return config