resnet.py 15.6 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
"""Contains definitions of ResNet and ResNet-RS models."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
16

Fan Yang's avatar
Fan Yang committed
17
18
from typing import Callable, Optional

Abdullah Rashwan's avatar
Abdullah Rashwan committed
19
20
# Import libraries
import tensorflow as tf
Fan Yang's avatar
Fan Yang committed
21
22

from official.modeling import hyperparams
Abdullah Rashwan's avatar
Abdullah Rashwan committed
23
from official.modeling import tf_utils
Yeqing Li's avatar
Yeqing Li committed
24
from official.vision.beta.modeling.backbones import factory
Abdullah Rashwan's avatar
Abdullah Rashwan committed
25
from official.vision.beta.modeling.layers import nn_blocks
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
26
from official.vision.beta.modeling.layers import nn_layers
Abdullah Rashwan's avatar
Abdullah Rashwan committed
27
28
29
30
31
32
33
34

layers = tf.keras.layers

# Specifications for different ResNet variants.
# Each entry specifies block configurations of the particular ResNet variant.
# Each element in the block configuration is in the following format:
# (block_fn, num_filters, block_repeats)
RESNET_SPECS = {
Fan Yang's avatar
Fan Yang committed
35
36
37
38
39
40
    10: [
        ('residual', 64, 1),
        ('residual', 128, 1),
        ('residual', 256, 1),
        ('residual', 512, 1),
    ],
Abdullah Rashwan's avatar
Abdullah Rashwan committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
    18: [
        ('residual', 64, 2),
        ('residual', 128, 2),
        ('residual', 256, 2),
        ('residual', 512, 2),
    ],
    34: [
        ('residual', 64, 3),
        ('residual', 128, 4),
        ('residual', 256, 6),
        ('residual', 512, 3),
    ],
    50: [
        ('bottleneck', 64, 3),
        ('bottleneck', 128, 4),
        ('bottleneck', 256, 6),
        ('bottleneck', 512, 3),
    ],
    101: [
        ('bottleneck', 64, 3),
        ('bottleneck', 128, 4),
        ('bottleneck', 256, 23),
        ('bottleneck', 512, 3),
    ],
    152: [
        ('bottleneck', 64, 3),
        ('bottleneck', 128, 8),
        ('bottleneck', 256, 36),
        ('bottleneck', 512, 3),
    ],
    200: [
        ('bottleneck', 64, 3),
        ('bottleneck', 128, 24),
        ('bottleneck', 256, 36),
        ('bottleneck', 512, 3),
    ],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
77
    270: [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
78
        ('bottleneck', 64, 4),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
79
80
        ('bottleneck', 128, 29),
        ('bottleneck', 256, 53),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
81
82
        ('bottleneck', 512, 4),
    ],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
83
84
85
86
87
88
    350: [
        ('bottleneck', 64, 4),
        ('bottleneck', 128, 36),
        ('bottleneck', 256, 72),
        ('bottleneck', 512, 4),
    ],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
89
90
91
92
93
94
    420: [
        ('bottleneck', 64, 4),
        ('bottleneck', 128, 44),
        ('bottleneck', 256, 87),
        ('bottleneck', 512, 4),
    ],
Abdullah Rashwan's avatar
Abdullah Rashwan committed
95
96
97
98
99
}


@tf.keras.utils.register_keras_serializable(package='Vision')
class ResNet(tf.keras.Model):
100
  """Creates ResNet and ResNet-RS family models.
Fan Yang's avatar
Fan Yang committed
101
102
103
104

  This implements the Deep Residual Network from:
    Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun.
    Deep Residual Learning for Image Recognition.
