# Copyright 2019 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 for post training quantization with calibration.""" import numpy as np from tensorflow.lite.python.convert_phase import Component from tensorflow.lite.python.convert_phase import convert_phase from tensorflow.lite.python.convert_phase import SubComponent from tensorflow.lite.python.interpreter import Interpreter from tensorflow.python.framework import dtypes from tensorflow.python.util.lazy_loader import LazyLoader # Lazy load since some of the performance benchmark skylark rules # break dependencies. Must use double quotes to match code internal rewrite # rule. _calibration_wrapper = LazyLoader( "_calibration_wrapper", globals(), "tensorflow.lite.python.optimize." "_pywrap_tensorflow_lite_calibration_wrapper") def add_intermediate_tensors(model_content): """Adds intermediate tensors to fused op if needed.""" return _calibration_wrapper.AddIntermediateTensors(model_content) class Calibrator(object): """Calibrates a floating point model and then quantizes it. This is an internal class, not a public interface. """ def __init__(self, model_content, custom_op_registerers_by_name=None, custom_op_registerers_by_func=None): """Constructor. Args: model_content: Content of a TF-Lite Flatbuffer file. custom_op_registerers_by_name: List of str (symbol names) that take a pointer to a MutableOpResolver and register custom ops. custom_op_registerers_by_func: List of functions that take a pointer to a MutableOpResolver and register custom ops. Raises: ValueError: If the calibrator was unable to open the model. """ if not model_content: raise ValueError("`model_content` must be specified.") if custom_op_registerers_by_name is None: custom_op_registerers_by_name = [] if custom_op_registerers_by_func is None: custom_op_registerers_by_func = [] try: self._calibrator = ( _calibration_wrapper.CalibrationWrapper( model_content, custom_op_registerers_by_name, custom_op_registerers_by_func)) self._model_content = model_content except Exception as e: raise ValueError("Failed to parse the model: %s." % e) if not self._calibrator: raise ValueError("Failed to parse the model.") self._interpreter = None def _create_input_array_from_dict(self, signature_key, inputs): input_array = [] signature_runner = self._interpreter.get_signature_runner(signature_key) input_details = sorted( signature_runner.get_input_details().items(), key=lambda item: item[1]["index"]) for input_name, _ in input_details: input_array.append(inputs[input_name]) return input_array def _feed_tensors(self, dataset_gen, resize_input): """Feed tensors to the calibrator.""" initialized = {} for sample in dataset_gen(): if isinstance(sample, tuple): if not isinstance(sample[1], dict): raise ValueError("You need to provide either a dictionary with input " "names and values in the second arugment in the " "tuple") # Convert signature based inputs to the tensor index based data. if self._interpreter is None: self._interpreter = Interpreter(model_content=self._model_content) signature_key = sample[0] input_array = self._create_input_array_from_dict( signature_key, sample[1]) elif isinstance(sample, dict): # Convert signature based inputs to the tensor index based data. if self._interpreter is None: self._interpreter = Interpreter(model_content=self._model_content) signature_key = None input_array = self._create_input_array_from_dict(None, sample) elif isinstance(sample, list): signature_key = None input_array = sample else: raise ValueError("You need to provide either a dictionary with input " "names and values, a tuple with signature key and a " "dictionary with input names and values, or an array " "with input values in the order of input tensors of " "the graph in the representative_dataset function. " "Unsupported value from dataset: {}.".format(sample)) if signature_key not in initialized: initialized[signature_key] = True if resize_input: if signature_key is not None: self._calibrator.Prepare([list(s.shape) for s in input_array], signature_key) else: self._calibrator.Prepare([list(s.shape) for s in input_array]) else: if signature_key is not None: self._calibrator.Prepare(signature_key) else: self._calibrator.Prepare() if signature_key is not None: self._calibrator.FeedTensor(input_array, signature_key) else: self._calibrator.FeedTensor(input_array) @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER) def calibrate_and_quantize(self, dataset_gen, input_type, output_type, allow_float, activations_type=dtypes.int8, bias_type=dtypes.int32, resize_input=True, disable_per_channel=False): """Calibrates the model with specified generator and then quantizes it. The input shapes of the calibrator are resized with the calibration data if `resize_input` is set. Returns: A quantized model. Args: dataset_gen: A generator that generates calibration samples. input_type: A tf.dtype representing the desired real-value input type. output_type: A tf.dtype representing the desired real-value output type. allow_float: A boolean. False if the resulting model cannot perform float computation, useful when targeting an integer-only backend. If False, an error will be thrown if an operation cannot be quantized, otherwise the model will fallback to float ops. activations_type: A tf.dtype representing the desired type for activations. bias_type: A tf.dtype representing the desired type for bias. resize_input: A boolean. True if the shape of the sample data is different from the input. disable_per_channel: A boolean. True if disabling per-channel quantization. """ self._feed_tensors(dataset_gen, resize_input) return self._calibrator.QuantizeModel( np.dtype(input_type.as_numpy_dtype()).num, np.dtype(output_type.as_numpy_dtype()).num, allow_float, np.dtype(activations_type.as_numpy_dtype()).num, np.dtype(bias_type.as_numpy_dtype()).num, disable_per_channel) @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER) def calibrate_and_quantize_single(self, dataset_gen, input_type, output_type, allow_float, op_output_name, resize_input=True): """Calibrates the model with specified generator and then quantizes it. Only the single op with output op_output_name will be quantized. The input shapes of the calibrator are resized with the calibration data. Returns: A quantized model. Args: dataset_gen: A generator that generates calibration samples. input_type: A tf.dtype representing the desired real-value input type. output_type: A tf.dtype representing the desired real-value output type. allow_float: A boolean. False if the resulting model cannot perform float computation, useful when targeting an integer-only backend. If False, an error will be thrown if an operation cannot be quantized, otherwise the model will fallback to float ops. op_output_name: A string, only this op will be quantized. resize_input: A boolean. True if the shape of the sample data is different from the input. """ self._feed_tensors(dataset_gen, resize_input) return self._calibrator.QuantizeModel( np.dtype(input_type.as_numpy_dtype()).num, np.dtype(output_type.as_numpy_dtype()).num, allow_float, op_output_name) @convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.CALIBRATE) def calibrate(self, dataset_gen): """Calibrates the model with specified generator. Returns: A model with min and max calibration stats. Args: dataset_gen: A generator that generates calibration samples. """ self._feed_tensors(dataset_gen, resize_input=True) return self._calibrator.Calibrate()