Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added a flag to allow skipping the first projection in small ResNets #11176

Closed
wants to merge 1 commit into from
Closed

Added a flag to allow skipping the first projection in small ResNets #11176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions official/vision/modeling/backbones/resnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,7 @@ def __init__(
kernel_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
bias_regularizer: Optional[tf_keras.regularizers.Regularizer] = None,
bn_trainable: bool = True,
use_first_projection: bool = True,
**kwargs):
"""Initializes a ResNet model.

Expand Down Expand Up @@ -164,6 +165,8 @@ def __init__(
Default to None.
bn_trainable: A `bool` that indicates whether batch norm layers should be
trainable. Default to True.
use_first_projection: A `bool` of whether to use the first projection
shortcut for small ResNets. See https://github.com/tensorflow/models/issues/10583.
**kwargs: Additional keyword arguments to be passed.
"""
self._model_id = model_id
Expand All @@ -184,6 +187,7 @@ def __init__(
self._kernel_regularizer = kernel_regularizer
self._bias_regularizer = bias_regularizer
self._bn_trainable = bn_trainable
self._use_first_projection = use_first_projection

if tf_keras.backend.image_data_format() == 'channels_last':
self._bn_axis = -1
Expand All @@ -202,12 +206,18 @@ def __init__(
block_fn = nn_blocks.BottleneckBlock
else:
raise ValueError('Block fn `{}` is not supported.'.format(spec[0]))
use_first_projection = (
spec[0] == 'bottleneck'
or i > 0
or self._use_first_projection
)
x = self._block_group(
inputs=x,
filters=int(spec[1] * self._depth_multiplier),
strides=(1 if i == 0 else 2),
block_fn=block_fn,
block_repeats=spec[2],
use_first_projection=use_first_projection,
stochastic_depth_drop_rate=nn_layers.get_stochastic_depth_rate(
self._init_stochastic_depth_rate, i + 2, 5),
name='block_group_l{}'.format(i + 2))
Expand Down Expand Up @@ -326,6 +336,7 @@ def _block_group(self,
strides: int,
block_fn: Callable[..., tf_keras.layers.Layer],
block_repeats: int = 1,
use_first_projection: bool = True,
stochastic_depth_drop_rate: float = 0.0,
name: str = 'block_group'):
"""Creates one group of blocks for the ResNet model.
Expand All @@ -339,6 +350,8 @@ def _block_group(self,
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.
use_first_projection: A `bool` to determine whether to use the first
projection shortcut.
stochastic_depth_drop_rate: A `float` of drop rate of the current block
group.
name: A `str` name for the block.
Expand All @@ -349,7 +362,7 @@ def _block_group(self,
x = block_fn(
filters=filters,
strides=strides,
use_projection=True,
use_projection=use_first_projection,
stochastic_depth_drop_rate=stochastic_depth_drop_rate,
se_ratio=self._se_ratio,
resnetd_shortcut=self._resnetd_shortcut,
Expand Down Expand Up @@ -400,7 +413,8 @@ def get_config(self):
'kernel_initializer': self._kernel_initializer,
'kernel_regularizer': self._kernel_regularizer,
'bias_regularizer': self._bias_regularizer,
'bn_trainable': self._bn_trainable
'bn_trainable': self._bn_trainable,
'use_first_projection': self._use_first_projection
}
return config_dict

Expand Down Expand Up @@ -441,4 +455,5 @@ def build_resnet(
norm_momentum=norm_activation_config.norm_momentum,
norm_epsilon=norm_activation_config.norm_epsilon,
kernel_regularizer=l2_regularizer,
bn_trainable=backbone_cfg.bn_trainable)
bn_trainable=backbone_cfg.bn_trainable,
use_first_projection=backbone_cfg.use_first_projection)
35 changes: 34 additions & 1 deletion official/vision/modeling/backbones/resnet_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,38 @@ def test_network_creation(self, input_size, model_id,
self.assertAllEqual(
[1, input_size / 2**5, input_size / 2**5, 512 * endpoint_filter_scale],
endpoints['5'].shape.as_list())

@parameterized.parameters(
(128, 18, 1),
(128, 34, 1),
)
def test_network_creation_no_first_shortcut(self, input_size, model_id,
endpoint_filter_scale):
"""Test creation of ResNet family models."""
resnet_params = {
18: 11186112,
34: 21301696,
}
tf.keras.backend.set_image_data_format('channels_last')

network = resnet.ResNet(model_id=model_id, use_first_projection=False)
self.assertEqual(network.count_params(), resnet_params[model_id])

inputs = tf.keras.Input(shape=(input_size, input_size, 3), batch_size=1)
endpoints = network(inputs)

self.assertAllEqual(
[1, input_size / 2**2, input_size / 2**2, 64 * endpoint_filter_scale],
endpoints['2'].shape.as_list())
self.assertAllEqual(
[1, input_size / 2**3, input_size / 2**3, 128 * endpoint_filter_scale],
endpoints['3'].shape.as_list())
self.assertAllEqual(
[1, input_size / 2**4, input_size / 2**4, 256 * endpoint_filter_scale],
endpoints['4'].shape.as_list())
self.assertAllEqual(
[1, input_size / 2**5, input_size / 2**5, 512 * endpoint_filter_scale],
endpoints['5'].shape.as_list())

@combinations.generate(
combinations.combine(
Expand Down Expand Up @@ -137,7 +169,8 @@ def test_serialize_deserialize(self):
kernel_initializer='VarianceScaling',
kernel_regularizer=None,
bias_regularizer=None,
bn_trainable=True)
bn_trainable=True,
use_first_projection=True)
network = resnet.ResNet(**kwargs)

expected_config = dict(kwargs)
Expand Down
Loading