Merge branch 'tensorflow:master' into exp_pr2

official/projects/edgetpu/vision/serving/export_tflite_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.
+
+"""Tests for export_tflite."""
+
+import itertools
+import os
+
+from absl.testing import parameterized
+import tensorflow as tf
+
+from official.projects.edgetpu.vision.serving import export_util
+
+
+def _build_model(config):
+  model = export_util.build_experiment_model(config.model_name)
+  model_input = tf.keras.Input(
+      shape=(config.image_size, config.image_size, 3), batch_size=1)
+  model_output = export_util.finalize_serving(model(model_input), config)
+  model_for_inference = tf.keras.Model(model_input, model_output)
+  return model_for_inference
+
+
+def _dump_tflite(model, config):
+  converter = tf.lite.TFLiteConverter.from_keras_model(model)
+  export_util.configure_tflite_converter(config, converter)
+  tflite_buffer = converter.convert()
+  tf.io.gfile.makedirs(os.path.dirname(config.output_dir))
+  tflite_path = os.path.join(config.output_dir, f'{config.model_name}.tflite')
+  tf.io.gfile.GFile(tflite_path, 'wb').write(tflite_buffer)
+  return tflite_path
+
+
+SEG_MODELS = [
+    'autoseg_edgetpu_xs',
+]
+FINALIZE_METHODS = [
+    'resize512,argmax,squeeze', 'resize256,argmax,resize512,squeeze',
+    'resize128,argmax,resize512,squeeze'
+]
+
+
+class ExportTfliteTest(tf.test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('mobilenet_edgetpu_v2_xs', 224),
+      ('autoseg_edgetpu_xs', 512),
+      ('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k', 512),
+      ('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k_32', 512),
+  )
+  def test_model_build_and_export_tflite(self, model_name, image_size):
+    tmp_dir = self.create_tempdir().full_path
+    config = export_util.ExportConfig(
+        model_name=model_name, image_size=image_size, output_dir=tmp_dir)
+    config.quantization_config.quantize = False
+    model = _build_model(config)
+    tflite_path = _dump_tflite(model, config)
+    self.assertTrue(tf.io.gfile.exists(tflite_path))
+
+  @parameterized.parameters(
+      ('mobilenet_edgetpu_v2_xs', 224),
+      ('autoseg_edgetpu_xs', 512),
+      ('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k', 512),
+      ('deeplabv3plus_mobilenet_edgetpuv2_xs_ade20k_32', 512),
+  )
+  def test_model_build_and_export_saved_model(self, model_name, image_size):
+    tmp_dir = self.create_tempdir().full_path
+    config = export_util.ExportConfig(
+        model_name=model_name, image_size=image_size, output_dir=tmp_dir)
+    model = _build_model(config)
+    saved_model_path = os.path.join(config.output_dir, config.model_name)
+    model.save(saved_model_path)
+    self.assertTrue(tf.saved_model.contains_saved_model(saved_model_path))
+
+  @parameterized.parameters(itertools.product(SEG_MODELS, FINALIZE_METHODS))
+  def test_segmentation_finalize_methods(self, model_name, finalize_method):
+    tmp_dir = self.create_tempdir().full_path
+    config = export_util.ExportConfig(
+        model_name=model_name,
+        image_size=512,
+        output_dir=tmp_dir,
+        finalize_method=finalize_method.split(','))
+    config.quantization_config.quantize = False
+    model = _build_model(config)
+    model_input = tf.random.normal([1, config.image_size, config.image_size, 3])
+    self.assertEqual(
+        model(model_input).get_shape().as_list(),
+        [1, config.image_size, config.image_size])
+
+
+if __name__ == '__main__':
+  tf.test.main()

