resnet.py 13.7 KB
Newer Older
Abdullah Rashwan's avatar
Abdullah Rashwan committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Copyright 2020 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.
# ==============================================================================
"""Contains definitions of Residual Networks.

Residual networks (ResNets) were proposed in:
[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun
    Deep Residual Learning for Image Recognition. arXiv:1512.03385
"""

# Import libraries
import tensorflow as tf
from official.modeling import tf_utils
Yeqing Li's avatar
Yeqing Li committed
25
from official.vision.beta.modeling.backbones import factory
Abdullah Rashwan's avatar
Abdullah Rashwan committed
26
from official.vision.beta.modeling.layers import nn_blocks
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
27
from official.vision.beta.modeling.layers import nn_layers
Abdullah Rashwan's avatar
Abdullah Rashwan committed
28
29
30
31
32
33
34
35
36
37
38
39
40
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

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 = {
    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
72
    270: [
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
73
        ('bottleneck', 64, 4),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
74
75
        ('bottleneck', 128, 29),
        ('bottleneck', 256, 53),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
76
77
        ('bottleneck', 512, 4),
    ],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
78
79
80
81
82
83
    350: [
        ('bottleneck', 64, 4),
        ('bottleneck', 128, 36),
        ('bottleneck', 256, 72),
        ('bottleneck', 512, 4),
    ],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
84
85
86
87
88
89
    420: [
        ('bottleneck', 64, 4),
        ('bottleneck', 128, 44),
        ('bottleneck', 256, 87),
        ('bottleneck', 512, 4),
    ],
Abdullah Rashwan's avatar
Abdullah Rashwan committed
90
91
92
93
94
95
96
97
98
99
}


@tf.keras.utils.register_keras_serializable(package='Vision')
class ResNet(tf.keras.Model):
  """Class to build ResNet family model."""

  def __init__(self,
               model_id,
               input_specs=layers.InputSpec(shape=[None, None, None, 3]),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
100
               depth_multiplier=1.0,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
101
               stem_type='v0',
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
102
103
               resnetd_shortcut=False,
               replace_stem_max_pool=False,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
104
               se_ratio=None,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
105
               init_stochastic_depth_rate=0.0,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
106
107
108
109
110
111
112
113
114
115
116
117
118
               activation='relu',
               use_sync_bn=False,
               norm_momentum=0.99,
               norm_epsilon=0.001,
               kernel_initializer='VarianceScaling',
               kernel_regularizer=None,
               bias_regularizer=None,
               **kwargs):
    """ResNet initialization function.

    Args:
      model_id: `int` depth of ResNet backbone model.
      input_specs: `tf.keras.layers.InputSpec` specs of the input tensor.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
119
120
      depth_multiplier: `float` a depth multiplier to uniformaly scale up all
        layers in channel size in ResNet.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
121
      stem_type: `str` stem type of ResNet. Default to `v0`. If set to `v1`,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
122
123
124
125
126
        use ResNet-D type stem (https://arxiv.org/abs/1812.01187).
      resnetd_shortcut: `bool` whether to use ResNet-D shortcut in downsampling
        blocks.
      replace_stem_max_pool: `bool` if True, replace the max pool in stem with
        a stride-2 conv,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
127
      se_ratio: `float` or None. Ratio of the Squeeze-and-Excitation layer.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
128
      init_stochastic_depth_rate: `float` initial stochastic depth rate.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
      activation: `str` name of the activation function.
      use_sync_bn: if True, use synchronized batch normalization.
      norm_momentum: `float` normalization omentum for the moving average.
      norm_epsilon: `float` small float added to variance to avoid dividing by
        zero.
      kernel_initializer: kernel_initializer for convolutional layers.
      kernel_regularizer: tf.keras.regularizers.Regularizer object for Conv2D.
                          Default to None.
      bias_regularizer: tf.keras.regularizers.Regularizer object for Conv2d.
                        Default to None.
      **kwargs: keyword arguments to be passed.
    """
    self._model_id = model_id
    self._input_specs = input_specs
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
143
    self._depth_multiplier = depth_multiplier
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
144
    self._stem_type = stem_type
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
145
146
    self._resnetd_shortcut = resnetd_shortcut
    self._replace_stem_max_pool = replace_stem_max_pool
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
147
    self._se_ratio = se_ratio
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
148
    self._init_stochastic_depth_rate = init_stochastic_depth_rate
