Add export_saved_model and export_tflite for object detection in Model Maker

PiperOrigin-RevId: 352517354

Author: ziyeqinghan

tensorflow_examples/lite/model_maker/core/task/configs.py

@@ -179,13 +179,13 @@ def create_float16_quantization(cls, optimizations=tf.lite.Optimize.DEFAULT):
     """Creates configuration for float16 quantization."""
     return QuantizationConfig(optimizations, supported_types=[tf.float16])
 
-  def get_converter_with_quantization(self, converter, preprocess=None):
+  def get_converter_with_quantization(self, converter, **kwargs):
     """Gets TFLite converter with settings for quantization."""
     converter.optimizations = self.optimizations
 
     if self.representative_data is not None:
       ds = self.representative_data.gen_dataset(
-          batch_size=1, is_training=False, preprocess=preprocess)
+          batch_size=1, is_training=False, **kwargs)
       converter.representative_dataset = tf.lite.RepresentativeDataset(
           _get_representative_dataset_gen(ds, self.quantization_steps))
 

tensorflow_examples/lite/model_maker/core/task/model_spec/object_detector_spec.py

@@ -24,10 +24,13 @@
 from tensorflow_examples.lite.model_maker.third_party.efficientdet import coco_metric
 from tensorflow_examples.lite.model_maker.third_party.efficientdet import hparams_config
 from tensorflow_examples.lite.model_maker.third_party.efficientdet import utils
+from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import efficientdet_keras
+from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import inference
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import label_util
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import postprocess
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import train
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import train_lib
+from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import util_keras
 
 
 def _get_ordered_label_map(label_map):
@@ -237,3 +240,92 @@ def _get_detections(images, labels):
         name = 'AP_/%s' % label_map[cid]
         metric_dict[name] = metrics[i + len(evaluator.metric_names)]
     return metric_dict