official/projects/edgetpu/vision/serving/export_util.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.
+
+"""Implements serving with custom post processing."""
+
+import dataclasses
+from typing import List, Optional
+
+import tensorflow as tf
+import tensorflow_datasets as tfds
+
+from official.core import exp_factory
+from official.core import task_factory
+from official.modeling.hyperparams import base_config
+# 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 import custom_layers
+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 as edgetpu_semantic_segmentation
+from official.vision.beta.tasks import semantic_segmentation
+# pylint: enable=unused-import
+
+MEAN_RGB = [127.5, 127.5, 127.5]
+STDDEV_RGB = [127.5, 127.5, 127.5]
+
+
+@dataclasses.dataclass
+class QuantizationConfig(base_config.Config):
+  """Configuration for post training quantization.
+
+  Attributes:
+    quantize: Whether to quantize model before exporting tflite.
+    quantize_less_restrictive: Allows non int8 based intermediate types,
+      automatic model output type.
+    use_experimental_quantizer: Enables experimental quantizer of
+      TFLiteConverter 2.0.
+    num_calibration_steps: Number of post-training quantization calibration
+      steps to run.
+    dataset_name: Name of the dataset to use for quantization calibration.
+    dataset_dir: Dataset location.
+    dataset_split: The dataset split (train, validation etc.) to use for
+      calibration.
+  """
+  quantize: bool = False
+  quantize_less_restrictive: bool = False
+  use_experimental_quantizer: bool = True
+  dataset_name: Optional[str] = None
+  dataset_dir: Optional[str] = None
+  dataset_split: Optional[str] = None
+  num_calibration_steps: int = 100
+
+
+@dataclasses.dataclass
+class ExportConfig(base_config.Config):
+  """Configuration for exporting models as tflite and saved_models.
+
+  Attributes:
+    model_name: One of the registered model names
+    ckpt_path: Path of the training checkpoint. If not provided tflite with
+      random parameters is exported.
+    ckpt_format: Format of the checkpoint. tf_checkpoint is for ckpt files from
+      tf.train.Checkpoint.save() method. keras_checkpoint is for ckpt files from
+      keras.Model.save_weights() method
+    output_dir: Directory to output exported files.
+    image_size: Size of the input image. Ideally should be the same as the
+      image_size used in training config
+    output_layer: Layer name to take the output from. Can be used to take the
+      output from an intermediate layer. None means use the original model
+      output.
+    finalize_method: 'Additional layers to be added to customize serving output
+      Supported are (none|(argmax|resize<?>)[,...]).
+      - none: do not add extra serving layers.
+      - argmax: adds argmax.
+      - squeeze: removes dimensions (except batch dim) of size 1 from the shape
+        of a tensor.
+      - resize<?> (for example resize512): adds resize bilinear|nn to <?> size.
+      For example: --finalize_method=resize128,argmax,resize512,squeeze will do
+        resize bilinear to 128x128, then argmax then resize nn to 512x512
+  """
+  quantization_config: QuantizationConfig = QuantizationConfig()
+  model_name: str = None
+  ckpt_path: Optional[str] = None
+  ckpt_format: Optional[str] = 'tf_checkpoint'
+  output_dir: str = '/tmp/'
+  image_size: int = 224
+  output_layer: Optional[str] = None
+  finalize_method: Optional[List[str]] = None
+
+
+def finalize_serving(model_output, export_config):
+  """Adds extra layers based on the provided configuration."""
+
+  finalize_method = export_config.finalize_method
+  output_layer = model_output
+  if not finalize_method or finalize_method[0] == 'none':
+    return output_layer
+  discrete = False
+  for i in range(len(finalize_method)):
+    if finalize_method[i] == 'argmax':
+      discrete = True
+      is_argmax_last = (i + 1) == len(finalize_method)
+      if is_argmax_last:
+        output_layer = tf.argmax(
+            output_layer, axis=3, output_type=tf.dtypes.int32)
+      else:
+        # TODO(tohaspiridonov): add first_match=False when cl/383951533 submited
+        output_layer = custom_layers.argmax(
+            output_layer, keepdims=True, epsilon=1e-3)
+    elif finalize_method[i] == 'squeeze':
+      output_layer = tf.squeeze(output_layer, axis=3)
+    else:
+      resize_params = finalize_method[i].split('resize')
+      if len(resize_params) != 2 or resize_params[0]:
+        raise ValueError('Cannot finalize with ' + finalize_method[i] + '.')
+      resize_to_size = int(resize_params[1])
+      if discrete:
+        output_layer = tf.image.resize(
+            output_layer, [resize_to_size, resize_to_size],
+            method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+      else:
+        output_layer = tf.image.resize(
+            output_layer, [resize_to_size, resize_to_size],
+            method=tf.image.ResizeMethod.BILINEAR)
+  return output_layer
+
+
+def preprocess_for_quantization(image_data, image_size, crop_padding=32):
+  """Crops to center of image with padding then scales, normalizes image_size.
+
+  Args:
+    image_data: A 3D Tensor representing the RGB image data. Image can be of
+      arbitrary height and width.
+    image_size: image height/width dimension.
+    crop_padding: the padding size to use when centering the crop.
+
+  Returns:
+    A decoded and cropped image Tensor. Image is normalized to [-1,1].
+  """
+  shape = tf.shape(image_data)
+  image_height = shape[0]
+  image_width = shape[1]
+
+  padded_center_crop_size = tf.cast(
+      (image_size * 1.0 / (image_size + crop_padding)) *
+      tf.cast(tf.minimum(image_height, image_width), tf.float32), tf.int32)
+
+  offset_height = ((image_height - padded_center_crop_size) + 1) // 2
+  offset_width = ((image_width - padded_center_crop_size) + 1) // 2
+
+  image = tf.image.crop_to_bounding_box(
+      image_data,
+      offset_height=offset_height,
+      offset_width=offset_width,
+      target_height=padded_center_crop_size,
+      target_width=padded_center_crop_size)
+
+  image = tf.image.resize([image], [image_size, image_size],
+                          method=tf.image.ResizeMethod.BILINEAR)[0]
+  image = tf.cast(image, tf.float32)
+  image -= tf.constant(MEAN_RGB)
+  image /= tf.constant(STDDEV_RGB)
+  return image
+
+
+def representative_dataset_gen(export_config):
+  """Gets a python generator of numpy arrays for the given dataset."""
+  quantization_config = export_config.quantization_config
+  dataset = tfds.builder(
+      quantization_config.dataset_name,
+      data_dir=quantization_config.dataset_dir)
+  dataset.download_and_prepare()
+  data = dataset.as_dataset()[quantization_config.dataset_split]
+  iterator = data.as_numpy_iterator()
+  for _ in range(quantization_config.num_calibration_steps):
+    features = next(iterator)
+    image = features['image']
+    image = preprocess_for_quantization(image, export_config.image_size)
+    image = tf.reshape(
+        image, [1, export_config.image_size, export_config.image_size, 3])
+    yield [image]
+
+
+def configure_tflite_converter(export_config, converter):
+  """Common code for picking up quantization parameters."""
+  quantization_config = export_config.quantization_config
+  if quantization_config.quantize:
+    if quantization_config.dataset_dir is None:
+      raise ValueError(
+          'Must provide a representative dataset when quantizing the model.')
+    converter.optimizations = [tf.lite.Optimize.DEFAULT]
+    converter.target_spec.supported_ops = [
+        tf.lite.OpsSet.TFLITE_BUILTINS_INT8
+    ]
+    converter.inference_input_type = tf.int8
+    converter.inference_output_type = tf.int8
+    if quantization_config.quantize_less_restrictive:
+      converter.target_spec.supported_ops += [
+          tf.lite.OpsSet.TFLITE_BUILTINS
+      ]
+      converter.inference_output_type = tf.float32
+    def _representative_dataset_gen():
+      return representative_dataset_gen(export_config)
+
+    converter.representative_dataset = _representative_dataset_gen
+
+
+def build_experiment_model(experiment_type):
+  """Builds model from experiment type configuration."""
+  params = exp_factory.get_exp_config(experiment_type)
+  params.validate()
+  params.lock()
+  task = task_factory.get_task(params.task)
+  return task.build_model()

official/projects/edgetpu/vision/serving/inference_visualization_tool.ipynb

+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Klhdy8pnk5J8"
+      },
+      "source": [
+        "**A tool to visualize the segmentation model inference output.**\\\n",
+        "This tool is used verify that the exported tflite can produce expected segmentation results.\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-vGHZSPWXbyu"
+      },
+      "outputs": [],
+      "source": [
+        "MODEL='gs://**/placeholder_for_edgetpu_models/autoseg/segmentation_search_edgetpu_s_not_fused.tflite'#@param\n",
+        "IMAGE_HOME = 'gs://**/PS_Compare/20190711'#@param\n",
+        "# Relative image file names separated by comas.\n",
+        "TEST_IMAGES = 'ADE_val_00001626.jpg,ADE_val_00001471.jpg,ADE_val_00000557.jpg'#@param\n",
+        "IMAGE_WIDTH = 512 #@param\n",
+        "IMAGE_HEIGHT = 512 #@param"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zzhF1ASDkxTU"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import tensorflow as tf\n",
+        "from PIL import Image as PILImage\n",
+        "import matplotlib.pyplot as plt\n",
+        "from scipy import ndimage"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "AXaJgLg1ml16"
+      },
+      "outputs": [],
+      "source": [
+        "# This block creates local copies of /cns and /x20 files.\n",
+        "TEST_IMAGES=','.join([IMAGE_HOME+'/'+image for image in TEST_IMAGES.split(',')])\n",
+        "\n",
+        "# The tflite interpreter only accepts model in local path.\n",
+        "def local_copy(awaypath):\n",
+        "  localpath = '/tmp/' + awaypath.split('/')[-1]\n",
+        "  !rm -f {localpath}\n",
+        "  !fileutil cp -f {awaypath} {localpath}\n",
+        "  !ls -lht {localpath}\n",
+        "  %download_file {localpath}\n",
+        "  return localpath\n",
+        "\n",
+        "IMAGES = [local_copy(image) for image in TEST_IMAGES.split(',')]\n",
+        "MODEL_COPY=local_copy(MODEL)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "KhS1lOrxHp5C"
+      },
+      "outputs": [],
+      "source": [
+        "# Creates a 6px wide boolean edge mask to highlight the segmentation.\n",
+        "def edge(mydata):\n",
+        "  mydata = mydata.reshape(512, 512)\n",
+        "  mydatat = mydata.transpose([1, 0])\n",
+        "  mydata = np.convolve(mydata.reshape(-1), [-1, 0, 1], mode='same').reshape(512, 512)\n",
+        "  mydatat = np.convolve(mydatat.reshape(-1), [-1, 0, 1], mode='same').reshape(512, 512).transpose([1, 0])\n",
+        "  mydata = np.maximum((mydata != 0).astype(np.int8), (mydatat != 0).astype(np.int8))\n",
+        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
+        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
+        "  mydata = ndimage.binary_dilation(mydata).astype(np.int8)\n",
+        "  return mydata"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "GdlsbiVqL5JZ"
+      },
+      "outputs": [],
+      "source": [
+        "def run_model(input_data):\n",
+        "  _input_data = input_data\n",
+        "  _input_data = (_input_data-128).astype(np.int8)\n",
+        "  # Load the tflite model and allocate tensors.\n",
+        "  interpreter_x = tf.lite.Interpreter(model_path=MODEL_COPY)\n",
+        "  interpreter_x.allocate_tensors()\n",
+        "  # Get input and output tensors.\n",
+        "  input_details = interpreter_x.get_input_details()\n",
+        "  output_details = interpreter_x.get_output_details()\n",
+        "  interpreter_x.set_tensor(input_details[0]['index'], _input_data)\n",
+        "  interpreter_x.invoke()\n",
+        "  output_data = interpreter_x.get_tensor(output_details[0]['index'])\n",
+        "  return output_data.reshape((512, 512, 1))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1mot5M_nl5P7"
+      },
+      "outputs": [],
+      "source": [
+        "# Set visualization wind sizes.\n",
+        "fig, ax = plt.subplots(max(len(IMAGES),2), 3)\n",
+        "fig.set_figwidth(30)\n",
+        "fig.set_figheight(10*max(len(IMAGES),2))\n",
+        "\n",
+        "# Read and test image.\n",
+        "for r, image in enumerate(IMAGES):\n",
+        "  im = PILImage.open(image).convert('RGB')\n",
+        "  min_dim=min(im.size[0], im.size[1])\n",
+        "  im = im.resize((IMAGE_WIDTH*im.size[0] // min_dim, IMAGE_HEIGHT*im.size[1] // min_dim))\n",
+        "  input_data = np.expand_dims(im, axis=0)\n",
+        "  input_data = input_data[:, :IMAGE_WIDTH,:IMAGE_HEIGHT]\n",
+        "  ax[r, 0].imshow(input_data.reshape([512, 512, 3]).astype(np.uint8))\n",
+        "  ax[r, 0].set_title('Original')\n",
+        "  ax[r, 0].grid(False)\n",
+        "\n",
+        "  # Test the model on random input data.\n",
+        "  output_data = run_model(input_data)\n",
+        "  ax[r, 1].imshow(output_data, vmin = 0, vmax = 32)\n",
+        "  ax[r, 1].set_title('Segmentation')\n",
+        "  ax[r, 1].grid(False)\n",
+        "\n",
+        "  output_data = np.reshape(np.minimum(output_data, 32), [512,512])\n",
+        "  output_edge = edge(output_data).reshape(512,512, 1)\n",
+        "  output_data = np.stack([output_data%3, (output_data//3)%3, (output_data//9)%3], axis = -1)\n",
+        "  \n",
+        "  output_data = input_data.reshape([512, 512, 3]).astype(np.float32) * (1-output_edge) + output_data * output_edge * 255\n",
+        "  ax[r, 2].imshow(output_data.astype(np.uint8), vmin = 0, vmax = 256)\n",
+        "  ax[r, 2].set_title('Segmentation \u0026 original')\n",
+        "  ax[r, 2].grid(False)\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "//quality/ranklab/experimental/notebook:rl_colab",
+        "kind": "private"
+      },
+      "name": "Inference_visualization_tool.ipynb",
+      "private_outputs": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

