From 1b611c66bb1577d4ecbbf4b4e97347261cbf9b78 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Tue, 11 Oct 2022 14:38:14 -0700 Subject: [PATCH] Improve quantization support in model_maker/image_classifier PiperOrigin-RevId: 480455944 --- mediapipe/model_maker/python/core/utils/model_util.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/mediapipe/model_maker/python/core/utils/model_util.py b/mediapipe/model_maker/python/core/utils/model_util.py index 8962f286..0899a9b1 100644 --- a/mediapipe/model_maker/python/core/utils/model_util.py +++ b/mediapipe/model_maker/python/core/utils/model_util.py @@ -94,7 +94,8 @@ def export_tflite( tflite_filepath: str, quantization_config: Optional[quantization.QuantizationConfig] = None, supported_ops: Tuple[tf.lite.OpsSet, - ...] = (tf.lite.OpsSet.TFLITE_BUILTINS,)): + ...] = (tf.lite.OpsSet.TFLITE_BUILTINS,), + preprocess: Optional[Callable[..., bool]] = None): """Converts the model to tflite format and saves it. Args: @@ -102,6 +103,9 @@ def export_tflite( tflite_filepath: File path to save tflite model. quantization_config: Configuration for post-training quantization. supported_ops: A list of supported ops in the converted TFLite file. + preprocess: A callable to preprocess the representative dataset for + quantization. The callable takes three arguments in order: feature, + label, and is_training. """ if tflite_filepath is None: raise ValueError( @@ -113,7 +117,8 @@ def export_tflite( converter = tf.lite.TFLiteConverter.from_saved_model(save_path) if quantization_config: - converter = quantization_config.set_converter_with_quantization(converter) + converter = quantization_config.set_converter_with_quantization( + converter, preprocess=preprocess) converter.target_spec.supported_ops = supported_ops tflite_model = converter.convert()