official/vision/modeling/backbones/mobiledet.py
# Copyright 2024 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.
"""Definitions of MobileDet Networks."""
import dataclasses
from typing import Any, Dict, Optional, Tuple, List
import tensorflow as tf, tf_keras
from official.modeling import hyperparams
from official.vision.modeling.backbones import factory
from official.vision.modeling.backbones import mobilenet
from official.vision.modeling.layers import nn_blocks
from official.vision.modeling.layers import nn_layers
layers = tf_keras.layers
# pylint: disable=pointless-string-statement
"""
Architecture: https://arxiv.org/abs/2004.14525.
"MobileDets: Searching for Object Detection Architectures for
Mobile Accelerators" Yunyang Xiong, Hanxiao Liu, Suyog Gupta, Berkin Akin,
Gabriel Bender, Yongzhe Wang, Pieter-Jan Kindermans, Mingxing Tan, Vikas Singh,
Bo Chen
Note that `round_down_protection` flag should be set to false for scaling
of the network.
"""
MD_CPU_BLOCK_SPECS = {
'spec_name': 'MobileDetCPU',
# [expand_ratio] is set to 1 and [use_residual] is set to false
# for inverted_bottleneck_no_expansion
# [se_ratio] is set to 0.25 for all inverted_bottleneck layers
# [activation] is set to 'hard_swish' for all applicable layers
'block_spec_schema': ['block_fn', 'kernel_size', 'strides', 'filters',
'activation', 'se_ratio', 'expand_ratio',
'use_residual', 'is_output'],
'block_specs': [
('convbn', 3, 2, 16, 'hard_swish', None, None, None, False),
# inverted_bottleneck_no_expansion
('invertedbottleneck', 3, 1, 8, 'hard_swish', 0.25, 1., False, True),
('invertedbottleneck', 3, 2, 16, 'hard_swish', 0.25, 4., False, True),
('invertedbottleneck', 3, 2, 32, 'hard_swish', 0.25, 8., False, False),
('invertedbottleneck', 3, 1, 32, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 32, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 32, 'hard_swish', 0.25, 4., True, True),
('invertedbottleneck', 5, 2, 72, 'hard_swish', 0.25, 8., False, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 8., True, False),
('invertedbottleneck', 5, 1, 72, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 8., False, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 8., True, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 8., True, False),
('invertedbottleneck', 3, 1, 72, 'hard_swish', 0.25, 8., True, True),
('invertedbottleneck', 5, 2, 104, 'hard_swish', 0.25, 8., False, False),
('invertedbottleneck', 5, 1, 104, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 5, 1, 104, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 104, 'hard_swish', 0.25, 4., True, False),
('invertedbottleneck', 3, 1, 144, 'hard_swish', 0.25, 8., False, True),
]
}
MD_DSP_BLOCK_SPECS = {
'spec_name': 'MobileDetDSP',
# [expand_ratio] is set to 1 and [use_residual] is set to false
# for inverted_bottleneck_no_expansion
# [use_depthwise] is set to False for fused_conv
# [se_ratio] is set to None for all inverted_bottleneck layers
# [activation] is set to 'relu6' for all applicable layers
'block_spec_schema': ['block_fn', 'kernel_size', 'strides', 'filters',
'activation', 'se_ratio', 'expand_ratio',
'input_compression_ratio', 'output_compression_ratio',
'use_depthwise', 'use_residual', 'is_output'],
'block_specs': [
('convbn', 3, 2, 32, 'relu6',
None, None, None, None, None, None, False),
# inverted_bottleneck_no_expansion
('invertedbottleneck', 3, 1, 24, 'relu6',
None, 1., None, None, True, False, True),
('invertedbottleneck', 3, 2, 32, 'relu6',
None, 4., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 32, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 32, 'relu6',
None, 4., None, None, True, True, False),
('tucker', 3, 1, 32, 'relu6',
None, None, 0.25, 0.75, None, True, True),
('invertedbottleneck', 3, 2, 64, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 4., None, None, True, True, False),
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 4., None, None, False, True, True), # fused_conv
('invertedbottleneck', 3, 2, 120, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 120, 'relu6',
None, 4., None, None, True, True, False),
('invertedbottleneck', 3, 1, 120, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 3, 1, 120, 'relu6',
None, 8., None, None, True, True, False),
('invertedbottleneck', 3, 1, 144, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 144, 'relu6',
None, 8., None, None, True, True, False),
('invertedbottleneck', 3, 1, 144, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 3, 1, 144, 'relu6',
None, 8., None, None, True, True, True),
('invertedbottleneck', 3, 2, 160, 'relu6',
None, 4, None, None, True, False, False),
('invertedbottleneck', 3, 1, 160, 'relu6',
None, 4, None, None, True, True, False),
('invertedbottleneck', 3, 1, 160, 'relu6',
None, 4., None, None, False, False, False), # fused_conv
('tucker', 3, 1, 160, 'relu6',
None, None, 0.75, 0.75, None, True, False),
('invertedbottleneck', 3, 1, 240, 'relu6',
None, 8, None, None, True, False, True),
]
}
MD_EdgeTPU_BLOCK_SPECS = {
'spec_name': 'MobileDetEdgeTPU',
# [use_depthwise] is set to False for fused_conv
# [se_ratio] is set to None for all inverted_bottleneck layers
# [activation] is set to 'relu6' for all applicable layers
'block_spec_schema': ['block_fn', 'kernel_size', 'strides', 'filters',
'activation', 'se_ratio', 'expand_ratio',
'input_compression_ratio', 'output_compression_ratio',
'use_depthwise', 'use_residual', 'is_output'],
'block_specs': [
('convbn', 3, 2, 32, 'relu6',
None, None, None, None, None, None, False),
('tucker', 3, 1, 16, 'relu6',
None, None, 0.25, 0.75, None, False, True),
('invertedbottleneck', 3, 2, 16, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 16, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 16, 'relu6',
None, 8., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 16, 'relu6',
None, 4., None, None, False, True, True), # fused_conv
('invertedbottleneck', 5, 2, 40, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 40, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 40, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 40, 'relu6',
None, 4., None, None, False, True, True), # fused_conv
('invertedbottleneck', 3, 2, 72, 'relu6',
None, 8, None, None, True, False, False),
('invertedbottleneck', 3, 1, 72, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 3, 1, 72, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 72, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 5, 1, 96, 'relu6',
None, 8, None, None, True, False, False),
('invertedbottleneck', 5, 1, 96, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 3, 1, 96, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 3, 1, 96, 'relu6',
None, 8, None, None, True, True, True),
('invertedbottleneck', 5, 2, 120, 'relu6',
None, 8, None, None, True, False, False),
('invertedbottleneck', 3, 1, 120, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 5, 1, 120, 'relu6',
None, 4, None, None, True, True, False),
('invertedbottleneck', 3, 1, 120, 'relu6',
None, 8, None, None, True, True, False),
('invertedbottleneck', 5, 1, 384, 'relu6',
None, 8, None, None, True, False, True),
]
}
MD_GPU_BLOCK_SPECS = {
'spec_name': 'MobileDetGPU',
# [use_depthwise] is set to False for fused_conv
# [se_ratio] is set to None for all inverted_bottleneck layers
# [activation] is set to 'relu6' for all applicable layers
'block_spec_schema': ['block_fn', 'kernel_size', 'strides', 'filters',
'activation', 'se_ratio', 'expand_ratio',
'input_compression_ratio', 'output_compression_ratio',
'use_depthwise', 'use_residual', 'is_output'],
'block_specs': [
# block 0
('convbn', 3, 2, 32, 'relu6',
None, None, None, None, None, None, False),
# block 1
('tucker', 3, 1, 16, 'relu6',
None, None, 0.25, 0.25, None, False, True),
# block 2
('invertedbottleneck', 3, 2, 32, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('tucker', 3, 1, 32, 'relu6',
None, None, 0.25, 0.25, None, True, False),
('tucker', 3, 1, 32, 'relu6',
None, None, 0.25, 0.25, None, True, False),
('tucker', 3, 1, 32, 'relu6',
None, None, 0.25, 0.25, None, True, True),
# block 3
('invertedbottleneck', 3, 2, 64, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 8., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 8., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 64, 'relu6',
None, 4., None, None, False, True, True), # fused_conv
# block 4
('invertedbottleneck', 3, 2, 128, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
# block 5
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 8., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 8., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 8., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 8., None, None, False, True, True), # fused_conv
# block 6
('invertedbottleneck', 3, 2, 128, 'relu6',
None, 4., None, None, False, False, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
('invertedbottleneck', 3, 1, 128, 'relu6',
None, 4., None, None, False, True, False), # fused_conv
# block 7
('invertedbottleneck', 3, 1, 384, 'relu6',
None, 8, None, None, True, False, True),
]
}
SUPPORTED_SPECS_MAP = {
'MobileDetCPU': MD_CPU_BLOCK_SPECS,
'MobileDetDSP': MD_DSP_BLOCK_SPECS,
'MobileDetEdgeTPU': MD_EdgeTPU_BLOCK_SPECS,
'MobileDetGPU': MD_GPU_BLOCK_SPECS,
}
@dataclasses.dataclass
class BlockSpec(hyperparams.Config):
"""A container class that specifies the block configuration for MobileDet."""
block_fn: str = 'convbn'
kernel_size: int = 3
strides: int = 1
filters: int = 32
use_bias: bool = False
use_normalization: bool = True
activation: str = 'relu6'
is_output: bool = True
# Used for block type InvertedResConv and TuckerConvBlock.
use_residual: bool = True
# Used for block type InvertedResConv only.
use_depthwise: bool = True
expand_ratio: Optional[float] = 8.
se_ratio: Optional[float] = None
# Used for block type TuckerConvBlock only.
input_compression_ratio: Optional[float] = None
output_compression_ratio: Optional[float] = None
def block_spec_decoder(
specs: Dict[Any, Any],
filter_size_scale: float,
divisible_by: int = 8) -> List[BlockSpec]:
"""Decodes specs for a block.
Args:
specs: A `dict` specification of block specs of a mobiledet version.
filter_size_scale: A `float` multiplier for the filter size for all
convolution ops. The value must be greater than zero. Typical usage will
be to set this value in (0, 1) to reduce the number of parameters or
computation cost of the model.
divisible_by: An `int` that ensures all inner dimensions are divisible by
this number.
Returns:
A list of `BlockSpec` that defines structure of the base network.
"""
spec_name = specs['spec_name']
block_spec_schema = specs['block_spec_schema']
block_specs = specs['block_specs']
if not block_specs:
raise ValueError(
'The block spec cannot be empty for {} !'.format(spec_name))
if len(block_specs[0]) != len(block_spec_schema):
raise ValueError('The block spec values {} do not match with '
'the schema {}'.format(block_specs[0], block_spec_schema))
decoded_specs = []
for s in block_specs:
kw_s = dict(zip(block_spec_schema, s))
decoded_specs.append(BlockSpec(**kw_s))
for ds in decoded_specs:
if ds.filters:
ds.filters = nn_layers.round_filters(filters=ds.filters,
multiplier=filter_size_scale,
divisor=divisible_by,
round_down_protect=False,
min_depth=8)
return decoded_specs
@tf_keras.utils.register_keras_serializable(package='Vision')
class MobileDet(tf_keras.Model):
"""Creates a MobileDet family model."""
def __init__(
self,
model_id: str = 'MobileDetCPU',
filter_size_scale: float = 1.0,
input_specs: tf_keras.layers.InputSpec = layers.InputSpec(
shape=[None, None, None, 3]),
# The followings are for hyper-parameter tuning.
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,
# The followings should be kept the same most of the times.
min_depth: int = 8,
divisible_by: int = 8,
regularize_depthwise: bool = False,
use_sync_bn: bool = False,
**kwargs):
"""Initializes a MobileDet model.
Args:
model_id: A `str` of MobileDet version. The supported values are
`MobileDetCPU`, `MobileDetDSP`, `MobileDetEdgeTPU`, `MobileDetGPU`.
filter_size_scale: A `float` of multiplier for the filters (number of
channels) for all convolution ops. The value must be greater than zero.
Typical usage will be to set this value in (0, 1) to reduce the number
of parameters or computation cost of the model.
input_specs: A `tf_keras.layers.InputSpec` of specs of the input tensor.
norm_momentum: A `float` of normalization momentum for the moving average.
norm_epsilon: A `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.
min_depth: An `int` of minimum depth (number of channels) for all
convolution ops. Enforced when filter_size_scale < 1, and not an active
constraint when filter_size_scale >= 1.
divisible_by: An `int` that ensures all inner dimensions are divisible by
this number.
regularize_depthwise: If Ture, apply regularization on depthwise.
use_sync_bn: If True, use synchronized batch normalization.
**kwargs: Additional keyword arguments to be passed.
"""
if model_id not in SUPPORTED_SPECS_MAP:
raise ValueError('The MobileDet version {} '
'is not supported'.format(model_id))
if filter_size_scale <= 0:
raise ValueError('filter_size_scale is not greater than zero.')
self._model_id = model_id
self._input_specs = input_specs
self._filter_size_scale = filter_size_scale
self._min_depth = min_depth
self._divisible_by = divisible_by
self._regularize_depthwise = regularize_depthwise
self._kernel_initializer = kernel_initializer
self._kernel_regularizer = kernel_regularizer
self._bias_regularizer = bias_regularizer
self._use_sync_bn = use_sync_bn
self._norm_momentum = norm_momentum
self._norm_epsilon = norm_epsilon
inputs = tf_keras.Input(shape=input_specs.shape[1:])
block_specs = SUPPORTED_SPECS_MAP.get(model_id)
self._decoded_specs = block_spec_decoder(
specs=block_specs,
filter_size_scale=self._filter_size_scale,
divisible_by=self._get_divisible_by())
x, endpoints, next_endpoint_level = self._mobiledet_base(inputs=inputs)
self._output_specs = {l: endpoints[l].get_shape() for l in endpoints}
super(MobileDet, self).__init__(
inputs=inputs, outputs=endpoints, **kwargs)
def _get_divisible_by(self):
return self._divisible_by
def _mobiledet_base(self,
inputs: tf.Tensor
) -> Tuple[tf.Tensor, Dict[str, tf.Tensor], int]:
"""Builds the base MobileDet architecture.
Args:
inputs: A `tf.Tensor` of shape `[batch_size, height, width, channels]`.
Returns:
A tuple of output Tensor and dictionary that collects endpoints.
"""
input_shape = inputs.get_shape().as_list()
if len(input_shape) != 4:
raise ValueError('Expected rank 4 input, was: %d' % len(input_shape))
net = inputs
endpoints = {}
endpoint_level = 1
for i, block_def in enumerate(self._decoded_specs):
block_name = 'block_group_{}_{}'.format(block_def.block_fn, i)
if block_def.block_fn == 'convbn':
net = mobilenet.Conv2DBNBlock(
filters=block_def.filters,
kernel_size=block_def.kernel_size,
strides=block_def.strides,
activation=block_def.activation,
use_bias=block_def.use_bias,
use_normalization=block_def.use_normalization,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
use_sync_bn=self._use_sync_bn,
norm_momentum=self._norm_momentum,
norm_epsilon=self._norm_epsilon
)(net)
elif block_def.block_fn == 'invertedbottleneck':
in_filters = net.shape.as_list()[-1]
net = nn_blocks.InvertedBottleneckBlock(
in_filters=in_filters,
out_filters=block_def.filters,
kernel_size=block_def.kernel_size,
strides=block_def.strides,
expand_ratio=block_def.expand_ratio,
se_ratio=block_def.se_ratio,
se_inner_activation=block_def.activation,
se_gating_activation='sigmoid',
se_round_down_protect=False,
expand_se_in_filters=True,
activation=block_def.activation,
use_depthwise=block_def.use_depthwise,
use_residual=block_def.use_residual,
regularize_depthwise=self._regularize_depthwise,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
use_sync_bn=self._use_sync_bn,
norm_momentum=self._norm_momentum,
norm_epsilon=self._norm_epsilon,
divisible_by=self._get_divisible_by()
)(net)
elif block_def.block_fn == 'tucker':
in_filters = net.shape.as_list()[-1]
net = nn_blocks.TuckerConvBlock(
in_filters=in_filters,
out_filters=block_def.filters,
kernel_size=block_def.kernel_size,
strides=block_def.strides,
input_compression_ratio=block_def.input_compression_ratio,
output_compression_ratio=block_def.output_compression_ratio,
activation=block_def.activation,
use_residual=block_def.use_residual,
kernel_initializer=self._kernel_initializer,
kernel_regularizer=self._kernel_regularizer,
bias_regularizer=self._bias_regularizer,
use_sync_bn=self._use_sync_bn,
norm_momentum=self._norm_momentum,
norm_epsilon=self._norm_epsilon,
divisible_by=self._get_divisible_by()
)(net)
else:
raise ValueError('Unknown block type {} for layer {}'.format(
block_def.block_fn, i))
net = tf_keras.layers.Activation('linear', name=block_name)(net)
if block_def.is_output:
endpoints[str(endpoint_level)] = net
endpoint_level += 1
return net, endpoints, endpoint_level
def get_config(self):
config_dict = {
'model_id': self._model_id,
'filter_size_scale': self._filter_size_scale,
'min_depth': self._min_depth,
'divisible_by': self._divisible_by,
'regularize_depthwise': self._regularize_depthwise,
'kernel_initializer': self._kernel_initializer,
'kernel_regularizer': self._kernel_regularizer,
'bias_regularizer': self._bias_regularizer,
'use_sync_bn': self._use_sync_bn,
'norm_momentum': self._norm_momentum,
'norm_epsilon': self._norm_epsilon,
}
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
@factory.register_backbone_builder('mobiledet')
def build_mobiledet(
input_specs: tf_keras.layers.InputSpec,
backbone_config: hyperparams.Config,
norm_activation_config: hyperparams.Config,
l2_regularizer: Optional[tf_keras.regularizers.Regularizer] = None
) -> tf_keras.Model:
"""Builds MobileDet backbone from a config."""
backbone_type = backbone_config.type
backbone_cfg = backbone_config.get()
assert backbone_type == 'mobiledet', (f'Inconsistent backbone type '
f'{backbone_type}')
return MobileDet(
model_id=backbone_cfg.model_id,
filter_size_scale=backbone_cfg.filter_size_scale,
input_specs=input_specs,
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)