official/projects/edgetpu/vision/serving/tflite_imagenet_evaluator.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.
+
+"""Evaluates image classification accuracy using TFLite Interpreter."""
+
+import dataclasses
+import multiprocessing.pool as mp
+from typing import Tuple
+from absl import logging
+import numpy as np
+import tensorflow as tf
+
+
+@dataclasses.dataclass
+class EvaluationInput():
+  """Contains image and its label as evaluation input."""
+  image: tf.Tensor
+  label: tf.Tensor
+
+
+class AccuracyEvaluator():
+  """Evaluates image classification accuracy using TFLite Interpreter.
+
+  Attributes:
+    model_content: The contents of a TFLite model.
+    num_threads: Number of threads used to evaluate images.
+    thread_batch_size: Batch size assigned to each thread.
+    image_size: Width/Height of the images.
+    num_classes: Number of classes predicted by the model.
+    resize_method: Resize method to use during image preprocessing.
+  """
+
+  def __init__(self,
+               model_content: bytes,
+               dataset: tf.data.Dataset,
+               num_threads: int = 16):
+    self._model_content: bytes = model_content
+    self._dataset = dataset
+    self._num_threads: int = num_threads
+
+  def evaluate_single_image(self, eval_input: EvaluationInput) -> bool:
+    """Evaluates a given single input.
+
+    Args:
+      eval_input: EvaluationInput holding image and label.
+
+    Returns:
+      Whether the estimation is correct.
+    """
+    interpreter = tf.lite.Interpreter(
+        model_content=self._model_content, num_threads=1)
+    interpreter.allocate_tensors()
+    # Get input and output tensors and quantization details.
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+    image_tensor = interpreter.tensor(input_details[0]['index'])
+    logits_tensor = interpreter.tensor(output_details[0]['index'])
+
+    # Handle quantization.
+    scale = 1.0
+    zero_point = 0.0
+    input_dtype = tf.as_dtype(input_details[0]['dtype'])
+    if input_dtype.is_quantized or input_dtype.is_integer:
+      input_quantization = input_details[0]['quantization']
+      scale = input_quantization[0]
+      zero_point = input_quantization[1]
+    image_tensor()[0, :] = (eval_input.image.numpy() / scale) + zero_point
+
+    interpreter.invoke()
+    return eval_input.label.numpy() == np.argmax(logits_tensor()[0])
+
+  def evaluate_all(self) -> Tuple[int, int]:
+    """Evaluates all of images in the default dataset.
+
+    Returns:
+      Total number of evaluations and correct predictions as tuple of ints.
+    """
+    num_evals = 0
+    num_corrects = 0
+    for image_batch, label_batch in self._dataset:
+      inputs = [
+          EvaluationInput(image, label)
+          for image, label in zip(image_batch, label_batch)
+      ]
+      pool = mp.ThreadPool(self._num_threads)
+      results = pool.map(self.evaluate_single_image, inputs)
+      pool.close()
+      pool.join()
+      num_evals += len(results)
+      num_corrects += results.count(True)
+      accuracy = 100.0 * num_corrects / num_evals if num_evals > 0 else 0
+      logging.info('Evaluated: %d, Correct: %d, Accuracy: %f', num_evals,
+                   num_corrects, accuracy)
+    return (num_evals, num_corrects)