Abdullah Rashwan's avatar
Abdullah Rashwan committed
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    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

    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:])

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
169
170
    if stem_type == 'v0':
      x = layers.Conv2D(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
171
          filters=int(64 * self._depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
172
173
174
175
176
177
178
179
180
181
182
183
184
185
          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(
          axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
              x)
      x = tf_utils.get_activation(activation)(x)
    elif stem_type == 'v1':
      x = layers.Conv2D(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
186
          filters=int(32 * self._depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
187
188
189
190
191
192
193
194
195
196
197
198
199
          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(
          axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
              x)
      x = tf_utils.get_activation(activation)(x)
      x = layers.Conv2D(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
200
          filters=int(32 * self._depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
201
202
203
204
205
206
207
208
209
210
211
212
213
          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(
          axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
              x)
      x = tf_utils.get_activation(activation)(x)
      x = layers.Conv2D(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
214
          filters=int(64 * self._depth_multiplier),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
          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(
          axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
              x)
      x = tf_utils.get_activation(activation)(x)
    else:
      raise ValueError('Stem type {} not supported.'.format(stem_type))

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
    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(
          axis=bn_axis, momentum=norm_momentum, epsilon=norm_epsilon)(
              x)
      x = tf_utils.get_activation(activation)(x)
    else:
      x = layers.MaxPool2D(pool_size=3, strides=2, padding='same')(x)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
247
248
249
250
251
252
253
254
255
256
257

    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
258
          filters=int(spec[1] * self._depth_multiplier),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
259
260
261
          strides=(1 if i == 0 else 2),
          block_fn=block_fn,
          block_repeats=spec[2],
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
262
263
          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
264
          name='block_group_l{}'.format(i + 2))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
265
      endpoints[str(i + 2)] = x
Abdullah Rashwan's avatar
Abdullah Rashwan committed
266
267
268
269
270
271
272
273
274
275
276

    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,
                   inputs,
                   filters,
                   strides,
                   block_fn,
                   block_repeats=1,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
277
                   stochastic_depth_drop_rate=0.0,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
278
279
280
281
282
283
284
285
286
287
                   name='block_group'):
    """Creates one group of blocks for the ResNet model.

    Args:
      inputs: `Tensor` of size `[batch, channels, height, width]`.
      filters: `int` number of filters for the first convolution of the layer.
      strides: `int` stride to use for the first convolution of the layer. If
        greater than 1, this layer will downsample the input.
      block_fn: Either `nn_blocks.ResidualBlock` or `nn_blocks.BottleneckBlock`.
      block_repeats: `int` number of blocks contained in the layer.
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
288
      stochastic_depth_drop_rate: `float` drop rate of the current block group.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
289
290
291
292
293
294
295
296
297
      name: `str`name for the block.

    Returns:
      The output `Tensor` of the block layer.
    """
    x = block_fn(
        filters=filters,
        strides=strides,
        use_projection=True,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
298
        stochastic_depth_drop_rate=stochastic_depth_drop_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
299
        se_ratio=self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
300
        resnetd_shortcut=self._resnetd_shortcut,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
        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,
        norm_epsilon=self._norm_epsilon)(
            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
315
          stochastic_depth_drop_rate=stochastic_depth_drop_rate,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
316
          se_ratio=self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
317
          resnetd_shortcut=self._resnetd_shortcut,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
318
319
320
321
322
323
324
325
326
327
328
329
330
331
          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,
          norm_epsilon=self._norm_epsilon)(
              x)

    return tf.identity(x, name=name)

  def get_config(self):
    config_dict = {
        'model_id': self._model_id,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
332
        'depth_multiplier': self._depth_multiplier,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
333
        'stem_type': self._stem_type,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
334
335
        'resnetd_shortcut': self._resnetd_shortcut,
        'replace_stem_max_pool': self._replace_stem_max_pool,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
336
        'activation': self._activation,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
337
        'se_ratio': self._se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
338
        'init_stochastic_depth_rate': self._init_stochastic_depth_rate,
Abdullah Rashwan's avatar
Abdullah Rashwan committed
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
        '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,
    }
    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
356
357
358
359
360
361
362


@factory.register_backbone_builder('resnet')
def build_resnet(
    input_specs: tf.keras.layers.InputSpec,
    model_config,
    l2_regularizer: tf.keras.regularizers.Regularizer = None) -> tf.keras.Model:
Abdullah Rashwan's avatar
Abdullah Rashwan committed
363
  """Builds ResNet backbone from a config."""
Yeqing Li's avatar
Yeqing Li committed
364
365
366
367
368
369
370
371
372
  backbone_type = model_config.backbone.type
  backbone_cfg = model_config.backbone.get()
  norm_activation_config = model_config.norm_activation
  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
373
      depth_multiplier=backbone_cfg.depth_multiplier,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
374
      stem_type=backbone_cfg.stem_type,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
375
376
      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
377
      se_ratio=backbone_cfg.se_ratio,
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
378
      init_stochastic_depth_rate=backbone_cfg.stochastic_depth_drop_rate,
Yeqing Li's avatar
Yeqing Li committed
379
380
381
382
383
      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,
      kernel_regularizer=l2_regularizer)