def _build_layer(params,
layer_spec,
input_filters,
output_filters,
filters_base):
"""Create a one of N possible types of layers within the body of a network."""
if isinstance(layer_spec, mobile_search_space_v3.ResidualSpec):
return _build_residual_spec(params, layer_spec, input_filters,
output_filters, filters_base)
elif isinstance(layer_spec, mobile_search_space_v3.ConvSpec):
return _build_conv(params, layer_spec, input_filters, output_filters)
elif isinstance(layer_spec, mobile_search_space_v3.SeparableConvSpec):
return _build_separable_conv(params, layer_spec, input_filters,
output_filters)
elif isinstance(layer_spec, mobile_search_space_v3.DepthwiseBottleneckSpec):
return _build_depthwise_bottleneck(params, layer_spec, input_filters,
output_filters, filters_base)
elif isinstance(layer_spec, mobile_search_space_v3.GlobalAveragePoolSpec):
return layers.GlobalAveragePool(keepdims=True)
elif isinstance(layer_spec, mobile_search_space_v3.ActivationSpec):
return _get_activation(layer_spec)
elif isinstance(layer_spec, basic_specs.ZeroSpec):
return layers.Zeros()
elif isinstance(layer_spec, schema.OneOf):
return _build_oneof(
params, layer_spec, input_filters, output_filters, filters_base)
else:
raise ValueError('Unsupported layer_spec type: {}'.format(
type(layer_spec)))
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 _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,
])
])