Internal change

PiperOrigin-RevId: 404080616

official/projects/edgetpu/vision/tasks/__init__.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.
+

official/projects/edgetpu/vision/tasks/image_classification.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.
+
+"""Image classification task definition."""
+import tempfile
+from typing import Any, List, Mapping, Optional, Tuple
+
+from absl import logging
+import tensorflow as tf
+from official.common import dataset_fn
+from official.core import base_task
+from official.core import task_factory
+from official.modeling import tf_utils
+from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config as edgetpu_cfg
+from official.projects.edgetpu.vision.dataloaders import classification_input
+from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v1_model
+from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v2_model
+from official.vision.beta.configs import image_classification as base_cfg
+from official.vision.beta.dataloaders import input_reader_factory
+from official.vision.beta.dataloaders.google import tfds_classification_decoders
+
+
+def get_models() -> Mapping[str, tf.keras.Model]:
+  """Returns the mapping from model type name to Keras model."""
+  model_mapping = {}
+
+  def add_models(name: str, constructor: Any):
+    if name in model_mapping:
+      raise ValueError(f'Model {name} already exists in the mapping.')
+    model_mapping[name] = constructor
+
+  for model in mobilenet_edgetpu_v1_model.MODEL_CONFIGS.keys():
+    add_models(model, mobilenet_edgetpu_v1_model.MobilenetEdgeTPU.from_name)
+
+  for model in mobilenet_edgetpu_v2_model.MODEL_CONFIGS.keys():
+    add_models(model, mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2.from_name)
+
+  return model_mapping
+
+
+def load_searched_model(saved_model_path: str) -> tf.keras.Model:
+  """Loads saved model from file.
+
+  Excepting loading MobileNet-EdgeTPU-V1/V2 models, we can also load searched
+  model directly from saved model path by changing the model path in
+  mobilenet_edgetpu_search (defined in mobilenet_edgetpu_config.py)
+
+  Args:
+    saved_model_path: Directory path for the saved searched model.
+  Returns:
+    Loaded keras model.
+  """
+  with tempfile.TemporaryDirectory() as tmp_dir:
+    if tf.io.gfile.IsDirectory(saved_model_path):
+      tf.io.gfile.RecursivelyCopyDir(saved_model_path, tmp_dir, overwrite=True)
+      load_path = tmp_dir
+    else:
+      raise ValueError('Saved model path is invalid.')
+    load_options = tf.saved_model.LoadOptions(
+        experimental_io_device='/job:localhost')
+    model = tf.keras.models.load_model(load_path, options=load_options)
+
+  return model
+
+
+@task_factory.register_task_cls(edgetpu_cfg.MobilenetEdgeTPUTaskConfig)
+class EdgeTPUTask(base_task.Task):
+  """A task for training MobileNet-EdgeTPU models."""
+
+  def build_model(self):
+    """Builds model for MobileNet-EdgeTPU Task."""
+    model_config = self.task_config.model
+    model_params = model_config.model_params.as_dict()
+    model_name = model_params['model_name']
+    registered_models = get_models()
+    if model_name in registered_models:
+      logging.info('Load MobileNet-EdgeTPU-V1/V2 model.')
+      logging.info(model_params)
+      model = registered_models[model_name](**model_params)
+    elif model_name == 'mobilenet_edgetpu_search':
+      if self.task_config.saved_model_path is None:
+        raise ValueError('If using MobileNet-EdgeTPU-Search model, please'
+                         'specify the saved model path via the'
+                         '--params_override flag.')
+      logging.info('Load saved model (model from search) directly.')
+      model = load_searched_model(self.task_config.saved_model_path)
+    else:
+      raise ValueError('Model has to be mobilenet-edgetpu model or searched'
+                       'model with given saved model path.')
+    model.summary()
+
+    return model
+
+  def initialize(self, model: tf.keras.Model):
+    """Loads pretrained checkpoint."""
+    if not self.task_config.init_checkpoint:
+      return
+
+    ckpt_dir_or_file = self.task_config.init_checkpoint
+    if tf.io.gfile.isdir(ckpt_dir_or_file):
+      ckpt_dir_or_file = tf.train.latest_checkpoint(ckpt_dir_or_file)
+
+    # Restoring checkpoint.
+    if self.task_config.init_checkpoint_modules == 'all':
+      ckpt = tf.train.Checkpoint(**model.checkpoint_items)
+      status = ckpt.read(ckpt_dir_or_file)
+      status.expect_partial().assert_existing_objects_matched()
+    elif self.task_config.init_checkpoint_modules == 'backbone':
+      ckpt = tf.train.Checkpoint(backbone=model.backbone)
+      status = ckpt.read(ckpt_dir_or_file)
+      status.expect_partial().assert_existing_objects_matched()
+    else:
+      raise ValueError(
+          "Only 'all' or 'backbone' can be used to initialize the model.")
+
+    logging.info('Finished loading pretrained checkpoint from %s',
+                 ckpt_dir_or_file)
+
+  def build_inputs(
+      self,
+      params: base_cfg.DataConfig,
+      input_context: Optional[tf.distribute.InputContext] = None
+  ) -> tf.data.Dataset:
+    """Builds classification input."""
+
+    num_classes = self.task_config.model.num_classes
+    input_size = self.task_config.model.input_size
+    image_field_key = self.task_config.train_data.image_field_key
+    label_field_key = self.task_config.train_data.label_field_key
+    is_multilabel = self.task_config.train_data.is_multilabel
+
+    if params.tfds_name:
+      if params.tfds_name in tfds_classification_decoders.TFDS_ID_TO_DECODER_MAP:
+        decoder = tfds_classification_decoders.TFDS_ID_TO_DECODER_MAP[
+            params.tfds_name]()
+      else:
+        raise ValueError('TFDS {} is not supported'.format(params.tfds_name))
+    else:
+      decoder = classification_input.Decoder(
+          image_field_key=image_field_key, label_field_key=label_field_key,
+          is_multilabel=is_multilabel)
+
+    parser = classification_input.Parser(
+        output_size=input_size[:2],
+        num_classes=num_classes,
+        image_field_key=image_field_key,
+        label_field_key=label_field_key,
+        decode_jpeg_only=params.decode_jpeg_only,
+        aug_rand_hflip=params.aug_rand_hflip,
+        aug_type=params.aug_type,
+        is_multilabel=is_multilabel,
+        dtype=params.dtype)
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+
+    dataset = reader.read(input_context=input_context)
+
+    return dataset
+
+  def build_losses(self,
+                   labels: tf.Tensor,
+                   model_outputs: tf.Tensor,
+                   aux_losses: Optional[Any] = None) -> tf.Tensor:
+    """Builds sparse categorical cross entropy loss.
+
+    Args:
+      labels: Input groundtruth labels.
+      model_outputs: Output logits of the classifier.
+      aux_losses: The auxiliarly loss tensors, i.e. `losses` in tf.keras.Model.
+
+    Returns:
+      The total loss tensor.
+    """
+    losses_config = self.task_config.losses
+    is_multilabel = self.task_config.train_data.is_multilabel
+
+    if not is_multilabel:
+      if losses_config.one_hot:
+        total_loss = tf.keras.losses.categorical_crossentropy(
+            labels,
+            model_outputs,
+            from_logits=False,
+            label_smoothing=losses_config.label_smoothing)
+      else:
+        total_loss = tf.keras.losses.sparse_categorical_crossentropy(
+            labels, model_outputs, from_logits=True)
+    else:
+      # Multi-label weighted binary cross entropy loss.
+      total_loss = tf.nn.sigmoid_cross_entropy_with_logits(
+          labels=labels, logits=model_outputs)
+      total_loss = tf.reduce_sum(total_loss, axis=-1)
+
+    total_loss = tf_utils.safe_mean(total_loss)
+    if aux_losses:
+      total_loss += tf.add_n(aux_losses)
+
+    return total_loss
+
+  def build_metrics(self,
+                    training: bool = True) -> List[tf.keras.metrics.Metric]:
+    """Gets streaming metrics for training/validation."""
+    is_multilabel = self.task_config.train_data.is_multilabel
+    if not is_multilabel:
+      k = self.task_config.evaluation.top_k
+      if self.task_config.losses.one_hot:
+        metrics = [
+            tf.keras.metrics.CategoricalAccuracy(name='accuracy'),
+            tf.keras.metrics.TopKCategoricalAccuracy(
+                k=k, name='top_{}_accuracy'.format(k))]
+      else:
+        metrics = [
+            tf.keras.metrics.SparseCategoricalAccuracy(name='accuracy'),
+            tf.keras.metrics.SparseTopKCategoricalAccuracy(
+                k=k, name='top_{}_accuracy'.format(k))]
+    else:
+      metrics = []
+      # These metrics destablize the training if included in training. The jobs
+      # fail due to OOM.
+      # TODO(arashwan): Investigate adding following metric to train.
+      if not training:
+        metrics = [
+            tf.keras.metrics.AUC(
+                name='globalPR-AUC',
+                curve='PR',
+                multi_label=False,
+                from_logits=True),
+            tf.keras.metrics.AUC(
+                name='meanPR-AUC',
+                curve='PR',
+                multi_label=True,
+                num_labels=self.task_config.model.num_classes,
+                from_logits=True),
+        ]
+    return metrics
+
+  def train_step(self,
+                 inputs: Tuple[Any, Any],
+                 model: tf.keras.Model,
+                 optimizer: tf.keras.optimizers.Optimizer,
+                 metrics: Optional[List[Any]] = None):
+    """Does forward and backward.
+
+    Args:
+      inputs: A tuple of of input tensors of (features, labels).
+      model: A tf.keras.Model instance.
+      optimizer: The optimizer for this training step.
+      metrics: A nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    features, labels = inputs
+    is_multilabel = self.task_config.train_data.is_multilabel
+    if self.task_config.losses.one_hot and not is_multilabel:
+      labels = tf.one_hot(labels, self.task_config.model.num_classes)
+
+    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
+    with tf.GradientTape() as tape:
+      outputs = model(features, training=True)
+
+      # Computes per-replica loss.
+      loss = self.build_losses(
+          model_outputs=outputs, labels=labels, aux_losses=model.losses)
+      # Scales loss as the default gradients allreduce performs sum inside the
+      # optimizer.
+      scaled_loss = loss / num_replicas
+
+      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
+      # scaled for numerical stability.
+      if isinstance(
+          optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
+        scaled_loss = optimizer.get_scaled_loss(scaled_loss)
+
+    tvars = model.trainable_variables
+    grads = tape.gradient(scaled_loss, tvars)
+    # Scales back gradient before apply_gradients when LossScaleOptimizer is
+    # used.
+    if isinstance(
+        optimizer, tf.keras.mixed_precision.LossScaleOptimizer):
+      grads = optimizer.get_unscaled_gradients(grads)
+    optimizer.apply_gradients(list(zip(grads, tvars)))
+
+    logs = {self.loss: loss}
+    if metrics:
+      self.process_metrics(metrics, labels, outputs)
+    elif model.compiled_metrics:
+      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
+      logs.update({m.name: m.result() for m in model.metrics})
+    return logs
+
+  def validation_step(self,
+                      inputs: Tuple[Any, Any],
+                      model: tf.keras.Model,
+                      metrics: Optional[List[Any]] = None):
+    """Runs validatation step.
+
+    Args:
+      inputs: A tuple of of input tensors of (features, labels).
+      model: A tf.keras.Model instance.
+      metrics: A nested structure of metrics objects.
+
+    Returns:
+      A dictionary of logs.
+    """
+    features, labels = inputs
+    is_multilabel = self.task_config.train_data.is_multilabel
+    if self.task_config.losses.one_hot and not is_multilabel:
+      labels = tf.one_hot(labels, self.task_config.model.num_classes)
+
+    outputs = self.inference_step(features, model)
+    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
+    loss = self.build_losses(model_outputs=outputs, labels=labels,
+                             aux_losses=model.losses)
+
+    logs = {self.loss: loss}
+    if metrics:
+      self.process_metrics(metrics, labels, outputs)
+    elif model.compiled_metrics:
+      self.process_compiled_metrics(model.compiled_metrics, labels, outputs)
+      logs.update({m.name: m.result() for m in model.metrics})
+    return logs
+
+  def inference_step(self, inputs: tf.Tensor, model: tf.keras.Model):
+    """Performs the forward step."""
+    return model(inputs, training=False)

official/projects/edgetpu/vision/tasks/image_classification_test.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.
+
+# Lint as: python3
+"""Tests for image classification task."""
+
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import orbit
+import tensorflow as tf
+
+from official.common import registry_imports
+from official.core import exp_factory
+from official.modeling import optimization
+from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config
+from official.projects.edgetpu.vision.tasks import image_classification
+
+
+class ImageClassificationTaskTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(('mobilenet_edgetpu_v2_xs'),
+                            ('mobilenet_edgetpu_v2_s'),
+                            ('mobilenet_edgetpu_v2_m'),
+                            ('mobilenet_edgetpu_v2_l'),
+                            ('mobilenet_edgetpu'),
+                            ('mobilenet_edgetpu_dm1p25'),
+                            ('mobilenet_edgetpu_dm1p5'),
+                            ('mobilenet_edgetpu_dm1p75'))
+  def test_task(self, config_name):
+    config = exp_factory.get_exp_config(config_name)
+    config.task.train_data.global_batch_size = 2
+
+    task = image_classification.EdgeTPUTask(config.task)
+    model = task.build_model()
+    metrics = task.build_metrics()
+    strategy = tf.distribute.get_strategy()
+
+    dataset = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
+                                                   config.task.train_data)
+
+    iterator = iter(dataset)
+    opt_factory = optimization.OptimizerFactory(config.trainer.optimizer_config)
+    optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
+    if isinstance(optimizer, optimization.ExponentialMovingAverage
+                 ) and not optimizer.has_shadow_copy:
+      optimizer.shadow_copy(model)
+
+    logs = task.train_step(next(iterator), model, optimizer, metrics=metrics)
+    for metric in metrics:
+      logs[metric.name] = metric.result()
+    self.assertIn('loss', logs)
+    self.assertIn('accuracy', logs)
+    self.assertIn('top_5_accuracy', logs)
+    logs = task.validation_step(next(iterator), model, metrics=metrics)
+    for metric in metrics:
+      logs[metric.name] = metric.result()
+    self.assertIn('loss', logs)
+    self.assertIn('accuracy', logs)
+    self.assertIn('top_5_accuracy', logs)
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/edgetpu/vision/tasks/semantic_segmentation.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.
+
+"""Image segmentation task definition."""
+from typing import Any, Mapping, Optional
+
+from absl import logging
+import tensorflow as tf
+
+from official.common import dataset_fn
+from official.core import task_factory
+from official.modeling.hyperparams import config_definitions as cfg
+from official.projects.edgetpu.vision.configs import semantic_segmentation_config as exp_cfg
+from official.projects.edgetpu.vision.configs import semantic_segmentation_searched_config as searched_cfg
+from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v1_model
+from official.projects.edgetpu.vision.modeling import mobilenet_edgetpu_v2_model
+from official.projects.edgetpu.vision.modeling.backbones import mobilenet_edgetpu  # pylint: disable=unused-import
+from official.projects.edgetpu.vision.modeling.heads import bifpn_head
+from official.vision.beta.dataloaders import input_reader_factory
+from official.vision.beta.dataloaders import segmentation_input
+from official.vision.beta.dataloaders import tfds_factory
+from official.vision.beta.ops import preprocess_ops
+from official.vision.beta.tasks import semantic_segmentation
+
+
+class ClassMappingParser(segmentation_input.Parser):
+  """Same parser but maps classes max_class+1... to class 0."""
+
+  max_class = 31
+
+  def _prepare_image_and_label(self, data):
+    """Prepare normalized image and label."""
+    image = tf.io.decode_image(data['image/encoded'], channels=3)
+    label = tf.io.decode_image(data['image/segmentation/class/encoded'],
+                               channels=1)
+    height = data['image/height']
+    width = data['image/width']
+    image = tf.reshape(image, (height, width, 3))
+
+    label = tf.reshape(label, (1, height, width))
+    label = tf.where(
+        tf.math.greater(label, self.max_class), tf.zeros_like(label), label)
+    label = tf.where(tf.math.equal(label, 0), tf.ones_like(label)*255, label)
+    label = tf.cast(label, tf.float32)
+    # Normalizes image with mean and std pixel values.
+    image = preprocess_ops.normalize_image(
+        image, offset=[0.5, 0.5, 0.5], scale=[0.5, 0.5, 0.5])
+    return image, label
+
+
+@task_factory.register_task_cls(exp_cfg.CustomSemanticSegmentationTaskConfig)
+class CustomSemanticSegmentationTask(
+    semantic_segmentation.SemanticSegmentationTask):
+  """A task for semantic segmentation."""
+
+  def build_inputs(self,
+                   params: cfg.DataConfig,
+                   input_context: Optional[tf.distribute.InputContext] = None):
+    """Builds classification input."""
+    ignore_label = self.task_config.losses.ignore_label
+
+    if params.tfds_name:
+      decoder = tfds_factory.get_segmentation_decoder(params.tfds_name)
+    else:
+      decoder = segmentation_input.Decoder()
+
+    parser = ClassMappingParser(
+        output_size=params.output_size,
+        crop_size=params.crop_size,
+        ignore_label=ignore_label,
+        resize_eval_groundtruth=params.resize_eval_groundtruth,
+        groundtruth_padded_size=params.groundtruth_padded_size,
+        aug_scale_min=params.aug_scale_min,
+        aug_scale_max=params.aug_scale_max,
+        aug_rand_hflip=params.aug_rand_hflip,
+        dtype=params.dtype)
+
+    parser.max_class = self.task_config.model.num_classes-1
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+
+    dataset = reader.read(input_context=input_context)
+
+    return dataset
+
+
+class AutosegEdgeTPU(tf.keras.Model):
+  """Segmentation keras network without pre/post-processing."""
+
+  def __init__(self,
+               model_params,
+               min_level=3,
+               max_level=8,
+               output_filters=96,
+               model_config=None,
+               use_original_backbone_features=False,
+               is_training_bn=True,
+               strategy='tpu',
+               data_format='channels_last',
+               pooling_type='avg',
+               fpn_num_filters=96,
+               apply_bn_for_resampling=True,
+               conv_after_downsample=True,
+               upsampling_type='bilinear',
+               conv_bn_act_pattern=True,
+               conv_type='sep_3',
+               head_conv_type='sep_3',
+               act_type='relu6',
+               fpn_weight_method='sum',
+               output_weight_method='sum',
+               fullres_output=False,
+               num_classes=32,
+               name='autoseg_edgetpu'):
+    """Initialize model."""
+    super().__init__()
+    self.min_level = min_level
+    self.max_level = max_level
+    self.use_original_backbone_features = use_original_backbone_features
+    self.strategy = strategy
+    self.data_format = data_format
+    model_name = model_params['model_name']
+    self.backbone = get_models()[model_name](**model_params)
+
+    # Feature network.
+    self.resample_layers = []  # additional resampling layers.
+    if use_original_backbone_features:
+      start_level = 6
+    else:
+      # Not using original backbone features will (1) Use convolutions to
+      # process all backbone features before feeding into FPN. (2) Use an extra
+      # convolution to get higher level features, while preserve the channel
+      # size from the last layer of backbone.
+      start_level = min_level
+      self.downsample_layers = []
+      for level in range(start_level, max_level + 1):
+        self.downsample_layers.append(
+            bifpn_head.ResampleFeatureMap(
+                feat_level=(level - min_level),
+                target_num_channels=fpn_num_filters,
+                is_training_bn=is_training_bn,
+                strategy=strategy,
+                data_format=data_format,
+                pooling_type=pooling_type,
+                name='downsample_p%d' % level,
+            ))
+    for level in range(start_level, max_level + 1):
+      # Adds a coarser level by downsampling the last feature map.
+      self.resample_layers.append(
+          bifpn_head.ResampleFeatureMap(
+              feat_level=(level - min_level),
+              target_num_channels=fpn_num_filters,
+              apply_bn=apply_bn_for_resampling,
+              is_training_bn=is_training_bn,
+              conv_after_downsample=conv_after_downsample,
+              strategy=strategy,
+              data_format=data_format,
+              pooling_type=pooling_type,
+              upsampling_type=upsampling_type,
+              name='resample_p%d' % level,
+          ))
+    self.fpn_cells = bifpn_head.FPNCells(
+        min_level=min_level,
+        max_level=max_level,
+        fpn_num_filters=fpn_num_filters,
+        apply_bn_for_resampling=apply_bn_for_resampling,
+        is_training_bn=is_training_bn,
+        conv_after_downsample=conv_after_downsample,
+        conv_bn_act_pattern=conv_bn_act_pattern,
+        conv_type=conv_type,
+        act_type=act_type,
+        strategy=strategy,
+        fpn_weight_method=fpn_weight_method,
+        data_format=data_format,
+        pooling_type=pooling_type,
+        upsampling_type=upsampling_type,
+        fpn_name='bifpn')
+
+    self.seg_class_net = bifpn_head.SegClassNet(
+        min_level=min_level,
+        max_level=max_level,
+        output_filters=output_filters,
+        apply_bn_for_resampling=apply_bn_for_resampling,
+        is_training_bn=is_training_bn,
+        conv_after_downsample=conv_after_downsample,
+        conv_bn_act_pattern=conv_bn_act_pattern,
+        head_conv_type=head_conv_type,
+        act_type=act_type,
+        strategy=strategy,
+        output_weight_method=output_weight_method,
+        data_format=data_format,
+        pooling_type=pooling_type,
+        upsampling_type=upsampling_type,
+        fullres_output=fullres_output,
+        num_classes=num_classes)
+
+  def call(self, inputs, training):
+    # call backbone network.
+    all_feats = self.backbone(inputs, training=training)
+    if self.use_original_backbone_features:
+      feats = all_feats[self.min_level:self.max_level + 1]
+      for resample_layer in self.resample_layers:
+        feats.append(resample_layer(feats[-1], training, None))
+    else:
+      feats = []
+      for downsample_layer in self.downsample_layers:
+        all_feats.append(downsample_layer(all_feats[-1], training, None))
+      for level in range(self.min_level - 1, self.max_level):
+        feats.append(self.resample_layers[level - self.min_level + 1](
+            all_feats[level], training, all_feats[self.min_level - 1:]))
+
+    # call feature network.
+    feats = self.fpn_cells(feats, training)
+
+    # call class/box output network.
+    class_outputs = self.seg_class_net(feats, all_feats, training)
+
+    return class_outputs
+
+
+def get_models() -> Mapping[str, tf.keras.Model]:
+  """Returns the mapping from model type name to Keras model."""
+  model_mapping = {}
+
+  def add_models(name: str, constructor: Any):
+    if name in model_mapping:
+      raise ValueError(f'Model {name} already exists in the mapping.')
+    model_mapping[name] = constructor
+
+  for model in mobilenet_edgetpu_v1_model.MODEL_CONFIGS.keys():
+    add_models(model, mobilenet_edgetpu_v1_model.MobilenetEdgeTPU.from_name)
+
+  for model in mobilenet_edgetpu_v2_model.MODEL_CONFIGS.keys():
+    add_models(model, mobilenet_edgetpu_v2_model.MobilenetEdgeTPUV2.from_name)
+
+  return model_mapping
+
+
+@task_factory.register_task_cls(searched_cfg.AutosegEdgeTPUTaskConfig)
+class AutosegEdgeTPUTask(semantic_segmentation.SemanticSegmentationTask):
+  """A task for training the AutosegEdgeTPU models."""
+
+  def build_model(self):
+    """Builds model for training task."""
+    model_config = self.task_config.model
+    model_params = model_config.model_params.as_dict()
+    model = AutosegEdgeTPU(
+        model_params,
+        min_level=model_config.head.min_level,
+        max_level=model_config.head.max_level,
+        fpn_num_filters=model_config.head.fpn_num_filters,
+        num_classes=model_config.num_classes)
+    logging.info(model_params)
+    return model
+
+  # TODO(suyoggupta): Dedup this function across tasks.
+  def build_inputs(self,
+                   params: cfg.DataConfig,
+                   input_context: Optional[tf.distribute.InputContext] = None):
+    """Builds inputs for the segmentation task."""
+    ignore_label = self.task_config.losses.ignore_label
+
+    if params.tfds_name:
+      decoder = tfds_factory.get_segmentation_decoder(params.tfds_name)
+    else:
+      decoder = segmentation_input.Decoder()
+
+    parser = ClassMappingParser(
+        output_size=params.output_size,
+        crop_size=params.crop_size,
+        ignore_label=ignore_label,
+        resize_eval_groundtruth=params.resize_eval_groundtruth,
+        groundtruth_padded_size=params.groundtruth_padded_size,
+        aug_scale_min=params.aug_scale_min,
+        aug_scale_max=params.aug_scale_max,
+        aug_rand_hflip=params.aug_rand_hflip,
+        dtype=params.dtype)
+
+    parser.max_class = self.task_config.model.num_classes - 1
+
+    reader = input_reader_factory.input_reader_generator(
+        params,
+        dataset_fn=dataset_fn.pick_dataset_fn(params.file_type),
+        decoder_fn=decoder.decode,
+        parser_fn=parser.parse_fn(params.is_training))
+
+    dataset = reader.read(input_context=input_context)
+    return dataset

official/projects/edgetpu/vision/tasks/semantic_segmentation_test.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.
+
+# Lint as: python3
+"""Tests for semantic segmentation task."""
+
+# pylint: disable=unused-import
+from absl.testing import parameterized
+import orbit
+import tensorflow as tf
+
+from official.core import exp_factory
+from official.modeling import optimization
+from official.projects.edgetpu.vision.configs import semantic_segmentation_config as seg_cfg
+from official.projects.edgetpu.vision.configs import semantic_segmentation_searched_config as autoseg_cfg
+from official.projects.edgetpu.vision.tasks import semantic_segmentation as img_seg_task
+from official.vision import beta
+
+
+class SemanticSegmentationTaskTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k_32',),
+                            ('deeplabv3plus_mobilenet_edgetpuv2_s_ade20k_32',),
+                            ('deeplabv3plus_mobilenet_edgetpuv2_m_ade20k_32',))
+  def test_task(self, config_name):
+    config_to_backbone_mapping = {
+        'deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k_32':
+            'mobilenet_edgetpu_v2_xs',
+        'deeplabv3plus_mobilenet_edgetpuv2_s_ade20k_32':
+            'mobilenet_edgetpu_v2_s',
+        'deeplabv3plus_mobilenet_edgetpuv2_m_ade20k_32':
+            'mobilenet_edgetpu_v2_m',
+    }
+    config = seg_cfg.seg_deeplabv3plus_ade20k_32(
+        config_to_backbone_mapping[config_name], init_backbone=False)
+    config.task.train_data.global_batch_size = 1
+    config.task.train_data.shuffle_buffer_size = 2
+    config.task.validation_data.shuffle_buffer_size = 2
+    config.task.validation_data.global_batch_size = 1
+    config.task.train_data.output_size = [32, 32]
+    config.task.validation_data.output_size = [32, 32]
+    config.task.model.decoder.aspp.pool_kernel_size = None
+    config.task.model.backbone.dilated_resnet.model_id = 50
+    config.task.model.backbone.dilated_resnet.output_stride = 16
+
+    task = img_seg_task.CustomSemanticSegmentationTask(config.task)
+    model = task.build_model()
+    metrics = task.build_metrics()
+    strategy = tf.distribute.get_strategy()
+
+    dataset = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
+                                                   config.task.train_data)
+
+    iterator = iter(dataset)
+    opt_factory = optimization.OptimizerFactory(config.trainer.optimizer_config)
+    optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
+    logs = task.train_step(next(iterator), model, optimizer, metrics=metrics)
+    self.assertIn('loss', logs)
+    logs = task.validation_step(next(iterator), model,
+                                metrics=task.build_metrics(training=False))
+    self.assertIn('loss', logs)
+
+
+class AutosegEdgeTPUTaskTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(('autoseg_edgetpu_xs',))
+  def test_task(self, config_name):
+    config_to_backbone_mapping = {
+        'autoseg_edgetpu_xs': 'autoseg_edgetpu_backbone_xs',
+        'autoseg_edgetpu_s': 'autoseg_edgetpu_backone_s'
+    }
+    config = autoseg_cfg.autoseg_edgetpu_experiment_config(
+        config_to_backbone_mapping[config_name], init_backbone=False)
+    config.task.train_data.global_batch_size = 2
+    config.task.train_data.shuffle_buffer_size = 2
+    config.task.validation_data.shuffle_buffer_size = 2
+    config.task.validation_data.global_batch_size = 2
+    config.task.train_data.output_size = [512, 512]
+    config.task.validation_data.output_size = [512, 512]
+
+    task = img_seg_task.AutosegEdgeTPUTask(config.task)
+    model = task.build_model()
+    metrics = task.build_metrics()
+    strategy = tf.distribute.get_strategy()
+
+    dataset = orbit.utils.make_distributed_dataset(strategy, task.build_inputs,
+                                                   config.task.train_data)
+
+    iterator = iter(dataset)
+    opt_factory = optimization.OptimizerFactory(config.trainer.optimizer_config)
+    optimizer = opt_factory.build_optimizer(opt_factory.build_learning_rate())
+    if isinstance(optimizer, optimization.ExponentialMovingAverage
+                 ) and not optimizer.has_shadow_copy:
+      optimizer.shadow_copy(model)
+
+    logs = task.train_step(next(iterator), model, optimizer, metrics=metrics)
+    self.assertIn('loss', logs)
+    logs = task.validation_step(
+        next(iterator), model, metrics=task.build_metrics(training=False))
+    self.assertIn('loss', logs)
+    model.summary()
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/edgetpu/vision/train.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.
+
+# Lint as: python3
+"""TensorFlow Model Garden Vision training for MobileNet-EdgeTPU."""
+
+from absl import app
+from absl import flags
+import gin
+
+# pylint: disable=unused-import
+from official.common import registry_imports
+# pylint: enable=unused-import
+from official.common import distribute_utils
+from official.common import flags as tfm_flags
+from official.core import task_factory
+from official.core import train_lib
+from official.core import train_utils
+from official.modeling import performance
+# pylint: disable=unused-import
+from official.projects.edgetpu.vision.configs import mobilenet_edgetpu_config
+from official.projects.edgetpu.vision.configs import semantic_segmentation_config
+from official.projects.edgetpu.vision.configs import semantic_segmentation_searched_config
+from official.projects.edgetpu.vision.modeling.backbones import mobilenet_edgetpu
+from official.projects.edgetpu.vision.tasks import image_classification
+from official.projects.edgetpu.vision.tasks import semantic_segmentation
+# pylint: enable=unused-import
+
+FLAGS = flags.FLAGS
+
+
+def main(_):
+  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params)
+  params = train_utils.parse_configuration(FLAGS)
+
+  model_dir = FLAGS.model_dir
+  if 'train' in FLAGS.mode:
+    # Pure eval modes do not output yaml files. Otherwise continuous eval job
+    # may race against the train job for writing the same file.
+    train_utils.serialize_config(params, model_dir)
+
+  # Sets mixed_precision policy. Using 'mixed_float16' or 'mixed_bfloat16'
+  # can have significant impact on model speeds by utilizing float16 in case of
+  # GPUs, and bfloat16 in the case of TPUs. loss_scale takes effect only when
+  # dtype is float16
+  if params.runtime.mixed_precision_dtype:
+    performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype)
+  distribution_strategy = distribute_utils.get_distribution_strategy(
+      distribution_strategy=params.runtime.distribution_strategy,
+      all_reduce_alg=params.runtime.all_reduce_alg,
+      num_gpus=params.runtime.num_gpus,
+      tpu_address=params.runtime.tpu)
+  with distribution_strategy.scope():
+    task = task_factory.get_task(params.task, logging_dir=model_dir)
+
+  train_lib.run_experiment(
+      distribution_strategy=distribution_strategy,
+      task=task,
+      mode=FLAGS.mode,
+      params=params,
+      model_dir=model_dir)
+
+  train_utils.save_gin_config(FLAGS.mode, model_dir)
+
+if __name__ == '__main__':
+  tfm_flags.define_flags()
+  flags.mark_flags_as_required(['mode', 'model_dir'])
+  app.run(main)