Add inference_input_type and inference_output_type for quantization in model maker.

ziyeqinghan · copybara-github · commit ce15879e289a · 2020-04-29T04:03:18.000-07:00
PiperOrigin-RevId: 308995760
diff --git a/tensorflow_examples/lite/model_maker/core/task/custom_model.py b/tensorflow_examples/lite/model_maker/core/task/custom_model.py
@@ -143,7 +143,9 @@ def _export_tflite(self,
                      tflite_filepath,
                      quantized=False,
                      quantization_steps=None,
-                     representative_data=None):
+                     representative_data=None,
+                     inference_input_type=tf.float32,
+                     inference_output_type=tf.float32):
     """Converts the retrained model to tflite format and saves it.
 
     Args:
@@ -153,6 +155,12 @@ def _export_tflite(self,
         to run. Used only if `quantized` is True.
       representative_data: Representative data used for post-training
         quantization. Used only if `quantized` is True.
+      inference_input_type: Target data type of real-number input arrays. Allows
+        for a different type for input arrays. Defaults to tf.float32. Must be
+        be `{tf.float32, tf.uint8, tf.int8}`
+      inference_output_type: Target data type of real-number output arrays.
+        Allows for a different type for output arrays. Defaults to tf.float32.
+         Must be `{tf.float32, tf.uint8, tf.int8}`
     """
     if tflite_filepath is None:
       raise ValueError(
@@ -181,8 +189,8 @@ def _export_tflite(self,
           get_representative_dataset_gen(ds, quantization_steps))
 
       converter.optimizations = [tf.lite.Optimize.DEFAULT]
-      converter.inference_input_type = tf.uint8
-      converter.inference_output_type = tf.uint8
+      converter.inference_input_type = inference_input_type
+      converter.inference_output_type = inference_output_type
       converter.target_spec.supported_ops = [
           tf.lite.OpsSet.TFLITE_BUILTINS_INT8
       ]
diff --git a/tensorflow_examples/lite/model_maker/core/task/image_classifier.py b/tensorflow_examples/lite/model_maker/core/task/image_classifier.py
@@ -268,6 +268,8 @@ def _export_tflite(self,
                      quantized=False,
                      quantization_steps=None,
                      representative_data=None,
+                     inference_input_type=tf.float32,
+                     inference_output_type=tf.float32,
                      with_metadata=False,
                      export_metadata_json_file=False):
     """Converts the retrained model to tflite format and saves it.
@@ -281,14 +283,21 @@ def _export_tflite(self,
         to run. Used only if `quantized` is True.
       representative_data: Representative data used for post-training
         quantization. Used only if `quantized` is True.
+      inference_input_type: Target data type of real-number input arrays. Allows
+        for a different type for input arrays. Defaults to tf.float32. Must be
+        be `{tf.float32, tf.uint8, tf.int8}`
+      inference_output_type: Target data type of real-number output arrays.
+        Allows for a different type for output arrays. Defaults to tf.float32.
+         Must be `{tf.float32, tf.uint8, tf.int8}`
       with_metadata: Whether the output tflite model contains metadata.
       export_metadata_json_file: Whether to export metadata in json file. If
         True, export the metadata in the same directory as tflite model.Used
         only if `with_metadata` is True.
     """
     super(ImageClassifier,
           self)._export_tflite(tflite_filepath, quantized, quantization_steps,
-                               representative_data)
+                               representative_data, inference_input_type,
+                               inference_output_type)
     if with_metadata:
       if not metadata.TFLITE_SUPPORT_TOOLS_INSTALLED:
         tf.compat.v1.logging.warning('Needs to install tflite-support package.')
diff --git a/tensorflow_examples/lite/model_maker/core/task/text_classifier.py b/tensorflow_examples/lite/model_maker/core/task/text_classifier.py
@@ -187,7 +187,9 @@ def _export_tflite(self,
                      tflite_filepath,
                      quantized=False,
                      quantization_steps=None,
-                     representative_data=None):
+                     representative_data=None,
+                     inference_input_type=tf.float32,
+                     inference_output_type=tf.float32):
     """Converts the retrained model to tflite format and saves it.
 
     Args:
@@ -197,11 +199,18 @@ def _export_tflite(self,
         to run. Used only if `quantized` is True.
       representative_data: Representative data used for post-training
         quantization. Used only if `quantized` is True.
+      inference_input_type: Target data type of real-number input arrays. Allows
+        for a different type for input arrays. Defaults to tf.float32. Must be
+        be `{tf.float32, tf.uint8, tf.int8}`
+      inference_output_type: Target data type of real-number output arrays.
+        Allows for a different type for output arrays. Defaults to tf.float32.
+         Must be `{tf.float32, tf.uint8, tf.int8}`
     """
     # Sets batch size from None to 1 when converting to tflite.
     self._set_batch_size(self.model, batch_size=1)
     super(TextClassifier,
           self)._export_tflite(tflite_filepath, quantized, quantization_steps,
-                               representative_data)
+                               representative_data, inference_input_type,
+                               inference_output_type)
     # Sets batch size back to None to support retraining later.
     self._set_batch_size(self.model, batch_size=None)
