From a07ca66517313dc0ed0f20c074f762e9d3129112 Mon Sep 17 00:00:00 2001 From: Meghna Natraj Date: Fri, 24 Apr 2020 17:52:11 -0700 Subject: [PATCH] Add python wrapper to modify model interface for full integer quantization PiperOrigin-RevId: 308359969 Change-Id: I5929e4778e5e12651b64665bfa9f3dbc6f82e72b --- tensorflow/lite/tools/optimize/python/BUILD | 67 +++++++++ .../optimize/python/modify_model_interface.cc | 40 ++++++ .../optimize/python/modify_model_interface.py | 78 +++++++++++ .../modify_model_interface_constants.py | 34 +++++ .../python/modify_model_interface_lib.py | 79 +++++++++++ .../python/modify_model_interface_lib_test.py | 129 ++++++++++++++++++ 6 files changed, 427 insertions(+) create mode 100644 tensorflow/lite/tools/optimize/python/BUILD create mode 100644 tensorflow/lite/tools/optimize/python/modify_model_interface.cc create mode 100644 tensorflow/lite/tools/optimize/python/modify_model_interface.py create mode 100644 tensorflow/lite/tools/optimize/python/modify_model_interface_constants.py create mode 100644 tensorflow/lite/tools/optimize/python/modify_model_interface_lib.py create mode 100644 tensorflow/lite/tools/optimize/python/modify_model_interface_lib_test.py diff --git a/tensorflow/lite/tools/optimize/python/BUILD b/tensorflow/lite/tools/optimize/python/BUILD new file mode 100644 index 00000000000..050c0008924 --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/BUILD @@ -0,0 +1,67 @@ +load("//tensorflow:tensorflow.bzl", "pybind_extension") + +package( + default_visibility = [ + "//visibility:public", + ], + licenses = ["notice"], # Apache 2.0 +) + +py_binary( + name = "modify_model_interface", + srcs = ["modify_model_interface.py"], + python_version = "PY3", + srcs_version = "PY3", + deps = [ + ":modify_model_interface_constants", + ":modify_model_interface_lib", + "//tensorflow/python:platform", + ], +) + +py_library( + name = "modify_model_interface_lib", + srcs = ["modify_model_interface_lib.py"], + srcs_version = "PY3", + deps = [ + ":_pywrap_modify_model_interface", + ":modify_model_interface_constants", + "//tensorflow:tensorflow_py", + "//tensorflow/lite/python:schema_py", + ], +) + +py_test( + name = "modify_model_interface_lib_test", + srcs = ["modify_model_interface_lib_test.py"], + python_version = "PY3", + srcs_version = "PY3", + tags = [ + "no_mac", # TODO(b/148247402): flatbuffers import broken on Mac OS. + ], + deps = [ + ":modify_model_interface_lib", + "//tensorflow:tensorflow_py", + "//tensorflow/python:client_testlib", + "//tensorflow/python:framework_test_lib", + "//third_party/py/numpy", + ], +) + +py_library( + name = "modify_model_interface_constants", + srcs = ["modify_model_interface_constants.py"], + srcs_version = "PY3", + deps = ["//tensorflow/lite/python:lite_constants"], +) + +pybind_extension( + name = "_pywrap_modify_model_interface", + srcs = ["modify_model_interface.cc"], + module_name = "_pywrap_modify_model_interface", + deps = [ + "//tensorflow/lite/schema:schema_fbs", + "//tensorflow/lite/tools/optimize:modify_model_interface", + "@pybind11", + ], +) diff --git a/tensorflow/lite/tools/optimize/python/modify_model_interface.cc b/tensorflow/lite/tools/optimize/python/modify_model_interface.cc new file mode 100644 index 00000000000..ed67b07cb0f --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/modify_model_interface.cc @@ -0,0 +1,40 @@ +/* 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. +==============================================================================*/ + +// Python wrapper to modify model interface. + +#include "tensorflow/lite/tools/optimize/modify_model_interface.h" + +#include + +#include "pybind11/pybind11.h" +#include "tensorflow/lite/schema/schema_generated.h" + +namespace pybind11 { + +PYBIND11_MODULE(_pywrap_modify_model_interface, m) { + // An anonymous function that invokes the C++ function + // after applying transformations to the python function arguments + m.def("modify_model_interface", + [](const std::string& input_file, const std::string& output_file, + const int input_type, const int output_type) -> int { + return tflite::optimize::ModifyModelInterface( + input_file, output_file, + static_cast(input_type), + static_cast(input_type)); + }); +} + +} // namespace pybind11 diff --git a/tensorflow/lite/tools/optimize/python/modify_model_interface.py b/tensorflow/lite/tools/optimize/python/modify_model_interface.py new file mode 100644 index 00000000000..938f353b0ae --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/modify_model_interface.py @@ -0,0 +1,78 @@ +# 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. +# ============================================================================== +r"""Modify a quantized model's interface from float to integer. + +Example usage: +python modify_model_interface_main.py \ + --input_file=float_model.tflite \ + --output_file=int_model.tflite \ + --input_type=INT8 \ + --output_type=INT8 +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import argparse +import sys + +from tensorflow.lite.tools.optimize.python import modify_model_interface_constants as mmi_constants +from tensorflow.lite.tools.optimize.python import modify_model_interface_lib as mmi_lib +from tensorflow.python.platform import app + + +def main(_): + """Application run loop.""" + parser = argparse.ArgumentParser( + description="Modify a quantized model's interface from float to integer.") + parser.add_argument( + '--input_file', + type=str, + required=True, + help='Full path name to the input tflite file.') + parser.add_argument( + '--output_file', + type=str, + required=True, + help='Full path name to the output tflite file.') + parser.add_argument( + '--input_type', + type=str.upper, + choices=mmi_constants.STR_TYPES, + default=mmi_constants.DEFAULT_STR_TYPE, + help='Modified input integer interface type.') + parser.add_argument( + '--output_type', + type=str.upper, + choices=mmi_constants.STR_TYPES, + default=mmi_constants.DEFAULT_STR_TYPE, + help='Modified output integer interface type.') + args = parser.parse_args() + + input_type = mmi_constants.STR_TO_TFLITE_TYPES[args.input_type] + output_type = mmi_constants.STR_TO_TFLITE_TYPES[args.output_type] + + mmi_lib.modify_model_interface(args.input_file, args.output_file, input_type, + output_type) + + print('Successfully modified the model input type from FLOAT to ' + '{input_type} and output type from FLOAT to {output_type}.'.format( + input_type=args.input_type, output_type=args.output_type)) + + +if __name__ == '__main__': + app.run(main=main, argv=sys.argv[:1]) diff --git a/tensorflow/lite/tools/optimize/python/modify_model_interface_constants.py b/tensorflow/lite/tools/optimize/python/modify_model_interface_constants.py new file mode 100644 index 00000000000..42767268e48 --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/modify_model_interface_constants.py @@ -0,0 +1,34 @@ +# 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. +# ============================================================================== +"""Constants for modify_model_interface.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.lite.python import lite_constants + +STR_TO_TFLITE_TYPES = { + 'INT8': lite_constants.INT8, + 'UINT8': lite_constants.QUANTIZED_UINT8 +} +TFLITE_TO_STR_TYPES = {v: k for k, v in STR_TO_TFLITE_TYPES.items()} + +STR_TYPES = STR_TO_TFLITE_TYPES.keys() +TFLITE_TYPES = STR_TO_TFLITE_TYPES.values() + +DEFAULT_STR_TYPE = 'INT8' +DEFAULT_TFLITE_TYPE = STR_TO_TFLITE_TYPES[DEFAULT_STR_TYPE] diff --git a/tensorflow/lite/tools/optimize/python/modify_model_interface_lib.py b/tensorflow/lite/tools/optimize/python/modify_model_interface_lib.py new file mode 100644 index 00000000000..5e4bf99ccdf --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/modify_model_interface_lib.py @@ -0,0 +1,79 @@ +# 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 modify a quantized model's interface from float to integer.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from tensorflow.lite.python import schema_py_generated as schema_fb +from tensorflow.lite.tools.optimize.python import _pywrap_modify_model_interface +from tensorflow.lite.tools.optimize.python import modify_model_interface_constants as mmi_constants + + +def _parse_type_to_int(dtype, flag): + """Converts a tflite type to it's integer representation. + + Args: + dtype: tf.DType representing the inference type. + flag: str representing the flag name. + + Returns: + integer, a tflite TensorType enum value. + + Raises: + ValueError: Unsupported tflite type. + """ + # Validate if dtype is supported in tflite and is a valid interface type. + if dtype not in mmi_constants.TFLITE_TYPES: + raise ValueError( + "Unsupported value '{0}' for {1}. Only {2} are supported.".format( + dtype, flag, mmi_constants.TFLITE_TYPES)) + + dtype_str = mmi_constants.TFLITE_TO_STR_TYPES[dtype] + dtype_int = schema_fb.TensorType.__dict__[dtype_str] + + return dtype_int + + +def modify_model_interface(input_file, output_file, input_type, output_type): + """Modify a quantized model's interface (input/output) from float to integer. + + Args: + input_file: Full path name to the input tflite file. + output_file: Full path name to the output tflite file. + input_type: Final input interface type. + output_type: Final output interface type. + + Raises: + RuntimeError: If the modification of the model interface was unsuccessful. + ValueError: If the input_type or output_type is unsupported. + + """ + # Map the interface types to integer values + input_type_int = _parse_type_to_int(input_type, 'input_type') + output_type_int = _parse_type_to_int(output_type, 'output_type') + + # Invoke the function to modify the model interface + status = _pywrap_modify_model_interface.modify_model_interface( + input_file, output_file, input_type_int, output_type_int) + + # Throw an exception if the return status is an error. + if status != 0: + raise RuntimeError( + 'Error occured when trying to modify the model input type from float ' + 'to {input_type} and output type from float to {output_type}.'.format( + input_type=input_type, output_type=output_type)) diff --git a/tensorflow/lite/tools/optimize/python/modify_model_interface_lib_test.py b/tensorflow/lite/tools/optimize/python/modify_model_interface_lib_test.py new file mode 100644 index 00000000000..e97f0db9bbb --- /dev/null +++ b/tensorflow/lite/tools/optimize/python/modify_model_interface_lib_test.py @@ -0,0 +1,129 @@ +# 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. +# ============================================================================== +"""Tests for modify_model_interface_lib.py.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import os +import numpy as np +import tensorflow as tf + +from tensorflow.lite.tools.optimize.python import modify_model_interface_lib +from tensorflow.python.framework import test_util +from tensorflow.python.platform import test + + +def build_tflite_model_with_full_integer_quantization(): + # Define TF model + input_size = 3 + model = tf.keras.Sequential([ + tf.keras.layers.InputLayer(input_shape=(input_size,), dtype=tf.float32), + tf.keras.layers.Dense(units=5, activation=tf.nn.relu), + tf.keras.layers.Dense(units=2, activation=tf.nn.softmax) + ]) + + # Convert TF Model to a Quantized TFLite Model + converter = tf.lite.TFLiteConverter.from_keras_model(model) + converter.optimizations = [tf.lite.Optimize.DEFAULT] + + def representative_dataset_gen(): + for i in range(10): + yield [np.array([i] * input_size, dtype=np.float32)] + + converter.representative_dataset = representative_dataset_gen + converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] + tflite_model = converter.convert() + + return tflite_model + + +class ModifyModelInterfaceTest(test_util.TensorFlowTestCase): + + def testInt8Interface(self): + # 1. SETUP + # Define the temporary directory and files + temp_dir = self.get_temp_dir() + initial_file = os.path.join(temp_dir, 'initial_model.tflite') + final_file = os.path.join(temp_dir, 'final_model.tflite') + # Define initial model + initial_model = build_tflite_model_with_full_integer_quantization() + with open(initial_file, 'wb') as model_file: + model_file.write(initial_model) + + # 2. INVOKE + # Invoke the modify_model_interface function + modify_model_interface_lib.modify_model_interface(initial_file, final_file, + tf.int8, tf.int8) + + # 3. VALIDATE + # Load TFLite model and allocate tensors. + initial_interpreter = tf.lite.Interpreter(model_path=initial_file) + initial_interpreter.allocate_tensors() + final_interpreter = tf.lite.Interpreter(model_path=final_file) + final_interpreter.allocate_tensors() + + # Get input and output types. + initial_input_dtype = initial_interpreter.get_input_details()[0]['dtype'] + initial_output_dtype = initial_interpreter.get_output_details()[0]['dtype'] + final_input_dtype = final_interpreter.get_input_details()[0]['dtype'] + final_output_dtype = final_interpreter.get_output_details()[0]['dtype'] + + # Validate the model interfaces + self.assertEqual(initial_input_dtype, np.float32) + self.assertEqual(initial_output_dtype, np.float32) + self.assertEqual(final_input_dtype, np.int8) + self.assertEqual(final_output_dtype, np.int8) + + def testUInt8Interface(self): + # 1. SETUP + # Define the temporary directory and files + temp_dir = self.get_temp_dir() + initial_file = os.path.join(temp_dir, 'initial_model.tflite') + final_file = os.path.join(temp_dir, 'final_model.tflite') + # Define initial model + initial_model = build_tflite_model_with_full_integer_quantization() + with open(initial_file, 'wb') as model_file: + model_file.write(initial_model) + + # 2. INVOKE + # Invoke the modify_model_interface function + modify_model_interface_lib.modify_model_interface(initial_file, final_file, + tf.uint8, tf.uint8) + + # 3. VALIDATE + # Load TFLite model and allocate tensors. + initial_interpreter = tf.lite.Interpreter(model_path=initial_file) + initial_interpreter.allocate_tensors() + final_interpreter = tf.lite.Interpreter(model_path=final_file) + final_interpreter.allocate_tensors() + + # Get input and output types. + initial_input_dtype = initial_interpreter.get_input_details()[0]['dtype'] + initial_output_dtype = initial_interpreter.get_output_details()[0]['dtype'] + final_input_dtype = final_interpreter.get_input_details()[0]['dtype'] + final_output_dtype = final_interpreter.get_output_details()[0]['dtype'] + + # Validate the model interfaces + self.assertEqual(initial_input_dtype, np.float32) + self.assertEqual(initial_output_dtype, np.float32) + self.assertEqual(final_input_dtype, np.uint8) + self.assertEqual(final_output_dtype, np.uint8) + + +if __name__ == '__main__': + test.main()