official/projects/edgetpu/vision/serving/tflite_imagenet_evaluator_run.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.
+
+r"""Evaluates image classification accuracy using tflite_imagenet_evaluator.
+
+Usage:
+tflite_imagenet_evaluator_run --tflite_model_path=/PATH/TO/MODEL.tflite
+"""
+
+from typing import Sequence
+from absl import app
+from absl import flags
+import tensorflow as tf
+
+from official.core import exp_factory
+from official.projects.edgetpu.vision.serving import tflite_imagenet_evaluator
+from official.projects.edgetpu.vision.tasks import image_classification
+
+
+flags.DEFINE_string('tflite_model_path', None,
+                    'Path to the tflite file to be evaluated.')
+flags.DEFINE_integer('num_threads', 16, 'Number of local threads.')
+flags.DEFINE_integer('batch_size', 256, 'Batch size per thread.')
+flags.DEFINE_string(
+    'model_name', 'mobilenet_edgetpu_v2_xs',
+    'Model name to identify a registered data pipeline setup and use as the '
+    'validation dataset.')
+
+FLAGS = flags.FLAGS
+
+
+def main(argv: Sequence[str]):
+  if len(argv) > 1:
+    raise app.UsageError('Too many command-line arguments.')
+
+  with tf.io.gfile.GFile(FLAGS.tflite_model_path, 'rb') as f:
+    model_content = f.read()
+
+  config = exp_factory.get_exp_config(FLAGS.model_name)
+  global_batch_size = FLAGS.num_threads * FLAGS.batch_size
+  config.task.validation_data.global_batch_size = global_batch_size
+  config.task.validation_data.dtype = 'float32'
+
+  task = image_classification.EdgeTPUTask(config.task)
+  dataset = task.build_inputs(config.task.validation_data)
+
+  evaluator = tflite_imagenet_evaluator.AccuracyEvaluator(
+      model_content=model_content,
+      dataset=dataset,
+      num_threads=FLAGS.num_threads)
+
+  evals, corrects = evaluator.evaluate_all()
+  accuracy = 100.0 * corrects / evals if evals > 0 else 0
+  print('Final accuracy: {}, Evaluated: {}, Correct: {} '.format(
+      accuracy, evals, corrects))
+
+
+if __name__ == '__main__':
+  flags.mark_flag_as_required('tflite_model_path')
+  app.run(main)

