Script to convert TF2 SSD models to TFLite, with documentation

PiperOrigin-RevId: 329938888

research/object_detection/README.md

@@ -73,6 +73,13 @@ documentation of the Object Detection API:
 
 ## Whats New
 
+### Mobile Inference for TF2 models
+
+TF2 OD API models can now be converted to TensorFlow Lite! Only SSD models
+currently supported. See <a href='running_on_mobile_tf2.md'>documentation</a>.
+
+**Thanks to contributors**: Sachin Joglekar
+
 ### TensorFlow 2 Support
 
 We are happy to announce that the TF OD API officially supports TF2! Our release

research/object_detection/export_tflite_graph_lib_tf2.py

+# Lint as: python3
+# Copyright 2020 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.
+# ==============================================================================
+"""Library to export TFLite-compatible SavedModel from TF2 detection models."""
+import os
+import numpy as np
+import tensorflow.compat.v1 as tf1
+import tensorflow.compat.v2 as tf
+
+from object_detection.builders import model_builder
+from object_detection.builders import post_processing_builder
+from object_detection.core import box_list
+
+_DEFAULT_NUM_CHANNELS = 3
+_DEFAULT_NUM_COORD_BOX = 4
+_MAX_CLASSES_PER_DETECTION = 1
+_DETECTION_POSTPROCESS_FUNC = 'TFLite_Detection_PostProcess'
+
+
+def get_const_center_size_encoded_anchors(anchors):
+  """Exports center-size encoded anchors as a constant tensor.
+
+  Args:
+    anchors: a float32 tensor of shape [num_anchors, 4] containing the anchor
+      boxes
+
+  Returns:
+    encoded_anchors: a float32 constant tensor of shape [num_anchors, 4]
+    containing the anchor boxes.
+  """
+  anchor_boxlist = box_list.BoxList(anchors)
+  y, x, h, w = anchor_boxlist.get_center_coordinates_and_sizes()
+  num_anchors = y.get_shape().as_list()
+
+  with tf1.Session() as sess:
+    y_out, x_out, h_out, w_out = sess.run([y, x, h, w])
+  encoded_anchors = tf1.constant(
+      np.transpose(np.stack((y_out, x_out, h_out, w_out))),
+      dtype=tf1.float32,
+      shape=[num_anchors[0], _DEFAULT_NUM_COORD_BOX],
+      name='anchors')
+  return num_anchors[0], encoded_anchors
+
+
+class SSDModule(tf.Module):
+  """Inference Module for TFLite-friendly SSD models."""
+
+  def __init__(self, pipeline_config, detection_model, max_detections,
+               use_regular_nms):
+    """Initialization.
+
+    Args:
+      pipeline_config: The original pipeline_pb2.TrainEvalPipelineConfig
+      detection_model: The detection model to use for inference.
+      max_detections: Max detections desired from the TFLite model.
+      use_regular_nms: If True, TFLite model uses the (slower) multi-class NMS.
+    """
+    self._process_config(pipeline_config)
+    self._pipeline_config = pipeline_config
+    self._model = detection_model
+    self._max_detections = max_detections
+    self._use_regular_nms = use_regular_nms
+
+  def _process_config(self, pipeline_config):
+    self._num_classes = pipeline_config.model.ssd.num_classes
+    self._nms_score_threshold = pipeline_config.model.ssd.post_processing.batch_non_max_suppression.score_threshold
+    self._nms_iou_threshold = pipeline_config.model.ssd.post_processing.batch_non_max_suppression.iou_threshold
+    self._scale_values = {}
+    self._scale_values[
+        'y_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.y_scale
+    self._scale_values[
+        'x_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.x_scale
+    self._scale_values[
+        'h_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.height_scale
+    self._scale_values[
+        'w_scale'] = pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.width_scale
+
+    image_resizer_config = pipeline_config.model.ssd.image_resizer
+    image_resizer = image_resizer_config.WhichOneof('image_resizer_oneof')
+    self._num_channels = _DEFAULT_NUM_CHANNELS
+
+    if image_resizer == 'fixed_shape_resizer':
+      self._height = image_resizer_config.fixed_shape_resizer.height
+      self._width = image_resizer_config.fixed_shape_resizer.width
+      if image_resizer_config.fixed_shape_resizer.convert_to_grayscale:
+        self._num_channels = 1
+    else:
+      raise ValueError(
+          'Only fixed_shape_resizer'
+          'is supported with tflite. Found {}'.format(
+              image_resizer_config.WhichOneof('image_resizer_oneof')))
+
+  def input_shape(self):
+    """Returns shape of TFLite model input."""
+    return [1, self._height, self._width, self._num_channels]
+
+  def postprocess_implements_signature(self):
+    """Returns tf.implements signature for MLIR legalization of TFLite NMS."""
+    implements_signature = [
+        'name: "%s"' % _DETECTION_POSTPROCESS_FUNC,
+        'attr { key: "max_detections" value { i: %d } }' % self._max_detections,
+        'attr { key: "max_classes_per_detection" value { i: %d } }' %
+        _MAX_CLASSES_PER_DETECTION,
+        'attr { key: "use_regular_nms" value { b: %s } }' %
+        str(self._use_regular_nms).lower(),
+        'attr { key: "nms_score_threshold" value { f: %f } }' %
+        self._nms_score_threshold,
+        'attr { key: "nms_iou_threshold" value { f: %f } }' %
+        self._nms_iou_threshold,
+        'attr { key: "y_scale" value { f: %f } }' %
+        self._scale_values['y_scale'],
+        'attr { key: "x_scale" value { f: %f } }' %
+        self._scale_values['x_scale'],
+        'attr { key: "h_scale" value { f: %f } }' %
+        self._scale_values['h_scale'],
+        'attr { key: "w_scale" value { f: %f } }' %
+        self._scale_values['w_scale'],
+        'attr { key: "num_classes" value { i: %d } }' % self._num_classes
+    ]
+    implements_signature = ' '.join(implements_signature)
+    return implements_signature
+
+  def _get_postprocess_fn(self, num_anchors, num_classes):
+    # There is no TF equivalent for TFLite's custom post-processing op.
+    # So we add an 'empty' composite function here, that is legalized to the
+    # custom op with MLIR.
+    @tf.function(
+        experimental_implements=self.postprocess_implements_signature())
+    # pylint: disable=g-unused-argument,unused-argument
+    def dummy_post_processing(box_encodings, class_predictions, anchors):
+      boxes = tf.constant(0.0, dtype=tf.float32, name='boxes')
+      scores = tf.constant(0.0, dtype=tf.float32, name='scores')
+      classes = tf.constant(0.0, dtype=tf.float32, name='classes')
+      num_detections = tf.constant(0.0, dtype=tf.float32, name='num_detections')
+      return boxes, scores, classes, num_detections
+
+    return dummy_post_processing
+
+  @tf.function
+  def inference_fn(self, image):
+    """Encapsulates SSD inference for TFLite conversion.
+
+    NOTE: The Args & Returns sections below indicate the TFLite model signature,
+    and not what the TF graph does (since the latter does not include the custom
+    NMS op used by TFLite)
+
+    Args:
+      image: a float32 tensor of shape [num_anchors, 4] containing the anchor
+        boxes
+
+    Returns:
+      num_detections: a float32 scalar denoting number of total detections.
+      classes: a float32 tensor denoting class ID for each detection.
+      scores: a float32 tensor denoting score for each detection.
+      boxes: a float32 tensor denoting coordinates of each detected box.
+    """
+    predicted_tensors = self._model.predict(image, true_image_shapes=None)
+    # The score conversion occurs before the post-processing custom op
+    _, score_conversion_fn = post_processing_builder.build(
+        self._pipeline_config.model.ssd.post_processing)
+    class_predictions = score_conversion_fn(
+        predicted_tensors['class_predictions_with_background'])
+
+    with tf.name_scope('raw_outputs'):
+      # 'raw_outputs/box_encodings': a float32 tensor of shape
+      #   [1, num_anchors, 4] containing the encoded box predictions. Note that
+      #   these are raw predictions and no Non-Max suppression is applied on
+      #   them and no decode center size boxes is applied to them.
+      box_encodings = tf.identity(
+          predicted_tensors['box_encodings'], name='box_encodings')
+      # 'raw_outputs/class_predictions': a float32 tensor of shape
+      #   [1, num_anchors, num_classes] containing the class scores for each
+      #   anchor after applying score conversion.
+      class_predictions = tf.identity(
+          class_predictions, name='class_predictions')
+    # 'anchors': a float32 tensor of shape
+    #   [4, num_anchors] containing the anchors as a constant node.
+    num_anchors, anchors = get_const_center_size_encoded_anchors(
+        predicted_tensors['anchors'])
+    anchors = tf.identity(anchors, name='anchors')
+
+    # tf.function@ seems to reverse order of inputs, so reverse them here.
+    return self._get_postprocess_fn(num_anchors,
+                                    self._num_classes)(box_encodings,
+                                                       class_predictions,
+                                                       anchors)[::-1]
+
+
+def export_tflite_model(pipeline_config, trained_checkpoint_dir,
+                        output_directory, max_detections, use_regular_nms):
+  """Exports inference SavedModel for TFLite conversion.
+
+  NOTE: Only supports SSD meta-architectures for now, and the output model will
+  have static-shaped, single-batch input.
+
+  This function creates `output_directory` if it does not already exist,
+  which will hold the intermediate SavedModel that can be used with the TFLite
+  converter.
+
+  Args:
+    pipeline_config: pipeline_pb2.TrainAndEvalPipelineConfig proto.
+    trained_checkpoint_dir: Path to the trained checkpoint file.
+    output_directory: Path to write outputs.
+    max_detections: Max detections desired from the TFLite model.
+    use_regular_nms: If True, TFLite model uses the (slower) multi-class NMS.
+
+  Raises:
+    ValueError: if pipeline is invalid.
+  """
+  output_saved_model_directory = os.path.join(output_directory, 'saved_model')
+
+  # Build the underlying model using pipeline config.
+  # TODO(b/162842801): Add support for other architectures.
+  if pipeline_config.model.WhichOneof('model') != 'ssd':
+    raise ValueError('Only ssd models are supported in tflite. '
+                     'Found {} in config'.format(
+                         pipeline_config.model.WhichOneof('model')))
+  detection_model = model_builder.build(
+      pipeline_config.model, is_training=False)
+
+  ckpt = tf.train.Checkpoint(model=detection_model)
+  manager = tf.train.CheckpointManager(
+      ckpt, trained_checkpoint_dir, max_to_keep=1)
+  status = ckpt.restore(manager.latest_checkpoint).expect_partial()
+
+  # The module helps build a TF SavedModel appropriate for TFLite conversion.
+  detection_module = SSDModule(pipeline_config, detection_model, max_detections,
+                               use_regular_nms)
+
+  # Getting the concrete function traces the graph and forces variables to
+  # be constructed; only after this can we save the saved model.
+  status.assert_existing_objects_matched()
+  concrete_function = detection_module.inference_fn.get_concrete_function(
+      tf.TensorSpec(
+          shape=detection_module.input_shape(), dtype=tf.float32, name='input'))
+  status.assert_existing_objects_matched()
+
+  # Export SavedModel.
+  tf.saved_model.save(
+      detection_module,
+      output_saved_model_directory,
+      signatures=concrete_function)

research/object_detection/export_tflite_graph_lib_tf2_test.py

+# Lint as: python3
+# Copyright 2020 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.
+# ==============================================================================
+"""Test for export_tflite_graph_lib_tf2.py."""
+
+from __future__ import division
+import os
+import unittest
+import six
+
+import tensorflow.compat.v2 as tf
+
+from object_detection import export_tflite_graph_lib_tf2
+from object_detection.builders import model_builder
+from object_detection.core import model
+from object_detection.protos import pipeline_pb2
+from object_detection.utils import tf_version
+
+if six.PY2:
+  import mock  # pylint: disable=g-importing-member,g-import-not-at-top
+else:
+  from unittest import mock  # pylint: disable=g-importing-member,g-import-not-at-top
+
+
+class FakeModel(model.DetectionModel):
+
+  def __init__(self):
+    super(FakeModel, self).__init__(num_classes=2)
+    self._conv = tf.keras.layers.Conv2D(
+        filters=1,
+        kernel_size=1,
+        strides=(1, 1),
+        padding='valid',
+        kernel_initializer=tf.keras.initializers.Constant(value=1.0))
+
+  def preprocess(self, inputs):
+    true_image_shapes = []  # Doesn't matter for the fake model.
+    return tf.identity(inputs), true_image_shapes
+
+  def predict(self, preprocessed_inputs, true_image_shapes):
+    prediction_tensors = {'image': self._conv(preprocessed_inputs)}
+    with tf.control_dependencies([prediction_tensors['image']]):
+      prediction_tensors['box_encodings'] = tf.constant(
+          [[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]]], tf.float32)
+      prediction_tensors['class_predictions_with_background'] = tf.constant(
+          [[[0.7, 0.6], [0.9, 0.0]]], tf.float32)
+    with tf.control_dependencies([
+        tf.convert_to_tensor(
+            prediction_tensors['image'].get_shape().as_list()[1:3])
+    ]):
+      prediction_tensors['anchors'] = tf.constant(
+          [[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 1.0, 1.0]], tf.float32)
+    return prediction_tensors
+
+  def postprocess(self, prediction_dict, true_image_shapes):
+    predict_tensor_sum = tf.reduce_sum(prediction_dict['image'])
+    with tf.control_dependencies(list(prediction_dict.values())):
+      postprocessed_tensors = {
+          'detection_boxes':
+              tf.constant([[[0.0, 0.0, 0.5, 0.5], [0.5, 0.5, 0.8, 0.8]],
+                           [[0.5, 0.5, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]]],
+                          tf.float32),
+          'detection_scores':
+              predict_tensor_sum +
+              tf.constant([[0.7, 0.6], [0.9, 0.0]], tf.float32),
+          'detection_classes':
+              tf.constant([[0, 1], [1, 0]], tf.float32),
+          'num_detections':
+              tf.constant([2, 1], tf.float32),
+      }
+    return postprocessed_tensors
+
+  def restore_map(self, checkpoint_path, from_detection_checkpoint):
+    pass
+
+  def restore_from_objects(self, fine_tune_checkpoint_type):
+    pass
+
+  def loss(self, prediction_dict, true_image_shapes):
+    pass
+
+  def regularization_losses(self):
+    pass
+
+  def updates(self):
+    pass
+
+
+@unittest.skipIf(tf_version.is_tf1(), 'Skipping TF2.X only test.')
+class ExportTfLiteGraphTest(tf.test.TestCase):
+
+  def _save_checkpoint_from_mock_model(self, checkpoint_dir):
+    mock_model = FakeModel()
+    fake_image = tf.zeros(shape=[1, 10, 10, 3], dtype=tf.float32)
+    preprocessed_inputs, true_image_shapes = mock_model.preprocess(fake_image)
+    predictions = mock_model.predict(preprocessed_inputs, true_image_shapes)
+    mock_model.postprocess(predictions, true_image_shapes)
+
+    ckpt = tf.train.Checkpoint(model=mock_model)
+    exported_checkpoint_manager = tf.train.CheckpointManager(
+        ckpt, checkpoint_dir, max_to_keep=1)
+    exported_checkpoint_manager.save(checkpoint_number=0)
+
+  def _get_ssd_config(self):
+    pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
+    pipeline_config.model.ssd.image_resizer.fixed_shape_resizer.height = 10
+    pipeline_config.model.ssd.image_resizer.fixed_shape_resizer.width = 10
+    pipeline_config.model.ssd.num_classes = 2
+    pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.y_scale = 10.0
+    pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.x_scale = 10.0
+    pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.height_scale = 5.0
+    pipeline_config.model.ssd.box_coder.faster_rcnn_box_coder.width_scale = 5.0
+    pipeline_config.model.ssd.post_processing.batch_non_max_suppression.iou_threshold = 0.5
+    return pipeline_config
+
+  # The tf.implements signature is important since it ensures MLIR legalization,
+  # so we test it here.
+  def test_postprocess_implements_signature(self):
+    tmp_dir = self.get_temp_dir()
+    self._save_checkpoint_from_mock_model(tmp_dir)
+    pipeline_config = self._get_ssd_config()
+
+    with mock.patch.object(
+        model_builder, 'build', autospec=True) as mock_builder:
+      mock_builder.return_value = FakeModel()
+
+      detection_model = model_builder.build(
+          pipeline_config.model, is_training=False)
+
+      ckpt = tf.train.Checkpoint(model=detection_model)
+      manager = tf.train.CheckpointManager(ckpt, tmp_dir, max_to_keep=1)
+      ckpt.restore(manager.latest_checkpoint).expect_partial()
+
+      # The module helps build a TF graph appropriate for TFLite conversion.
+      detection_module = export_tflite_graph_lib_tf2.SSDModule(
+          pipeline_config=pipeline_config,
+          detection_model=detection_model,
+          max_detections=20,
+          use_regular_nms=True)
+
+      expected_signature = ('name: "TFLite_Detection_PostProcess" attr { key: '
+                            '"max_detections" value { i: 20 } } attr { key: '
+                            '"max_classes_per_detection" value { i: 1 } } attr '
+                            '{ key: "use_regular_nms" value { b: true } } attr '
+                            '{ key: "nms_score_threshold" value { f: 0.000000 }'
+                            ' } attr { key: "nms_iou_threshold" value { f: '
+                            '0.500000 } } attr { key: "y_scale" value { f: '
+                            '10.000000 } } attr { key: "x_scale" value { f: '
+                            '10.000000 } } attr { key: "h_scale" value { f: '
+                            '5.000000 } } attr { key: "w_scale" value { f: '
+                            '5.000000 } } attr { key: "num_classes" value { i: '
+                            '2 } }')
+
+      self.assertEqual(expected_signature,
+                       detection_module.postprocess_implements_signature())
+
+  def test_unsupported_architecture(self):
+    tmp_dir = self.get_temp_dir()
+    self._save_checkpoint_from_mock_model(tmp_dir)
+
+    pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
+    pipeline_config.model.faster_rcnn.num_classes = 10
+
+    with mock.patch.object(
+        model_builder, 'build', autospec=True) as mock_builder:
+      mock_builder.return_value = FakeModel()
+      output_directory = os.path.join(tmp_dir, 'output')
+      expected_message = 'Only ssd models are supported in tflite'
+      try:
+        export_tflite_graph_lib_tf2.export_tflite_model(
+            pipeline_config=pipeline_config,
+            trained_checkpoint_dir=tmp_dir,
+            output_directory=output_directory,
+            max_detections=10,
+            use_regular_nms=False)
+      except ValueError as e:
+        if expected_message not in str(e):
+          raise
+      else:
+        raise AssertionError('Exception not raised: %s' % expected_message)
+
+  def test_export_yields_saved_model(self):
+    tmp_dir = self.get_temp_dir()
+    self._save_checkpoint_from_mock_model(tmp_dir)
+    with mock.patch.object(
+        model_builder, 'build', autospec=True) as mock_builder:
+      mock_builder.return_value = FakeModel()
+      output_directory = os.path.join(tmp_dir, 'output')
+      export_tflite_graph_lib_tf2.export_tflite_model(
+          pipeline_config=self._get_ssd_config(),
+          trained_checkpoint_dir=tmp_dir,
+          output_directory=output_directory,
+          max_detections=10,
+          use_regular_nms=False)
+      self.assertTrue(
+          os.path.exists(
+              os.path.join(output_directory, 'saved_model', 'saved_model.pb')))
+      self.assertTrue(
+          os.path.exists(
+              os.path.join(output_directory, 'saved_model', 'variables',
+                           'variables.index')))
+      self.assertTrue(
+          os.path.exists(
+              os.path.join(output_directory, 'saved_model', 'variables',
+                           'variables.data-00000-of-00001')))
+
+  def test_exported_model_inference(self):
+    tmp_dir = self.get_temp_dir()
+    output_directory = os.path.join(tmp_dir, 'output')
+    self._save_checkpoint_from_mock_model(tmp_dir)
+    with mock.patch.object(
+        model_builder, 'build', autospec=True) as mock_builder:
+      mock_builder.return_value = FakeModel()
+      export_tflite_graph_lib_tf2.export_tflite_model(
+          pipeline_config=self._get_ssd_config(),
+          trained_checkpoint_dir=tmp_dir,
+          output_directory=output_directory,
+          max_detections=10,
+          use_regular_nms=False)
+
+    saved_model_path = os.path.join(output_directory, 'saved_model')
+    detect_fn = tf.saved_model.load(saved_model_path)
+    detect_fn_sig = detect_fn.signatures['serving_default']
+    image = tf.zeros(shape=[1, 10, 10, 3], dtype=tf.float32)
+    detections = detect_fn_sig(image)
+
+    # The exported graph doesn't have numerically correct outputs, but there
+    # should be 4.
+    self.assertEqual(4, len(detections))
+
+
+if __name__ == '__main__':
+  tf.test.main()

research/object_detection/export_tflite_graph_tf2.py

+# Lint as: python2, python3
+# Copyright 2020 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"""Exports TF2 detection SavedModel for conversion to TensorFlow Lite.
+
+Link to the TF2 Detection Zoo:
+https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md
+The output folder will contain an intermediate SavedModel that can be used with
+the TfLite converter.
+
+NOTE: This only supports SSD meta-architectures for now.
+
+One input:
+  image: a float32 tensor of shape[1, height, width, 3] containing the
+  *normalized* input image.
+  NOTE: See the `preprocess` function defined in the feature extractor class
+  in the object_detection/models directory.
+
+Four Outputs:
+  detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box
+  locations
+  detection_classes: a float32 tensor of shape [1, num_boxes]
+  with class indices
+  detection_scores: a float32 tensor of shape [1, num_boxes]
+  with class scores
+  num_boxes: a float32 tensor of size 1 containing the number of detected boxes
+
+Example Usage:
+--------------
+python object_detection/export_tflite_graph_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_prefix path/to/ssd_model/checkpoint \
+    --output_directory path/to/exported_model_directory
+
+The expected output SavedModel would be in the directory
+path/to/exported_model_directory (which is created if it does not exist).
+
+Config overrides (see the `config_override` flag) are text protobufs
+(also of type pipeline_pb2.TrainEvalPipelineConfig) which are used to override
+certain fields in the provided pipeline_config_path.  These are useful for
+making small changes to the inference graph that differ from the training or
+eval config.
+
+Example Usage (in which we change the NMS iou_threshold to be 0.5 and
+NMS score_threshold to be 0.0):
+python object_detection/export_tflite_model_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_prefix path/to/ssd_model/checkpoint \
+    --output_directory path/to/exported_model_directory
+    --config_override " \
+            model{ \
+            ssd{ \
+              post_processing { \
+                batch_non_max_suppression { \
+                        score_threshold: 0.0 \
+                        iou_threshold: 0.5 \
+                } \
+             } \
+          } \
+       } \
+       "
+"""
+from absl import app
+from absl import flags
+
+import tensorflow.compat.v2 as tf
+from google.protobuf import text_format
+from object_detection import export_tflite_graph_lib_tf2
+from object_detection.protos import pipeline_pb2
+
+tf.enable_v2_behavior()
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'pipeline_config_path', None,
+    'Path to a pipeline_pb2.TrainEvalPipelineConfig config '
+    'file.')
+flags.DEFINE_string('trained_checkpoint_dir', None,
+                    'Path to trained checkpoint directory')
+flags.DEFINE_string('output_directory', None, 'Path to write outputs.')
+flags.DEFINE_string(
+    'config_override', '', 'pipeline_pb2.TrainEvalPipelineConfig '
+    'text proto to override pipeline_config_path.')
+# SSD-specific flags
+flags.DEFINE_integer('ssd_max_detections', 10,
+                     'Maximum number of detections (boxes) to return.')
+flags.DEFINE_bool(
+    'ssd_use_regular_nms', False,
+    'Flag to set postprocessing op to use Regular NMS instead of Fast NMS '
+    '(Default false).')
+
+
+def main(argv):
+  del argv  # Unused.
+  flags.mark_flag_as_required('pipeline_config_path')
+  flags.mark_flag_as_required('trained_checkpoint_dir')
+  flags.mark_flag_as_required('output_directory')
+
+  pipeline_config = pipeline_pb2.TrainEvalPipelineConfig()
+
+  with tf.io.gfile.GFile(FLAGS.pipeline_config_path, 'r') as f:
+    text_format.Parse(f.read(), pipeline_config)
+  text_format.Parse(FLAGS.config_override, pipeline_config)
+
+  export_tflite_graph_lib_tf2.export_tflite_model(pipeline_config,
+                                                  FLAGS.trained_checkpoint_dir,
+                                                  FLAGS.output_directory,
+                                                  FLAGS.ssd_max_detections,
+                                                  FLAGS.ssd_use_regular_nms)
+
+
+if __name__ == '__main__':
+  app.run(main)

research/object_detection/g3doc/release_notes.md

 # Release Notes
 
+### September 3rd, 2020
+
+TF2 OD API models can now be converted to TensorFlow Lite! Only SSD models
+currently supported. See <a href='running_on_mobile_tf2.md'>documentation</a>.
+
+**Thanks to contributors**: Sachin Joglekar
+
 ### July 10th, 2020
 
 We are happy to announce that the TF OD API officially supports TF2! Our release

research/object_detection/g3doc/running_on_mobile_tf2.md

+# Running TF2 Detection API Models on mobile
+
+[![TensorFlow 2.2](https://img.shields.io/badge/TensorFlow-2.2-FF6F00?logo=tensorflow)](https://github.com/tensorflow/tensorflow/releases/tag/v2.2.0)
+[![Python 3.6](https://img.shields.io/badge/Python-3.6-3776AB)](https://www.python.org/downloads/release/python-360/)
+
+[TensorFlow Lite](https://www.tensorflow.org/mobile/tflite/)(TFLite) is
+TensorFlow’s lightweight solution for mobile and embedded devices. It enables
+on-device machine learning inference with low latency and a small binary size.
+TensorFlow Lite uses many techniques for this such as quantized kernels that
+allow smaller and faster (fixed-point math) models.
+
+This document shows how elgible models from the
+[TF2 Detection zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md)
+can be converted for inference with TFLite.
+
+**NOTE:** TFLite currently only supports **SSD Architectures** (excluding
+EfficientDet) for boxes-based detection. Support for EfficientDet is coming
+soon.
+
+The output model has the following inputs & outputs:
+
+```
+One input:
+  image: a float32 tensor of shape[1, height, width, 3] containing the
+  *normalized* input image.
+  NOTE: See the `preprocess` function defined in the feature extractor class
+  in the object_detection/models directory.
+
+Four Outputs:
+  detection_boxes: a float32 tensor of shape [1, num_boxes, 4] with box
+  locations
+  detection_classes: a float32 tensor of shape [1, num_boxes]
+  with class indices
+  detection_scores: a float32 tensor of shape [1, num_boxes]
+  with class scores
+  num_boxes: a float32 tensor of size 1 containing the number of detected boxes
+```
+
+There are two steps to TFLite conversion:
+
+### Step 1: Export TFLite inference graph
+
+This step generates an intermediate SavedModel that can be used with the
+[TFLite Converter](https://www.tensorflow.org/lite/convert) via commandline or
+Python API.
+
+To use the script:
+
+```bash
+# From the tensorflow/models/research/ directory
+python object_detection/export_tflite_graph_tf2.py \
+    --pipeline_config_path path/to/ssd_model/pipeline.config \
+    --trained_checkpoint_prefix path/to/ssd_model/checkpoint \
+    --output_directory path/to/exported_model_directory
+```
+
+Use `--help` with the aboev script to get the full list of supported parameters.
+
+### Step 2: Convert to TFLite
+
+Use the [TensorFlow Lite Converter](https://www.tensorflow.org/lite/convert) to
+convert the `SavedModel` to TFLite. You can also leverage
+[Post-training Quantization](https://www.tensorflow.org/lite/performance/post_training_quantization)
+to
+[optimize performance](https://www.tensorflow.org/lite/performance/model_optimization)
+and obtain a smaller model.
+
+## Running our model on Android
+
+To run our TensorFlow Lite model on device, we will use Android Studio to build
+and run the TensorFlow Lite detection example with the new model. The example is
+found in the
+[TensorFlow examples repository](https://github.com/tensorflow/examples) under
+`/lite/examples/object_detection`. The example can be built with
+[Android Studio](https://developer.android.com/studio/index.html), and requires
+the
+[Android SDK with build tools](https://developer.android.com/tools/revisions/build-tools.html)
+that support API >= 21. Additional details are available on the
+[TensorFlow Lite example page](https://github.com/tensorflow/examples/tree/master/lite/examples/object_detection/android).
+
+Next we need to point the app to our new detect.tflite file and give it the
+names of our new labels. Specifically, we will copy our TensorFlow Lite
+flatbuffer to the app assets directory with the following command:
+
+```shell
+mkdir $TF_EXAMPLES/lite/examples/object_detection/android/app/src/main/assets
+cp /tmp/tflite/detect.tflite \
+  $TF_EXAMPLES/lite/examples/object_detection/android/app/src/main/assets
+```
+
+You will also need to copy your new labelmap labelmap.txt to the assets
+directory.
+
+We will now edit the gradle build file to use these assets. First, open the
+`build.gradle` file
+`$TF_EXAMPLES/lite/examples/object_detection/android/app/build.gradle`. Comment
+out the model download script to avoid your assets being overwritten: `// apply
+from:'download_model.gradle'` ```
+
+If your model is named `detect.tflite`, and your labels file `labelmap.txt`, the
+example will use them automatically as long as they've been properly copied into
+the base assets directory. If you need to use a custom path or filename, open up
+the
+$TF_EXAMPLES/lite/examples/object_detection/android/app/src/main/java/org/tensorflow/demo/DetectorActivity.java
+file in a text editor and find the definition of TF_OD_API_LABELS_FILE. Update
+this path to point to your new label map file: "labels_list.txt". Note that if
+your model is quantized, the flag TF_OD_API_IS_QUANTIZED is set to true, and if
+your model is floating point, the flag TF_OD_API_IS_QUANTIZED is set to false.
+This new section of DetectorActivity.java should now look as follows for a
+quantized model:
+
+```shell
+  private static final boolean TF_OD_API_IS_QUANTIZED = true;
+  private static final String TF_OD_API_MODEL_FILE = "detect.tflite";
+  private static final String TF_OD_API_LABELS_FILE = "labels_list.txt";
+```
+
+Once you’ve copied the TensorFlow Lite model and edited the gradle build script
+to not use the downloaded assets, you can build and deploy the app using the
+usual Android Studio build process.

research/object_detection/g3doc/tf2_detection_zoo.md

@@ -15,6 +15,8 @@ They are also useful for initializing your models when training on novel
 datasets. You can try this out on our few-shot training
 [colab](../colab_tutorials/eager_few_shot_od_training_tf2_colab.ipynb).
 
+Please look at [this guide](../running_on_mobile_tf2.md) for mobile inference.
+
 <!-- mdlint on -->
 
 Finally, if you would like to train these models from scratch, you can find the