+
+  def export_saved_model(self,
+                         saved_model_dir,
+                         batch_size=None,
+                         pre_mode='infer',
+                         post_mode='global'):
+    """Saves the model to Tensorflow SavedModel.
+
+    Args:
+      saved_model_dir: Folder path for saved model.
+      batch_size: Batch size to be saved in saved_model.
+      pre_mode: Pre-processing Mode in ExportModel, must be {None, 'infer'}.
+      post_mode: Post-processing Mode in ExportModel, must be {None, 'global',
+        'per_class'}.
+    """
+    # Create EfficientDetModel with latest checkpoint.
+    config = self.config
+    model = efficientdet_keras.EfficientDetModel(config=config)
+    model.build((batch_size, *config.image_size, 3))
+    if config.model_dir:
+      util_keras.restore_ckpt(
+          model,
+          config.model_dir,
+          config['moving_average_decay'],
+          skip_mismatch=False)
+    else:
+      # EfficientDetModel is random initialized without restoring the
+      # checkpoint. This is mainly used in object_detector_test and shouldn't be
+      #  used if we want to export trained model.
+      tf.compat.v1.logging.warn('Need to restore the checkpoint for '
+                                'EfficientDet.')
+    # Gets tf.TensorSpec.
+    if pre_mode is None:
+      # Input is the preprocessed image that's already resized to a certain
+      # input shape.
+      input_spec = tf.TensorSpec(
+          shape=[batch_size, *config.image_size, 3],
+          dtype=tf.float32,
+          name='images')
+    else:
+      # Input is that raw image that can be in any input shape,
+      input_spec = tf.TensorSpec(
+          shape=[batch_size, None, None, 3], dtype=tf.uint8, name='images')
+
+    export_model = inference.ExportModel(
+        model, pre_mode=pre_mode, post_mode=post_mode)
+    tf.saved_model.save(
+        export_model,
+        saved_model_dir,
+        signatures=export_model.__call__.get_concrete_function(input_spec))
+
+  def export_tflite(self, tflite_filepath, quantization_config=None):
+    """Converts the retrained model to tflite format and saves it.
+
+    The exported TFLite model has the following inputs & outputs:
+    One input:
+      image: a float32 tensor of shape[1, height, width, 3] containing the
+        normalized input image. `self.config.image_size` is [height, width].
+
+    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.
+
+    Args:
+      tflite_filepath: File path to save tflite model.
+      quantization_config: Configuration for post-training quantization.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+      self.export_saved_model(
+          temp_dir, batch_size=1, pre_mode=None, post_mode='tflite')
+      converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir)
+      if quantization_config:
+        converter = quantization_config.get_converter_with_quantization(
+            converter, model_spec=self)
+
+      # TFLITE_BUILTINS is needed for TFLite's custom NMS op for integer only
+      # quantization.
+      if tf.lite.OpsSet.TFLITE_BUILTINS not in converter.target_spec.supported_ops:
+        converter.target_spec.supported_ops += [tf.lite.OpsSet.TFLITE_BUILTINS]
+      tflite_model = converter.convert()
+
+      with tf.io.gfile.GFile(tflite_filepath, 'wb') as f:
+        f.write(tflite_model)

tensorflow_examples/lite/model_maker/core/task/model_util.py

@@ -91,7 +91,7 @@ def export_tflite(model,
 
     if quantization_config:
       converter = quantization_config.get_converter_with_quantization(
-          converter, preprocess)
+          converter, preprocess=preprocess)
 
     converter.target_spec.supported_ops = supported_ops
     tflite_model = converter.convert()

tensorflow_examples/lite/model_maker/core/task/object_detector.py

@@ -117,3 +117,16 @@ def evaluate(self, data, batch_size=None):
 
     return self.model_spec.evaluate(self.model, ds, steps,
                                     data.annotations_json_file)
+
+  def _export_saved_model(self, saved_model_dir):
+    """Saves the model to Tensorflow SavedModel."""
+    self.model_spec.export_saved_model(saved_model_dir)
+
+  def _export_tflite(self, tflite_filepath, quantization_config=None):
+    """Converts the retrained model to tflite format and saves it.
+
+    Args:
+      tflite_filepath: File path to save tflite model.
+      quantization_config: Configuration for post-training quantization.
+    """
+    self.model_spec.export_tflite(tflite_filepath, quantization_config)

tensorflow_examples/lite/model_maker/core/task/object_detector_test.py

@@ -22,6 +22,8 @@
 from tensorflow_examples.lite.model_maker.core import compat
 from tensorflow_examples.lite.model_maker.core import test_util
 from tensorflow_examples.lite.model_maker.core.data_util import object_detector_dataloader
+from tensorflow_examples.lite.model_maker.core.export_format import ExportFormat
+from tensorflow_examples.lite.model_maker.core.task import configs
 from tensorflow_examples.lite.model_maker.core.task import object_detector
 from tensorflow_examples.lite.model_maker.core.task.model_spec import object_detector_spec
 
@@ -48,6 +50,38 @@ def testEfficientDetLite0(self):
     self.assertIsInstance(metrics, dict)
     self.assertGreaterEqual(metrics['AP'], 0)
 
+    # Export the model to saved model.
+    spec.config.model_dir = None  # Don't restore checkpoint.
+    output_path = os.path.join(self.get_temp_dir(), 'saved_model')
+    task.export(self.get_temp_dir(), export_format=ExportFormat.SAVED_MODEL)
+    self.assertTrue(os.path.isdir(output_path))
+    self.assertNotEqual(len(os.listdir(output_path)), 0)
+
+    # Export the model to TFLite model.
+    output_path = os.path.join(self.get_temp_dir(), 'float.tflite')
+    task.export(
+        self.get_temp_dir(),
+        tflite_filename='float.tflite',
+        export_format=ExportFormat.TFLITE)
+    self.assertTrue(tf.io.gfile.exists(output_path))
+    self.assertGreater(os.path.getsize(output_path), 0)
+
+    # Export the model to quantized TFLite model.
+    # TODO(b/175173304): Skips the test for stable tensorflow 2.4 for now since
+    # it fails. Will revert this change after TF upgrade.
+    if tf.__version__.startswith('2.4'):
+      return
+    output_path = os.path.join(self.get_temp_dir(), 'model_quantized.tflite')
+    config = configs.QuantizationConfig.create_full_integer_quantization(
+        data, is_integer_only=True)
+    task.export(
+        self.get_temp_dir(),
+        tflite_filename='model_quantized.tflite',
+        quantization_config=config,
+        export_format=ExportFormat.TFLITE)
+    self.assertTrue(os.path.isfile(output_path))
+    self.assertGreater(os.path.getsize(output_path), 0)
+
 
 if __name__ == '__main__':
   # Load compressed models from tensorflow_hub

tensorflow_examples/lite/model_maker/pip_package/setup.py

@@ -91,14 +91,10 @@ def _read_required_packages(fpath):
 
 def get_required_packages():
   """Gets packages inside requirements.txt."""
-  # Gets third party's required packages.
-  fpath = BASE_DIR.joinpath('third_party', 'efficientdet', 'requirements.txt')
-  required_pkgs = _read_required_packages(fpath)
-
   # Gets model maker's required packages
   filename = 'requirements_nightly.txt' if nightly else 'requirements.txt'
   fpath = BASE_DIR.joinpath(filename)
-  required_pkgs += _read_required_packages(fpath)
+  required_pkgs = _read_required_packages(fpath)
 
   return required_pkgs
 

tensorflow_examples/lite/model_maker/requirements.txt

@@ -14,3 +14,11 @@ tensorflow>=2.4.0
 librosa>=0.5
 lxml>=4.6.1
 PyYAML>=5.1
+# The following are the requirements of efficientdet.
+matplotlib>=3.0.3
+six>=1.12.0
+tensorflow-addons>=0.11.2
+neural-structured-learning>=1.3.1
+tensorflow-model-optimization>=0.5
+Cython>=0.29.13
+pycocotools>=2.0.2

tensorflow_examples/lite/model_maker/requirements_nightly.txt

@@ -15,3 +15,11 @@ tensorflow>=2.4.0
 librosa>=0.5
 lxml>=4.6.1
 PyYAML>=5.1
+# The following are the requirements of efficientdet.
+matplotlib>=3.0.3
+six>=1.12.0
+tfa-nightly
+neural-structured-learning>=1.3.1
+tensorflow-model-optimization>=0.5
+Cython>=0.29.13
+pycocotools>=2.0.2