105
106
107
108
109
    (https://arxiv.org/pdf/1512.03385) and
    Irwan Bello, William Fedus, Xianzhi Du, Ekin D. Cubuk, Aravind Srinivas,
    Tsung-Yi Lin, Jonathon Shlens, Barret Zoph.
    Revisiting ResNets: Improved Training and Scaling Strategies.
    (https://arxiv.org/abs/2103.07579).
Fan Yang's avatar
Fan Yang committed
110
  """
Abdullah Rashwan's avatar
Abdullah Rashwan committed
111

Fan Yang's avatar
Fan Yang committed
112
113
114
115
116
117
118
119
120
121
122
  def __init__(
      self,
      model_id: int,
      input_specs: tf.keras.layers.InputSpec = layers.InputSpec(
          shape=[None, None, None, 3]),
      depth_multiplier: float = 1.0,
      stem_type: str = 'v0',
      resnetd_shortcut: bool = False,
      replace_stem_max_pool: bool = False,
      se_ratio: Optional[float] = None,
      init_stochastic_depth_rate: float = 0.0,
Xianzhi Du's avatar
Xianzhi Du committed
123
      scale_stem: bool = True,
Fan Yang's avatar
Fan Yang committed
124
125
126
127
128
129
130
      activation: str = 'relu',
      use_sync_bn: bool = False,
      norm_momentum: float = 0.99,
      norm_epsilon: float = 0.001,
      kernel_initializer: str = 'VarianceScaling',
      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
      bias_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
131
      bn_trainable: bool = True,
Fan Yang's avatar
Fan Yang committed
132
      **kwargs):
Fan Yang's avatar
Fan Yang committed
133
    """Initializes a ResNet model.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
134
135

    Args:
Fan Yang's avatar
Fan Yang committed
136
137
138
      model_id: An `int` of the depth of ResNet backbone model.
      input_specs: A `tf.keras.layers.InputSpec` of the input tensor.
      depth_multiplier: A `float` of the depth multiplier to uniformaly scale up
139
140
        all layers in channel size. This argument is also referred to as
        `width_multiplier` in (https://arxiv.org/abs/2103.07579).
Fan Yang's avatar
Fan Yang committed
141
142
143
144
145
146
147
148
      stem_type: A `str` of stem type of ResNet. Default to `v0`. If set to
        `v1`, use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
      resnetd_shortcut: A `bool` of whether to use ResNet-D shortcut in
        downsampling blocks.
      replace_stem_max_pool: A `bool` of whether to replace the max pool in stem
        with a stride-2 conv,
      se_ratio: A `float` or None. Ratio of the Squeeze-and-Excitation layer.
      init_stochastic_depth_rate: A `float` of initial stochastic depth rate.
Xianzhi Du's avatar
Xianzhi Du committed
149
      scale_stem: A `bool` of whether to scale stem layers.
Fan Yang's avatar
Fan Yang committed
150
151
152
153
154
155
156
157
158
      activation: A `str` name of the activation function.
      use_sync_bn: If True, use synchronized batch normalization.
      norm_momentum: A `float` of normalization momentum for the moving average.
      norm_epsilon: A small `float` added to variance to avoid dividing by zero.
      kernel_initializer: A str for kernel initializer of convolutional layers.
      kernel_regularizer: A `tf.keras.regularizers.Regularizer` object for
        Conv2D. Default to None.
      bias_regularizer: A `tf.keras.regularizers.Regularizer` object for Conv2D.
        Default to None.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
159
160
      bn_trainable: A `bool` that indicates whether batch norm layers should be
        trainable. Default to True.
Fan Yang's avatar
Fan Yang committed
161
      **kwargs: Additional keyword arguments to be passed.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
162
163
164
    """
    self._model_id = model_id
    self._input_specs = input_specs
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
165
    self._depth_multiplier = depth_multiplier
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
166
    self._stem_type = stem_type
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
167
168
    self._resnetd_shortcut = resnetd_shortcut
    self._replace_stem_max_pool = replace_stem_max_pool
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
169
    self._se_ratio = se_ratio
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
170
    self._init_stochastic_depth_rate = init_stochastic_depth_rate
Xianzhi Du's avatar
Xianzhi Du committed
171
    self._scale_stem = scale_stem
Abdullah Rashwan's avatar
Abdullah Rashwan committed
172
173
174
175
176
177
178
179
180
181
182
    self._use_sync_bn = use_sync_bn
    self._activation = activation
    self._norm_momentum = norm_momentum
    self._norm_epsilon = norm_epsilon
    if use_sync_bn:
      self._norm = layers.experimental.SyncBatchNormalization
    else:
      self._norm = layers.BatchNormalization
    self._kernel_initializer = kernel_initializer
    self._kernel_regularizer = kernel_regularizer
    self._bias_regularizer = bias_regularizer
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
183
    self._bn_trainable = bn_trainable
Abdullah Rashwan's avatar
Abdullah Rashwan committed
184
185
186
187
188
189
190
191
192

    if tf.keras.backend.image_data_format() == 'channels_last':
      bn_axis = -1
    else:
      bn_axis = 1

    # Build ResNet.
    inputs = tf.keras.Input(shape=input_specs.shape[1:])

Xianzhi Du's avatar
Xianzhi Du committed
193
    stem_depth_multiplier = self._depth_multiplier if scale_stem else 1.0
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
194
195
    if stem_type == 'v0':
      x = layers.Conv2D(
Xianzhi Du's avatar
Xianzhi Du committed
196
          filters=int(64 * stem_depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
197
198
199
200
201
202
203
204
205
          kernel_size=7,
          strides=2,
          use_bias=False,
          padding='same',
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)(
              inputs)
      x = self._norm(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
206
207
208
209
          axis=bn_axis,
          momentum=norm_momentum,
          epsilon=norm_epsilon,
          trainable=bn_trainable)(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
210
              x)
211
      x = tf_utils.get_activation(activation, use_keras_layer=True)(x)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
212
213
    elif stem_type == 'v1':
      x = layers.Conv2D(
Xianzhi Du's avatar
Xianzhi Du committed
214
          filters=int(32 * stem_depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
215
216
217
218
219
220
221
222
223
          kernel_size=3,
          strides=2,
          use_bias=False,
          padding='same',
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)(
              inputs)
      x = self._norm(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
224
225
226
227
          axis=bn_axis,
          momentum=norm_momentum,
          epsilon=norm_epsilon,
          trainable=bn_trainable)(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
228
              x)
229
      x = tf_utils.get_activation(activation, use_keras_layer=True)(x)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
230
      x = layers.Conv2D(
Xianzhi Du's avatar
Xianzhi Du committed
231
          filters=int(32 * stem_depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
232
233
234
235
236
237
238
239
240
          kernel_size=3,
          strides=1,
          use_bias=False,
          padding='same',
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)(
              x)
      x = self._norm(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
241
242
243
244
          axis=bn_axis,
          momentum=norm_momentum,
          epsilon=norm_epsilon,
          trainable=bn_trainable)(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
245
              x)
246
      x = tf_utils.get_activation(activation, use_keras_layer=True)(x)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
247
      x = layers.Conv2D(
Xianzhi Du's avatar
Xianzhi Du committed
248
          filters=int(64 * stem_depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
249
250
251
252
253
254
255
256
257
          kernel_size=3,
          strides=1,
          use_bias=False,
          padding='same',
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)(
              x)
      x = self._norm(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
258
259
260
261
          axis=bn_axis,
          momentum=norm_momentum,
          epsilon=norm_epsilon,
          trainable=bn_trainable)(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
262
              x)
263
      x = tf_utils.get_activation(activation, use_keras_layer=True)(x)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
264
265
266
    else:
      raise ValueError('Stem type {} not supported.'.format(stem_type))

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
267
268
269
270
271
272
273
274
275
276
277
278
    if replace_stem_max_pool:
      x = layers.Conv2D(
          filters=int(64 * self._depth_multiplier),
          kernel_size=3,
          strides=2,
          use_bias=False,
          padding='same',
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer)(
              x)
      x = self._norm(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
279
280
281
282
          axis=bn_axis,
          momentum=norm_momentum,
          epsilon=norm_epsilon,
          trainable=bn_trainable)(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
283
              x)
284
      x = tf_utils.get_activation(activation, use_keras_layer=True)(x)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
285
286
    else:
      x = layers.MaxPool2D(pool_size=3, strides=2, padding='same')(x)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
287
288
289
290
291
292
293
294
295
296
297

    endpoints = {}
    for i, spec in enumerate(RESNET_SPECS[model_id]):
      if spec[0] == 'residual':
        block_fn = nn_blocks.ResidualBlock
      elif spec[0] == 'bottleneck':
        block_fn = nn_blocks.BottleneckBlock
      else:
        raise ValueError('Block fn `{}` is not supported.'.format(spec[0]))
      x = self._block_group(
          inputs=x,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
298
          filters=int(spec[1] * self._depth_multiplier),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
299
300
301
          strides=(1 if i == 0 else 2),
          block_fn=block_fn,
          block_repeats=spec[2],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
302
303
          stochastic_depth_drop_rate=nn_layers.get_stochastic_depth_rate(
              self._init_stochastic_depth_rate, i + 2, 5),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
304
          name='block_group_l{}'.format(i + 2))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
305
      endpoints[str(i + 2)] = x
Abdullah Rashwan's avatar
Abdullah Rashwan committed
306
307
308
309
310
311

    self._output_specs = {l: endpoints[l].get_shape() for l in endpoints}

    super(ResNet, self).__init__(inputs=inputs, outputs=endpoints, **kwargs)

  def _block_group(self,
Fan Yang's avatar
Fan Yang committed
312
313
314
315
316
317
318
                   inputs: tf.Tensor,
                   filters: int,
                   strides: int,
                   block_fn: Callable[..., tf.keras.layers.Layer],
                   block_repeats: int = 1,
                   stochastic_depth_drop_rate: float = 0.0,
                   name: str = 'block_group'):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
319
320
321
    """Creates one group of blocks for the ResNet model.

    Args:
Fan Yang's avatar
Fan Yang committed
322
323
324
325
326
327
328
329
330
331
332
      inputs: A `tf.Tensor` of size `[batch, channels, height, width]`.
      filters: An `int` number of filters for the first convolution of the
        layer.
      strides: An `int` stride to use for the first convolution of the layer.
        If greater than 1, this layer will downsample the input.
      block_fn: The type of block group. Either `nn_blocks.ResidualBlock` or
        `nn_blocks.BottleneckBlock`.
      block_repeats: An `int` number of blocks contained in the layer.
      stochastic_depth_drop_rate: A `float` of drop rate of the current block
        group.
      name: A `str` name for the block.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
333
334

    Returns:
Fan Yang's avatar
Fan Yang committed
335
      The output `tf.Tensor` of the block layer.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
336
337
338
339
340
    """
    x = block_fn(
        filters=filters,
        strides=strides,
        use_projection=True,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
341
        stochastic_depth_drop_rate=stochastic_depth_drop_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
342
        se_ratio=self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
343
        resnetd_shortcut=self._resnetd_shortcut,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
344
345
346
347
348
349
        kernel_initializer=self._kernel_initializer,
        kernel_regularizer=self._kernel_regularizer,
        bias_regularizer=self._bias_regularizer,
        activation=self._activation,
        use_sync_bn=self._use_sync_bn,
        norm_momentum=self._norm_momentum,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
350
351
        norm_epsilon=self._norm_epsilon,
        bn_trainable=self._bn_trainable)(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
352
353
354
355
356
357
358
            inputs)

    for _ in range(1, block_repeats):
      x = block_fn(
          filters=filters,
          strides=1,
          use_projection=False,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
359
          stochastic_depth_drop_rate=stochastic_depth_drop_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
360
          se_ratio=self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
361
          resnetd_shortcut=self._resnetd_shortcut,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
362
363
364
365
366
367
          kernel_initializer=self._kernel_initializer,
          kernel_regularizer=self._kernel_regularizer,
          bias_regularizer=self._bias_regularizer,
          activation=self._activation,
          use_sync_bn=self._use_sync_bn,
          norm_momentum=self._norm_momentum,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
368
369
          norm_epsilon=self._norm_epsilon,
          bn_trainable=self._bn_trainable)(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
370
371
              x)

372
    return tf.keras.layers.Activation('linear', name=name)(x)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
373
374
375
376

  def get_config(self):
    config_dict = {
        'model_id': self._model_id,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
377
        'depth_multiplier': self._depth_multiplier,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
378
        'stem_type': self._stem_type,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
379
380
        'resnetd_shortcut': self._resnetd_shortcut,
        'replace_stem_max_pool': self._replace_stem_max_pool,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
381
        'activation': self._activation,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
382
        'se_ratio': self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
383
        'init_stochastic_depth_rate': self._init_stochastic_depth_rate,
Xianzhi Du's avatar
Xianzhi Du committed
384
        'scale_stem': self._scale_stem,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
385
386
387
388
389
390
        'use_sync_bn': self._use_sync_bn,
        'norm_momentum': self._norm_momentum,
        'norm_epsilon': self._norm_epsilon,
        'kernel_initializer': self._kernel_initializer,
        'kernel_regularizer': self._kernel_regularizer,
        'bias_regularizer': self._bias_regularizer,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
391
        'bn_trainable': self._bn_trainable
Abdullah Rashwan's avatar
Abdullah Rashwan committed
392
393
394
395
396
397
398
399
400
401
402
    }
    return config_dict

  @classmethod
  def from_config(cls, config, custom_objects=None):
    return cls(**config)

  @property
  def output_specs(self):
    """A dict of {level: TensorShape} pairs for the model output."""
    return self._output_specs
Yeqing Li's avatar
Yeqing Li committed
403
404
405
406
407


@factory.register_backbone_builder('resnet')
def build_resnet(
    input_specs: tf.keras.layers.InputSpec,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
408
409
    backbone_config: hyperparams.Config,
    norm_activation_config: hyperparams.Config,
Rebecca Chen's avatar
Rebecca Chen committed
410
    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:  # pytype: disable=annotation-type-mismatch  # typed-keras
Abdullah Rashwan's avatar
Abdullah Rashwan committed
411
  """Builds ResNet backbone from a config."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
412
413
  backbone_type = backbone_config.type
  backbone_cfg = backbone_config.get()
Yeqing Li's avatar
Yeqing Li committed
414
415
416
417
418
419
  assert backbone_type == 'resnet', (f'Inconsistent backbone type '
                                     f'{backbone_type}')

  return ResNet(
      model_id=backbone_cfg.model_id,
      input_specs=input_specs,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
420
      depth_multiplier=backbone_cfg.depth_multiplier,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
421
      stem_type=backbone_cfg.stem_type,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
422
423
      resnetd_shortcut=backbone_cfg.resnetd_shortcut,
      replace_stem_max_pool=backbone_cfg.replace_stem_max_pool,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
424
      se_ratio=backbone_cfg.se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
425
      init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate,
Xianzhi Du's avatar
Xianzhi Du committed
426
      scale_stem=backbone_cfg.scale_stem,
Yeqing Li's avatar
Yeqing Li committed
427
428
429
430
      activation=norm_activation_config.activation,
      use_sync_bn=norm_activation_config.use_sync_bn,
      norm_momentum=norm_activation_config.norm_momentum,
      norm_epsilon=norm_activation_config.norm_epsilon,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
431
432
      kernel_regularizer=l2_regularizer,
      bn_trainable=backbone_cfg.bn_trainable)