mobile_models.py

# Copyright 2021 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.

"""Mobile model builder."""

from typing import List, Optional, Sequence, Text

import pyglove as pg
from pyglove.tensorflow import keras
from pyglove.tensorflow import selections

from official.projects.tunas.modeling.layers import nn_blocks


class _MobileModel(keras.Model):
  """Mobile model."""

  def __init__(
      self,
      stem_conv_filters: nn_blocks.Filters,
      blocks: Sequence[keras.layers.Layer],
      feature_conv_filters: nn_blocks.Filters,
      kernel_initializer=keras.initializers.he_normal(),
      dense_initializer=keras.initializers.random_normal(stddev=0.01),
      # NOTE(daiyip): Keras L2 implementation is 2x of
      # tf.contrib.keras.layers.l2_regularizer.
      kernel_regularizer=keras.regularizers.l2(4e-5 * 0.5),
      normalization=keras.layers.BatchNormalization(
          momentum=0.99, epsilon=0.001),
      activation=keras.layers.ReLU(),
      dropout_rate: float = 0.,
      num_classes: Optional[int] = 1001,
      name: Optional[Text] = None):
    """Mobile model.

    Args:
      stem_conv_filters: Filter size for the stem conv unit.
      blocks: A list of layers as residual blocks after the stem layer.
      feature_conv_filters: Number of penultimate features.
      kernel_initializer: Kernel initializer used for the stem and featurizer.
      dense_initializer: Kernel initializer used for the classification layer.
      kernel_regularizer: Regularizer for the layers in the network.
      normalization: Normalization layer used in the network.
      activation: Activation layer used in the network.
      dropout_rate: Dropout rate for the penultimate features, applicable only
        when `num_classes` is not None.
      num_classes: Number of classes for the classification model. If None,
        the classification layer will be excluded.
      name: Name of the model.

    Returns:
      A list of tensors as model outputs.
        If `num_classes` is not None, the list is [logits, penultimate_features]
          plus lower-level features.
        Otherwise the list is [penultimate_features] plus lower-level features.
    """
    super().__init__(name=name)
    self._stem = nn_blocks.conv2d(
        filters=stem_conv_filters,
        kernel_size=(3, 3),
        strides=(2, 2),
        padding='same',
        kernel_initializer=kernel_initializer,
        kernel_regularizer=kernel_regularizer,
        use_bias=False,
        normalization=normalization,
        activation=activation,
        name='stem')

    # An ugly hack to track each layer from the `blocks`, since Keras does not
    # track `tf.keras.layers.Layer` objects from container-type members.
    for i, block in enumerate(blocks):
      setattr(self, '_blocks_{:d}'.format(i), block)
    self._blocks = blocks

    self._featurizer = nn_blocks.conv2d(
        filters=feature_conv_filters,
        kernel_size=(1, 1),
        padding='same',
        kernel_initializer=kernel_initializer,
        kernel_regularizer=kernel_regularizer,
        use_bias=False,
        normalization=normalization,
        activation=activation,
        name='features')

    self._global_pooling = keras.layers.GlobalAveragePooling2D()
    if num_classes is not None:
      self._dropout = keras.layers.Dropout(dropout_rate)
      self._classification_head = keras.layers.Dense(
          num_classes,
          kernel_initializer=dense_initializer,
          use_bias=True,
          name='classification_head')
    self.num_classes = num_classes

  def call(self, inputs):
    x = self._stem(inputs)
    lower_level_features = []
    for block in self._blocks:
      x = block(x)
      if isinstance(x, (list, tuple)):
        x, features = x[0], x[1:]
        lower_level_features.extend(list(features))

    x = self._featurizer(x)
    x = self._global_pooling(x)
    penultimate_features = x
    if self.num_classes is None:
      return [penultimate_features] + lower_level_features
    else:
      x = self._dropout(x)
      x = self._classification_head(x)
      return [x, penultimate_features] + lower_level_features


MobileModel = pg.symbolize(_MobileModel, class_name='MobileModel')


