def _depthwise_conv_with_fixed_kernel(
params,
filters_or_mask,
kernel_size,
strides = (1, 1),
activation = None):
"""Build a depthwise conv, batch norm, and optional activation."""
result = []
if isinstance(filters_or_mask, int):
result.append(
layers.DepthwiseConv2D(
kernel_size=kernel_size,
depthwise_initializer=params['depthwise_initializer'],
depthwise_regularizer=params['kernel_regularizer'],
strides=strides))
else:
result.append(
layers.MaskedDepthwiseConv2D(
kernel_size=kernel_size,
mask=filters_or_mask,
depthwise_initializer=params['depthwise_initializer'],
depthwise_regularizer=params['kernel_regularizer'],
strides=strides))
result.append(_batch_norm(params, filters_or_mask))
if activation is not None:
result.append(activation)
return layers.Sequential(result)
def _build_model(params,
model_spec):
"""Translate a ConvTowerSpec namedtuple into a rematlib Layer."""
input_filters = schema.OneOf([3], basic_specs.FILTERS_TAG)
layer = None
result = []
endpoints = []
for block_spec in model_spec.blocks:
for layer_spec in block_spec.layers:
if isinstance(layer_spec, mobile_search_space_v3.DetectionEndpointSpec):
if layer is None:
raise ValueError(
'The first layer of the network cannot be a detection endpoint.')
endpoints.append(layer)
else:
output_filters = block_spec.filters
if isinstance(output_filters, int):
output_filters = schema.OneOf(
[output_filters], basic_specs.FILTERS_TAG)
layer = _build_layer(
params, layer_spec, input_filters, output_filters,
model_spec.filters_base)
input_filters = output_filters
result.append(layer)
result.append(_make_head(params))
# Build the model
model = layers.Sequential(result, aux_outputs=endpoints)
return model
def _conv_with_fixed_kernel(params,
input_filters_or_mask,
output_filters_or_mask,
kernel_size,
strides = (1, 1),
activation = None,
use_batch_norm = True):
"""Construct a Conv2D, followed by a batch norm and optional activation."""
result = []
if (isinstance(input_filters_or_mask, int) and
isinstance(output_filters_or_mask, int)):
result.append(
layers.Conv2D(
filters=output_filters_or_mask,
kernel_size=kernel_size,
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=not use_batch_norm,
strides=strides))
else:
result.append(
layers.MaskedConv2D(
input_mask=_to_filters_mask(input_filters_or_mask),
output_mask=_to_filters_mask(output_filters_or_mask),
kernel_size=kernel_size,
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=not use_batch_norm,
strides=strides))
if use_batch_norm:
result.append(_batch_norm(params, output_filters_or_mask))
if activation is not None:
result.append(activation)
return layers.Sequential(result)
def _make_head(params):
"""Construct a classification model head."""
result = []
if params['dropout_rate'] > 0:
result.append(layers.Dropout(params['dropout_rate']))
result.append(
layers.Conv2D(filters=params['num_classes'],
kernel_size=(1, 1),
kernel_initializer=params['dense_initializer'],
# kernel_regularizer is not used for the final dense layer:
kernel_regularizer=None,
use_bias=True))
result.append(layers.GlobalAveragePool(keepdims=False))
return layers.Sequential(result)
def _build_separable_conv(params, layer_spec, input_filters, output_filters):
return layers.Sequential([
_depthwise_conv(
params,
filters_or_mask=_oneof_filters_to_int_or_mask(input_filters),
kernel_size=layer_spec.kernel_size,
strides=layer_spec.strides,
activation=_get_activation(layer_spec.activation)),
_conv(
params,
input_filters_or_mask=_oneof_filters_to_int_or_mask(input_filters),
output_filters_or_mask=_oneof_filters_to_int_or_mask(
output_filters),
kernel_size=(1, 1),
activation=None),
])
def _build_depthwise_bottleneck(
params,
layer_spec,
input_filters,
output_filters,
filters_base):
"""Construct a bottleneck layer with a depthwise conv in the middle."""
input_filters_or_mask = _oneof_filters_to_int_or_mask(input_filters)
output_filters_or_mask = _oneof_filters_to_int_or_mask(output_filters)
expansion_filters_or_mask = _oneof_filters_to_int_or_mask(
layer_spec.expansion_filters, input_filters_or_mask, filters_base)
result = [
_conv(
params,
input_filters_or_mask=input_filters_or_mask,
output_filters_or_mask=expansion_filters_or_mask,
kernel_size=(1, 1),
activation=_get_activation(layer_spec.activation)),
_depthwise_conv(
params,
filters_or_mask=expansion_filters_or_mask,
kernel_size=layer_spec.kernel_size,
strides=layer_spec.strides,
activation=_get_activation(layer_spec.activation)),
]
result.append(
_maybe_squeeze_and_excite(
params,
expansion_filters_or_mask,
_get_activation(layer_spec.se_inner_activation),
_get_activation(layer_spec.se_gating_activation),
layer_spec.use_squeeze_and_excite))
result.append(
_conv(
params,
input_filters_or_mask=expansion_filters_or_mask,
output_filters_or_mask=output_filters_or_mask,
kernel_size=(1, 1),
activation=None))
return layers.Sequential(result)
def _squeeze_and_excite(params,
input_filters_or_mask,
inner_activation,
gating_activation):
"""Generate a squeeze-and-excite layer."""
# We provide two code paths:
# 1. For the case where the number of input filters is known at graph
# construction time, and input_filters_or_mask is an int. This typically
# happens during stand-alone model training.
# 2. For the case where the number of input filters is not known until
# runtime, and input_filters_or_mask is a 1D float tensor. This often
# happens during an architecture search.
if isinstance(input_filters_or_mask, int):
input_filters = input_filters_or_mask
hidden_filters = search_space_utils.make_divisible(
input_filters * _SQUEEZE_AND_EXCITE_RATIO,
divisor=params['filters_base'])
return layers.ParallelProduct([
layers.Identity(),
layers.Sequential([
layers.GlobalAveragePool(keepdims=True),
layers.Conv2D(
filters=hidden_filters,
kernel_size=(1, 1),
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=True),
inner_activation,
layers.Conv2D(
filters=input_filters,
kernel_size=(1, 1),
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=True),
gating_activation,
]),
])
else:
input_mask = input_filters_or_mask
input_filters = tf.reduce_sum(input_mask)
hidden_filters = search_space_utils.tf_make_divisible(
input_filters * _SQUEEZE_AND_EXCITE_RATIO,
divisor=params['filters_base'])
max_input_filters = int(input_mask.shape[0])
max_hidden_filters = search_space_utils.make_divisible(
max_input_filters * _SQUEEZE_AND_EXCITE_RATIO,
divisor=params['filters_base'])
hidden_mask = tf.sequence_mask(
hidden_filters, max_hidden_filters, dtype=tf.float32)
return layers.ParallelProduct([
layers.Identity(),
layers.Sequential([
layers.GlobalAveragePool(keepdims=True),
layers.MaskedConv2D(
input_mask=input_mask,
output_mask=hidden_mask,
kernel_size=(1, 1),
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=True),
inner_activation,
layers.MaskedConv2D(
input_mask=hidden_mask,
output_mask=input_mask,
kernel_size=(1, 1),
kernel_initializer=params['kernel_initializer'],
kernel_regularizer=params['kernel_regularizer'],
use_bias=True),
gating_activation,
])
])
