Add an option to convert a full INT8 TFLite model with FP32 input and output.

PiperOrigin-RevId: 468808165

official/vision/serving/export_tflite.py

@@ -69,8 +69,9 @@ flags.DEFINE_string(
 flags.DEFINE_string(
     'quant_type',
     default=None,
-    help='Post training quantization type. Support `int8`, `int8_full`, '
-    '`fp16`, and `default`. See '
+    help='Post training quantization type. Support `int8_fp32_fallback`, '
+    '`int8_fp32_input_output`, `int8_full`, `fp16`, `qat`, '
+    '`qat_fp32_input_output`, and `default`. See '
     'https://www.tensorflow.org/lite/performance/post_training_quantization '
     'for more details.')
 flags.DEFINE_integer('calibration_steps', 500,

official/vision/serving/export_tflite_lib.py

@@ -107,16 +107,17 @@ def convert_tflite_model(saved_model_dir: str,
           representative_dataset,
           params=params,
           calibration_steps=calibration_steps)
-      if quant_type == 'int8_full':
+      if quant_type in ('int8_full', 'int8_fp32_input_output'):
         converter.target_spec.supported_ops = [
             tf.lite.OpsSet.TFLITE_BUILTINS_INT8
         ]
+      if quant_type == 'int8_full':
         converter.inference_input_type = tf.uint8  # or tf.int8
         converter.inference_output_type = tf.uint8  # or tf.int8
     elif quant_type == 'fp16':
       converter.optimizations = [tf.lite.Optimize.DEFAULT]
       converter.target_spec.supported_types = [tf.float16]
-    elif quant_type == 'default':
+    elif quant_type in ('default', 'qat_fp32_input_output'):
       converter.optimizations = [tf.lite.Optimize.DEFAULT]
     elif quant_type == 'qat':
       converter.optimizations = [tf.lite.Optimize.DEFAULT]

official/vision/serving/export_tflite_lib_test.py

@@ -76,7 +76,14 @@ class ExportTfliteLibTest(tf.test.TestCase, parameterized.TestCase):
   @combinations.generate(
       combinations.combine(
           experiment=['mobilenet_imagenet'],
-          quant_type=[None, 'default', 'fp16', 'int8', 'int8_full']))
+          quant_type=[
+              None,
+              'default',
+              'fp16',
+              'int8_fp32_fallback',
+              'int8_full',
+              'int8_fp32_input_output',
+          ]))
   def test_export_tflite_image_classification(self, experiment, quant_type):
 
     params = exp_factory.get_exp_config(experiment)
@@ -105,7 +112,14 @@ class ExportTfliteLibTest(tf.test.TestCase, parameterized.TestCase):
   @combinations.generate(
       combinations.combine(
           experiment=['retinanet_mobile_coco'],
-          quant_type=[None, 'default', 'fp16', 'int8', 'int8_full']))
+          quant_type=[
+              None,
+              'default',
+              'fp16',
+              'int8_fp32_fallback',
+              'int8_full',
+              'int8_fp32_input_output',
+          ]))
   def test_export_tflite_detection(self, experiment, quant_type):
 
     params = exp_factory.get_exp_config(experiment)
@@ -138,7 +152,14 @@ class ExportTfliteLibTest(tf.test.TestCase, parameterized.TestCase):
   @combinations.generate(
       combinations.combine(
           experiment=['mnv2_deeplabv3_pascal'],
-          quant_type=[None, 'default', 'fp16', 'int8', 'int8_full']))
+          quant_type=[
+              None,
+              'default',
+              'fp16',
+              'int8_fp32_fallback',
+              'int8_full',
+              'int8_fp32_input_output',
+          ]))
   def test_export_tflite_semantic_segmentation(self, experiment, quant_type):
 
     params = exp_factory.get_exp_config(experiment)