class _MobileBlock(keras.Model):
  """Mobile block."""

  def __init__(
      self,
      # We use List instead of Sequence here since Tuple cannot be
      # modified using rebind.
      sublayers: List[keras.layers.Layer],
      filters: nn_blocks.Filters,
      name: Optional[Text] = None):
    """Mobile block.

    Args:
      sublayers: Sublayers for the block.
      filters: Number of filters for the block. All sublayers will be using
        this filters.
      name: Name of the block.

    Returns:
      A tuple of 2 tensors (block output, features-before-downsampling)
    """
    super().__init__(name=name)
    self.sublayers = [s.clone(override={'filters': filters}) for s in sublayers]
    self.filters = filters

  def call(self, inputs):
    get_image_size = lambda x: (int(x.shape[1]), int(x.shape[2]))
    x = inputs
    image_size = get_image_size(x)
    features = []
    for layer in self.sublayers:
      x = layer(x)
      new_image_size = get_image_size(x)
      if new_image_size != image_size:
        features.append(x)
        image_size = new_image_size
    return tuple([x] + features)


MobileBlock = pg.symbolize(_MobileBlock, class_name='MobileBlock')


def search_model_v2(
    init_filters: Sequence[int],
    filters_multipliers: Optional[Sequence[float]] = None,
    filters_scale_factor: float = 1.0,
    filters_base: int = 8,
    se_ratios: Optional[List[float]] = None,
    num_classes: Optional[int] = 1001,
    normalization=keras.layers.BatchNormalization(momentum=0.0, epsilon=0.001),
    activation=keras.layers.ReLU(),
    dropout_rate: float = 0.,
    kernel_initializer=keras.initializers.he_normal(),
    depthwise_initializer=keras.initializers.depthwise_he_normal(),
    dense_initializer=keras.initializers.random_normal(stddev=0.01),
    # NOTE(daiyip): Keras L2 implementation is 2x of
    # tf.contrib.keras.layers.l2_regularizer.
    kernel_regularizer=keras.regularizers.l2(4e-5 * 0.5),
    name: Optional[Text] = 'search_mobile_model_v2'):
  """A searchable model derived from MobileNetV2.

  Args:
    init_filters: A list of integers (size=9) as the initial filter size of
      each mobile block.
    filters_multipliers: An optional list of floats as multipliers for the
      filters. If the list size is larger than 1, it is a search space including
      searching the filter sizes per block.
    filters_scale_factor: Additional scaling factor on top of
      filters_multipliers, this is to align with existing TuNAS codebase.
    filters_base: An integer as base to compute multiplied filters.
      Please see `layers.scale_filters` for details.
    se_ratios: Squeeze-and-excite ratios. If empty, SE is not used.
    num_classes: Number of classes for the classification model. If None,
      the classification layer will be excluded.
    normalization: Normalization layer used in the model.
    activation: Activation layer used in the model.
    dropout_rate: Dropout rate for the penultimate features, applicable only
      when `num_classes` is not None.
    kernel_initializer: Kernel initializer used for the Conv2D units
      in the model.
    depthwise_initializer: Kernel initializer used for DepthwiseConv2D units
      in the model.
    dense_initializer: Kernel initializer used for the classification layer.
    kernel_regularizer: Regularizer for the layers in the network.
    name: Name of the model, which will be used as the top name scope.

  Returns:
    A `MobileModel` object (a tf.keras.Model subclass) as the search model.
  """
  if not isinstance(init_filters, (tuple, list)) or len(init_filters) != 9:
    raise ValueError(
        '`init_filters` must be a sequence of 9. '
        'Encountered: %r.' % init_filters)

  se_ratios = [None] + (se_ratios if se_ratios else [])

  def _filters(x):
    filters = keras.layers.get_filters(x, filters_multipliers, filters_base)
    if filters_scale_factor != 1.0:
      # Up to now, filters will contain non-duplicated values. We will then
      # apply an additional filters scaling based on the candidates.
      # Please be aware of that this round of rescaling may result in duplicated
      # entries. We do not dedup these values to be compatible with original
      # TuNAS implementation.
      filters = keras.layers.maybe_oneof([
          keras.layers.get_filters(x, [filters_scale_factor], filters_base)
          for x in filters.candidates
      ], choice_type=keras.layers.ChoiceType.FILTERS)
    return filters

  def _mobile_layer(layer_index, strides, kernel_size,
                    expansion_factor, skippable=True):
    # Note(luoshixin): collapsed search space is not supported currently.
    candidates = []
    for i, (sr, f, k) in enumerate(
        selections.map_candidates([se_ratios, expansion_factor, kernel_size])):
      candidates.append(nn_blocks.inverted_bottleneck_with_se(
          # Placeholder, which will be modified at mobile_block level.
          filters=1,
          strides=strides,
          kernel_size=k,
          expansion_factor=f,
          se_ratio=sr,
          normalization=normalization,
          activation=activation,
          kernel_initializer=kernel_initializer,
          kernel_regularizer=kernel_regularizer,
          depthwise_initializer=depthwise_initializer,
          depthwise_regularizer=kernel_regularizer,
          name='inverted_bottleneck%d' % (i + 1)))

    if skippable:
      candidates.append(keras.layers.zeros())
    op = keras.layers.maybe_oneof(
        candidates,
        name=('switch' if skippable else ('switch%d' % (layer_index + 1))))
    if skippable:
      op = keras.layers.Residual(op, name='residual%d' % (layer_index + 1))
    return op

  # pylint: disable=unexpected-keyword-arg
  blocks = [
      MobileBlock([
          _mobile_layer(0, (1, 1), [(3, 3), (5, 5), (7, 7)], [1], False),
      ], _filters(init_filters[1]), name='block1'),
      MobileBlock([
          _mobile_layer(0, (2, 2), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
          _mobile_layer(1, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(2, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(3, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
      ], _filters(init_filters[2]), name='block2'),
      MobileBlock([
          _mobile_layer(0, (2, 2), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
          _mobile_layer(1, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(2, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(3, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
      ], _filters(init_filters[3]), name='block3'),
      MobileBlock([
          _mobile_layer(0, (2, 2), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
          _mobile_layer(1, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(2, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(3, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
      ], _filters(init_filters[4]), name='block4'),
      MobileBlock([
          _mobile_layer(0, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
          _mobile_layer(1, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(2, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(3, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
      ], _filters(init_filters[5]), name='block5'),
      MobileBlock([
          _mobile_layer(0, (2, 2), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
          _mobile_layer(1, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(2, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
          _mobile_layer(3, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6]),
      ], _filters(init_filters[6]), name='block6'),
      MobileBlock([
          _mobile_layer(0, (1, 1), [(3, 3), (5, 5), (7, 7)], [3, 6], False),
      ], _filters(init_filters[7]), name='block7'),
  ]
  return MobileModel(   # pylint: disable=unexpected-keyword-arg
      stem_conv_filters=_filters(init_filters[0]),
      blocks=blocks,
      feature_conv_filters=_filters(init_filters[8]),
      kernel_initializer=kernel_initializer,
      dense_initializer=dense_initializer,
      kernel_regularizer=kernel_regularizer,
      normalization=normalization,
      activation=activation,
      dropout_rate=dropout_rate,
      num_classes=num_classes,
      name=name)


def static_model(
    search_model: MobileModel,
    dna: pg.DNA,
    use_stateful_batch_norm: bool = True) -> MobileModel:
  """Returns a static model from a search model and a DNA."""
  model = pg.template(search_model).decode(dna)
  return pg.patch_on_member(
      model,
      keras.layers.BatchNormalization, 'momentum',
      0.99 if use_stateful_batch_norm else 0.0)


def static_mobile_model(
    op_indices: Sequence[int],
    init_filters: Sequence[int],
    num_classes: int,
    weight_decay: float,
    dropout: float = 0.0,
    filters_multiplier: float = 1.0,
    name: Optional[Text] = None) -> MobileModel:
  """Create static mobile model."""
  assert len(op_indices) == 22
  # NOTE(daiyip): Regularizer value of Keras L2 implementation is 2x of
  # tf.contrib.keras.layers.l2_regularizer.
  search_model = search_model_v2(
      init_filters,
      num_classes=num_classes,
      filters_multipliers=[filters_multiplier],
      dropout_rate=dropout,
      kernel_regularizer=keras.regularizers.l2(weight_decay * 0.5),
      name=name)
  return static_model(search_model, pg.DNA.parse(list(op_indices)))


MOBILENET_V2_FILTERS = (32, 16, 24, 32, 64, 96, 160, 320, 1280)
MNASNET_FILTERS = (32, 16, 24, 40, 80, 96, 192, 320, 1280)
PROXYLESSNAS_GPU_FILTERS = (40, 24, 32, 56, 112, 128, 256, 432, 1728)
PROXYLESSNAS_CPU_FILTERS = (40, 24, 32, 48, 88, 104, 216, 360, 1432)
PROXYLESSNAS_MOBILE_FILTERS = (32, 16, 32, 40, 80, 96, 192, 320, 1280)
MOBILEDET_EDGE_TPU_FILTERS = (32, 16, 32, 48, 96, 96, 160, 192, 192)

MOBILENET_V2_OPERATIONS = (0, 3, 3, 6, 6, 3, 3, 3, 6, 3, 3,
                           3, 3, 3, 3, 3, 6, 3, 3, 3, 6, 3)
MNASNET_OPERATIONS = (0, 0, 0, 0, 6, 1, 1, 1, 6, 4, 4,
                      4, 6, 3, 3, 6, 6, 4, 4, 4, 4, 3)
PROXYLESSNAS_GPU_OPERATIONS = (0, 1, 6, 6, 6, 2, 6, 6, 0, 5, 6,
                               6, 1, 4, 6, 0, 1, 5, 5, 5, 4, 5)
PROXYLESSNAS_CPU_OPERATIONS = (0, 3, 0, 0, 0, 3, 0, 0, 1, 3, 6,
                               6, 0, 4, 0, 0, 0, 4, 1, 1, 0, 4)
PROXYLESSNAS_MOBILE_OPERATIONS = (0, 1, 0, 6, 6, 2, 0, 1, 1, 5, 1,
                                  1, 1, 4, 1, 1, 1, 5, 5, 2, 2, 5)

MOBILE_DEFAULT_FILTER_MULTIPLIERS = (0.5, 0.625, 0.75, 1.0, 1.25, 1.5, 2.0)
MOBILEDET_EDGE_TPU_FILTER_MULTIPLIERS = MOBILE_DEFAULT_FILTER_MULTIPLIERS


def mobilenet_v2(num_classes: Optional[int] = 1001,
                 weight_decay: float = 4e-5,
                 dropout: float = 0.0) -> MobileModel:
  """MobileNet v2."""
  return static_mobile_model(
      op_indices=MOBILENET_V2_OPERATIONS,
      init_filters=MOBILENET_V2_FILTERS,
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      name='mobilenet_v2')


def mnasnet(num_classes: Optional[int] = 1001,
            weight_decay: float = 4e-5,
            dropout: float = 0.0) -> MobileModel:
  """MNASNet."""
  return static_mobile_model(
      op_indices=MNASNET_OPERATIONS,
      init_filters=MNASNET_FILTERS,
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      name='mnasnet')


def proxyless_nas_gpu(num_classes: Optional[int] = 1001,
                      weight_decay: float = 4e-5,
                      dropout: float = 0.0) -> MobileModel:
  """ProxylessNAS searched for GPU."""
  return static_mobile_model(
      op_indices=PROXYLESSNAS_GPU_OPERATIONS,
      init_filters=PROXYLESSNAS_GPU_FILTERS,
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      name='proxyless_nas_gpu')


def proxyless_nas_cpu(num_classes: Optional[int] = 1001,
                      weight_decay: float = 4e-5,
                      dropout: float = 0.0) -> MobileModel:
  """ProxylessNAS searched for CPU."""
  return static_mobile_model(
      op_indices=PROXYLESSNAS_CPU_OPERATIONS,
      init_filters=PROXYLESSNAS_CPU_FILTERS,
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      name='proxyless_nas_cpu')


def proxyless_nas_mobile(num_classes: Optional[int] = 1001,
                         weight_decay: float = 4e-5,
                         dropout: float = 0.0) -> MobileModel:
  """ProxylessNAS searched for mobile device."""
  return static_mobile_model(
      op_indices=PROXYLESSNAS_MOBILE_OPERATIONS,
      init_filters=PROXYLESSNAS_MOBILE_FILTERS,
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      name='proxyless_nas_mobile')


# Search spaces from the TuNAS paper:
# pylint:disable=line-too-long
# Reference:
#  Gabriel Bender & Hanxiao Liu, et al. Can Weight Sharing Outperform Random Architecture Search? An Investigation With TuNAS
#  https://openaccess.thecvf.com/content_CVPR_2020/html/Bender_Can_Weight_Sharing_Outperform_Random_Architecture_Search_An_Investigation_With_CVPR_2020_paper.html
# pylint:enable=line-too-long


def proxylessnas_search(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """Original Proxyless NAS search space."""
  return search_model_v2(
      init_filters=PROXYLESSNAS_MOBILE_FILTERS,
      num_classes=num_classes,
      name='proxylessnas_search')


def proxylessnas_with_filters_doubled_every_block(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """A variant search space of `proxylessnas_search`.

  In this search space the number of filters is doubled in each consecutive
  block. This search space is a baseline in the Tunas paper to evaluate the
  effect of searching over filter sizes compared to traditional heuristics.

  Args:
    num_classes: Number of classes for the classification model. If None,
      the classification layer will be excluded.

  Returns:
    A `MobileModel` object (a tf.keras.Model subclass) as the search model.
  """
  return search_model_v2(
      init_filters=(16, 16, 16, 32, 64, 128, 256, 512, 1024),
      num_classes=num_classes,
      name='proxylessnas_with_filters_doubled_every_block_search')


def proxylessnas_with_filters_doubled_every_stride2(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """A variant search space of `proxylessnas_search`.

  This search space is an extension of the ProxylessNas search space where it is
  made possible to search over the output filter sizes.

  Args:
    num_classes: Number of classes for the classification model. If None,
      the classification layer will be excluded.

  Returns:
    A `MobileModel` object (a tf.keras.Model subclass) as the search model.
  """
  return search_model_v2(
      init_filters=(16, 16, 32, 64, 128, 128, 256, 256, 512),
      num_classes=num_classes,
      name='proxylessnas_with_filters_doubled_every_stride2_search')


def proxylessnas_outfilters_search(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """A variant search space of `proxylessnas_search`.

  This search space is an extension of the ProxylessNas search space where it is
  made possible to search over the output filter sizes.

  Args:
    num_classes: Number of classes for the classification model. If None,
      the classification layer will be excluded.

  Returns:
    A `MobileModel` object (a tf.keras.Model subclass) as the search model.
  """
  return search_model_v2(
      init_filters=(16, 16, 16, 32, 64, 128, 256, 512, 1024),
      filters_multipliers=(0.5, 0.625, 0.75, 1.0, 1.25, 1.5, 2.0),
      num_classes=num_classes,
      name='proxylessnas_outfilters_search')


def proxylessnas_with_mobilenet_v2_filters_search(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """Original Proxyless NAS search space."""
  return search_model_v2(
      init_filters=MOBILENET_V2_FILTERS,
      num_classes=num_classes,
      name='proxylessnas_with_mobilenet_v2_filters_search')


def mobilenet_v2_filters_search(
    num_classes: Optional[int] = 1001) -> MobileModel:
  """MobileNetV2 filters search."""
  search_model = search_model_v2(
      init_filters=MOBILENET_V2_FILTERS,
      filters_multipliers=MOBILE_DEFAULT_FILTER_MULTIPLIERS,
      num_classes=num_classes,
      name='mobilenet_v2_filters_search')
  def select_ops(x):
    return keras.layers.get_choice_type(x) == keras.layers.ChoiceType.OP
  return pg.materialize(
      search_model,
      pg.DNA.parse(list(MOBILENET_V2_OPERATIONS)),
      where=select_ops)


def tunas_search_model(ssd: Text) -> MobileModel:  # pytype: disable=invalid-annotation
  """Get TuNAS search model by search space name."""
  # Note(luoshixin): collapsed search space is not supported, and hence
  # not migrated currently.
  ssd_map = {
      'proxylessnas_search': proxylessnas_search,
      'proxylessnas_with_filters_doubled_every_block':
          proxylessnas_with_filters_doubled_every_block,
      'proxylessnas_with_filters_doubled_every_stride2':
          proxylessnas_with_filters_doubled_every_stride2,
      'proxylessnas_outfilters_search': proxylessnas_outfilters_search,
  }
  if ssd not in ssd_map:
    raise ValueError('Unsupported TuNAS search space %r.' % ssd)
  return ssd_map[ssd]()


def _swap_op_choices(
    model,
    initial_op_choices,
    body_op_choices):
  """Helper method to swap op choices in a MobileModel."""
  context = dict(initial_op=True)
  def swap_ops(key_path: pg.KeyPath, value, parent):
    del key_path
    # Skip static values and non-operation choices.
    if (not isinstance(value, pg.hyper.OneOf)
        or keras.layers.get_choice_type(value) != keras.layers.ChoiceType.OP):
      return value

    sample_ibn = value.candidates[0]
    assert isinstance(sample_ibn, nn_blocks.inverted_bottleneck), sample_ibn

    if context['initial_op']:
      candidates = pg.clone(initial_op_choices, deep=True)
      context['initial_op'] = False
    else:
      candidates = pg.clone(body_op_choices, deep=True)

    for c in candidates:
      keras.layers.inherit_hyperparameters_from(c, sample_ibn, [
          'filters', 'strides', 'kernel_initializer', 'kernel_regularizer',
          'depthwise_initializer', 'depthwise_regularizer',
          'normalization', 'activation'
      ])
    if isinstance(parent, keras.layers.Residual):
      candidates.append(keras.layers.zeros())
    return keras.layers.maybe_oneof(candidates)
  return model.rebind(swap_ops)


def mobiledet_edge_tpu_search(num_classes: Optional[int] = 1001,
                              weight_decay: float = 4e-5,
                              dropout: float = 0.0,
                              filters_scale_factor: float = 1.0,
                              filters_base: int = 8,
                              filters_multipliers: Sequence[float] = (
                                  MOBILEDET_EDGE_TPU_FILTER_MULTIPLIERS),
                              expansion_multipliers: Sequence[int] = (4, 8),
                              name: Text = 'mobiledet_edge_tpu_search'):
  """Return search model for MOBILEDET_EDGE_TPU search space from TuNAS."""
  def _op_choices(
      kernel_sizes,
      expansion_factors,
      tucker_kernel_sizes,
      tucker_input_ratios,
      tucker_output_ratios):
    ops = []
    # Add choices from regular inverted bottleneck.
    for i, (ef, ks) in enumerate(
        selections.map_candidates([expansion_factors, kernel_sizes])):
      ops.append(nn_blocks.inverted_bottleneck.partial(
          kernel_size=ks, expansion_factor=ef,
          name='inverted_bottleneck%d' % i))

    # Add choices from fused inverted bottleneck.
    for i, (ef, ks) in enumerate(selections.map_candidates(
        [expansion_factors, kernel_sizes])):
      ops.append(nn_blocks.fused_inverted_bottleneck.partial(
          kernel_size=ks, expansion_factor=ef,
          name='fused_inverted_bottleneck%d' % i))

    # Add choices from tucker bottleneck.
    for i, (iratio, ks, oratio) in enumerate(selections.map_candidates([
        tucker_input_ratios, tucker_kernel_sizes, tucker_output_ratios])):
      ops.append(nn_blocks.tucker_bottleneck.partial(
          kernel_size=ks, input_scale_ratio=iratio,
          output_scale_ratio=oratio, name='tucker_bottleneck%d' % i))
    return ops

  initial_op_choices = _op_choices(
      kernel_sizes=[(3, 3), (5, 5)],
      expansion_factors=[1],
      tucker_kernel_sizes=[(3, 3)],
      tucker_input_ratios=[0.25, 0.75],
      tucker_output_ratios=[0.25, 0.75])
  body_op_choices = _op_choices(
      kernel_sizes=[(3, 3), (5, 5)],
      expansion_factors=expansion_multipliers,
      tucker_kernel_sizes=[(3, 3)],
      tucker_input_ratios=[0.25, 0.75],
      tucker_output_ratios=[0.25, 0.75])
  search_model = search_model_v2(
      init_filters=MOBILEDET_EDGE_TPU_FILTERS,
      filters_multipliers=filters_multipliers,
      filters_scale_factor=filters_scale_factor,
      filters_base=filters_base,
      num_classes=num_classes,
      dropout_rate=dropout,
      kernel_regularizer=keras.regularizers.l2(weight_decay * 0.5),
      name=name)
  return _swap_op_choices(search_model, initial_op_choices, body_op_choices)


# This arch string is copied from: tunas/detection_search_space.py
DEFAULT_MOBILEDET_EDGE_TPU_ARCH_STRING = (
    '2:5:1:6:4:6:4:0:7:4:4:4:2:2:2:4:4:2:3:3:2:2:3:3:2:1:2:2:3:6:5')


def mobiledet_edge_tpu(
    arch_string: Text = DEFAULT_MOBILEDET_EDGE_TPU_ARCH_STRING,
    num_classes: Optional[int] = 1001,
    filters_scale_factor: float = 1.0,
    filters_base: int = 8,
    weight_decay: float = 4e-5,
    dropout: float = 0.0,
    filters_multipliers: Sequence[float] = (
        MOBILEDET_EDGE_TPU_FILTER_MULTIPLIERS),
    expansion_multipliers: Sequence[int] = (4, 8),
    name: Text = 'mobiledet_edge_tpu') -> MobileModel:
  """Static model from MobileDet Edge TPU search space."""
  dna = pg.DNA.parse([int(v) for v in arch_string.split(':')])
  search_model = mobiledet_edge_tpu_search(
      num_classes=num_classes,
      weight_decay=weight_decay,
      dropout=dropout,
      filters_scale_factor=filters_scale_factor,
      filters_base=filters_base,
      filters_multipliers=filters_multipliers,
      expansion_multipliers=expansion_multipliers,
      name=name)
  return static_model(search_model, dna)