official/nlp/keras_nlp/encoders/bert_encoder.py → official/projects/edgetpu/vision/serving/tflite_imagenet_evaluator_test.py

@@ -12,19 +12,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Bert encoder network."""
-# pylint: disable=g-classes-have-attributes
+"""Tests for tflite_imagenet_evaluator."""
 
+from unittest import mock
 import tensorflow as tf
 
-from official.nlp.modeling import networks
+from official.projects.edgetpu.vision.serving import tflite_imagenet_evaluator
 
 
-@tf.keras.utils.register_keras_serializable(package='keras_nlp')
-class BertEncoder(networks.BertEncoder):
-  """Deprecated."""
+class TfliteImagenetEvaluatorTest(tf.test.TestCase):
 
-  def __init__(self, *args, **kwargs):
-    if 'dict_outputs' in kwargs:
-      kwargs.pop('dict_outputs')
-    super().__init__(*args, dict_outputs=True, **kwargs)
+  # Only tests the parallelization aspect. Mocks image evaluation and dataset.
+  def test_evaluate_all(self):
+    batch_size = 8
+    num_threads = 4
+    num_batches = 5
+
+    labels = tf.data.Dataset.range(batch_size * num_threads * num_batches)
+    images = tf.data.Dataset.range(batch_size * num_threads * num_batches)
+    dataset = tf.data.Dataset.zip((images, labels))
+    dataset = dataset.batch(batch_size)
+
+    with mock.patch.object(
+        tflite_imagenet_evaluator.AccuracyEvaluator,
+        'evaluate_single_image',
+        return_value=True,
+        autospec=True):
+      evaluator = tflite_imagenet_evaluator.AccuracyEvaluator(
+          model_content='MockModelContent'.encode('utf-8'),
+          dataset=dataset,
+          num_threads=num_threads)
+      num_evals, num_corrects = evaluator.evaluate_all()
+
+    expected_evals = num_batches * num_threads * batch_size
+
+    self.assertEqual(num_evals, expected_evals)
+    self.assertEqual(num_corrects, expected_evals)
+
+
+if __name__ == '__main__':
+  tf.test.main()

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)

official/vision/beta/dataloaders/classification_input.py

@@ -253,3 +253,21 @@ class Parser(parser.Parser):
     image = tf.image.convert_image_dtype(image, self._dtype)
 
     return image
+
+  @classmethod
+  def inference_fn(cls,
+                   image: tf.Tensor,
+                   input_image_size: List[int],
+                   num_channels: int = 3) -> tf.Tensor:
+    """Builds image model inputs for serving."""
+
+    image = tf.cast(image, dtype=tf.float32)
+    image = preprocess_ops.center_crop_image(image)
+    image = tf.image.resize(
+        image, input_image_size, method=tf.image.ResizeMethod.BILINEAR)
+
+    # Normalizes image with mean and std pixel values.
+    image = preprocess_ops.normalize_image(
+        image, offset=MEAN_RGB, scale=STDDEV_RGB)
+    image.set_shape(input_image_size + [num_channels])
+    return image

official/vision/beta/dataloaders/parser.py

@@ -67,3 +67,15 @@ class Parser(object):
         return self._parse_eval_data(decoded_tensors)
 
     return parse
+
+  @classmethod
+  def inference_fn(cls, inputs):
+    """Parses inputs for predictions.
+
+    Args:
+      inputs: A Tensor, or dictionary of Tensors.
+
+    Returns:
+      processed_inputs: An input tensor to the model.
+    """
+    pass

official/vision/beta/modeling/backbones/spinenet.py

@@ -60,6 +60,7 @@ SPINENET_BLOCK_SPECS = [
     (5, 'bottleneck', (7, 12), True),
     (7, 'bottleneck', (5, 14), True),
     (6, 'bottleneck', (12, 14), True),
+    (2, 'bottleneck', (2, 13), True),
 ]
 
 SCALING_MAP = {

official/vision/beta/modeling/layers/nn_layers.py

@@ -970,3 +970,213 @@ class Conv3D(tf.keras.layers.Conv3D, CausalConvMixin):
     """Computes the spatial output shape from the input shape."""
     shape = super(Conv3D, self)._spatial_output_shape(spatial_input_shape)
     return self._buffered_spatial_output_shape(shape)
+
+
+@tf.keras.utils.register_keras_serializable(package='Vision')
+class SpatialPyramidPooling(tf.keras.layers.Layer):
+  """Implements the Atrous Spatial Pyramid Pooling.
+
+  References:
+    [Rethinking Atrous Convolution for Semantic Image Segmentation](
+      https://arxiv.org/pdf/1706.05587.pdf)
+    [Encoder-Decoder with Atrous Separable Convolution for Semantic Image
+    Segmentation](https://arxiv.org/pdf/1802.02611.pdf)
+  """
+
+  def __init__(
+      self,
+      output_channels: int,
+      dilation_rates: List[int],
+      pool_kernel_size: Optional[List[int]] = None,
+      use_sync_bn: bool = False,
+      batchnorm_momentum: float = 0.99,
+      batchnorm_epsilon: float = 0.001,
+      activation: str = 'relu',
+      dropout: float = 0.5,
+      kernel_initializer: str = 'GlorotUniform',
+      kernel_regularizer: Optional[tf.keras.regularizers.Regularizer] = None,
+      interpolation: str = 'bilinear',
+      use_depthwise_convolution: bool = False,
+      **kwargs):
+    """Initializes `SpatialPyramidPooling`.
+
+    Args:
+      output_channels: Number of channels produced by SpatialPyramidPooling.
+      dilation_rates: A list of integers for parallel dilated conv.
+      pool_kernel_size: A list of integers or None. If None, global average
+        pooling is applied, otherwise an average pooling of pool_kernel_size is
+        applied.
+      use_sync_bn: A bool, whether or not to use sync batch normalization.
+      batchnorm_momentum: A float for the momentum in BatchNorm. Defaults to
+        0.99.
+      batchnorm_epsilon: A float for the epsilon value in BatchNorm. Defaults to
+        0.001.
+      activation: A `str` for type of activation to be used. Defaults to 'relu'.
+      dropout: A float for the dropout rate before output. Defaults to 0.5.
+      kernel_initializer: Kernel initializer for conv layers. Defaults to
+        `glorot_uniform`.
+      kernel_regularizer: Kernel regularizer for conv layers. Defaults to None.
+      interpolation: The interpolation method for upsampling. Defaults to
+        `bilinear`.
+      use_depthwise_convolution: Allows spatial pooling to be separable
+        depthwise convolusions. [Encoder-Decoder with Atrous Separable
+        Convolution for Semantic Image Segmentation](
+         https://arxiv.org/pdf/1802.02611.pdf)
+      **kwargs: Other keyword arguments for the layer.
+    """
+    super().__init__(**kwargs)
+
+    self._output_channels = output_channels
+    self._dilation_rates = dilation_rates
+    self._use_sync_bn = use_sync_bn
+    self._batchnorm_momentum = batchnorm_momentum
+    self._batchnorm_epsilon = batchnorm_epsilon
+    self._activation = activation
+    self._dropout = dropout
+    self._kernel_initializer = kernel_initializer
+    self._kernel_regularizer = kernel_regularizer
+    self._interpolation = interpolation
+    self._pool_kernel_size = pool_kernel_size
+    self._use_depthwise_convolution = use_depthwise_convolution
+    self._activation_fn = tf_utils.get_activation(activation)
+    if self._use_sync_bn:
+      self._bn_op = tf.keras.layers.experimental.SyncBatchNormalization
+    else:
+      self._bn_op = tf.keras.layers.BatchNormalization
+
+    if tf.keras.backend.image_data_format() == 'channels_last':
+      self._bn_axis = -1
+    else:
+      self._bn_axis = 1
+
+  def build(self, input_shape):
+    height = input_shape[1]
+    width = input_shape[2]
+    channels = input_shape[3]
+
+    self.aspp_layers = []
+
+    conv1 = tf.keras.layers.Conv2D(
+        filters=self._output_channels,
+        kernel_size=(1, 1),
+        kernel_initializer=self._kernel_initializer,
+        kernel_regularizer=self._kernel_regularizer,
+        use_bias=False)
+    norm1 = self._bn_op(
+        axis=self._bn_axis,
+        momentum=self._batchnorm_momentum,
+        epsilon=self._batchnorm_epsilon)
+
+    self.aspp_layers.append([conv1, norm1])
+
+    for dilation_rate in self._dilation_rates:
+      leading_layers = []
+      kernel_size = (3, 3)
+      if self._use_depthwise_convolution:
+        leading_layers += [
+            tf.keras.layers.DepthwiseConv2D(
+                depth_multiplier=1,
+                kernel_size=kernel_size,
+                padding='same',
+                depthwise_regularizer=self._kernel_regularizer,
+                depthwise_initializer=self._kernel_initializer,
+                dilation_rate=dilation_rate,
+                use_bias=False)
+        ]
+        kernel_size = (1, 1)
+      conv_dilation = leading_layers + [
+          tf.keras.layers.Conv2D(
+              filters=self._output_channels,
+              kernel_size=kernel_size,
+              padding='same',
+              kernel_regularizer=self._kernel_regularizer,
+              kernel_initializer=self._kernel_initializer,
+              dilation_rate=dilation_rate,
+              use_bias=False)
+      ]
+      norm_dilation = self._bn_op(
+          axis=self._bn_axis,
+          momentum=self._batchnorm_momentum,
+          epsilon=self._batchnorm_epsilon)
+
+      self.aspp_layers.append(conv_dilation + [norm_dilation])
+
+    if self._pool_kernel_size is None:
+      pooling = [
+          tf.keras.layers.GlobalAveragePooling2D(),
+          tf.keras.layers.Reshape((1, 1, channels))
+      ]
+    else:
+      pooling = [tf.keras.layers.AveragePooling2D(self._pool_kernel_size)]
+
+    conv2 = tf.keras.layers.Conv2D(
+        filters=self._output_channels,
+        kernel_size=(1, 1),
+        kernel_initializer=self._kernel_initializer,
+        kernel_regularizer=self._kernel_regularizer,
+        use_bias=False)
+    norm2 = self._bn_op(
+        axis=self._bn_axis,
+        momentum=self._batchnorm_momentum,
+        epsilon=self._batchnorm_epsilon)
+
+    self.aspp_layers.append(pooling + [conv2, norm2])
+
+    self._resize_layer = tf.keras.layers.Resizing(
+        height, width, interpolation=self._interpolation, dtype=tf.float32)
+
+    self._projection = [
+        tf.keras.layers.Conv2D(
+            filters=self._output_channels,
+            kernel_size=(1, 1),
+            kernel_initializer=self._kernel_initializer,
+            kernel_regularizer=self._kernel_regularizer,
+            use_bias=False),
+        self._bn_op(
+            axis=self._bn_axis,
+            momentum=self._batchnorm_momentum,
+            epsilon=self._batchnorm_epsilon)
+    ]
+    self._dropout_layer = tf.keras.layers.Dropout(rate=self._dropout)
+    self._concat_layer = tf.keras.layers.Concatenate(axis=-1)
+
+  def call(self,
+           inputs: tf.Tensor,
+           training: Optional[bool] = None) -> tf.Tensor:
+    if training is None:
+      training = tf.keras.backend.learning_phase()
+    result = []
+    for i, layers in enumerate(self.aspp_layers):
+      x = inputs
+      for layer in layers:
+        # Apply layers sequentially.
+        x = layer(x, training=training)
+      x = self._activation_fn(x)
+
+      # Apply resize layer to the end of the last set of layers.
+      if i == len(self.aspp_layers) - 1:
+        x = self._resize_layer(x)
+
+      result.append(tf.cast(x, inputs.dtype))
+    x = self._concat_layer(result)
+    for layer in self._projection:
+      x = layer(x, training=training)
+    x = self._activation_fn(x)
+    return self._dropout_layer(x)
+
+  def get_config(self):
+    config = {
+        'output_channels': self._output_channels,
+        'dilation_rates': self._dilation_rates,
+        'pool_kernel_size': self._pool_kernel_size,
+        'use_sync_bn': self._use_sync_bn,
+        'batchnorm_momentum': self._batchnorm_momentum,
+        'batchnorm_epsilon': self._batchnorm_epsilon,
+        'activation': self._activation,
+        'dropout': self._dropout,
+        'kernel_initializer': self._kernel_initializer,
+        'kernel_regularizer': self._kernel_regularizer,
+        'interpolation': self._interpolation,
+    }
+    base_config = super().get_config()
+    return dict(list(base_config.items()) + list(config.items()))

official/vision/beta/modeling/layers/nn_layers_test.py

@@ -406,5 +406,19 @@ class NNLayersTest(parameterized.TestCase, tf.test.TestCase):
                         [[[[[1.]]],
                           [[[3.]]]]])
 
+  @parameterized.parameters(
+      (None, []),
+      (None, [6, 12, 18]),
+      ([32, 32], [6, 12, 18]),
+  )
+  def test_aspp(self, pool_kernel_size, dilation_rates):
+    inputs = tf.keras.Input(shape=(64, 64, 128), dtype=tf.float32)
+    layer = nn_layers.SpatialPyramidPooling(
+        output_channels=256,
+        dilation_rates=dilation_rates,
+        pool_kernel_size=pool_kernel_size)
+    output = layer(inputs)
+    self.assertAllEqual([None, 64, 64, 256], output.shape)
+
 if __name__ == '__main__':
   tf.test.main()