o i e%@s|dZddlZddlmZddlmZddlmZddlmZddl m Z ddl m Z e d e d Zd d ZGd ddZdS)z?Python wrapper for post training quantization with calibration.N) Component) convert_phase) SubComponent) Interpreter)dtypes) LazyLoader_calibration_wrapperzKtensorflow.lite.python.optimize._pywrap_tensorflow_lite_calibration_wrappercCs t|S)z0Adds intermediate tensors to fused op if needed.)rAddIntermediateTensors model_contentr Y/var/www/myenv/lib/python3.10/site-packages/tensorflow/lite/python/optimize/calibrator.pyadd_intermediate_tensors&s rc@seZdZdZ  dddZddZddZeej e j e j e jd d fd d Zeej e j  dd dZeej e jddZdS) CalibratorzqCalibrates a floating point model and then quantizes it. This is an internal class, not a public interface. Nc Cs||std|dur g}|durg}z t||||_||_Wnty1}ztd|d}~ww|js9tdd|_dS)aConstructor. 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. z"`model_content` must be specified.NzFailed to parse the model: %s.zFailed to parse the model.) ValueErrorrCalibrationWrapper _calibrator_model_content Exception _interpreter)selfr custom_op_registerers_by_namecustom_op_registerers_by_funcer r r __init__1s(   zCalibrator.__init__cCsHg}|j|}t|ddd}|D] \}}|||q|S)NcSs |ddS)Nindexr )itemr r r [s z:Calibrator._create_input_array_from_dict..)key)rget_signature_runnersortedget_input_detailsitemsappend)r signature_keyinputs input_arraysignature_runner input_details input_name_r r r _create_input_array_from_dictVs   z(Calibrator._create_input_array_from_dictcCsNi}|D]}t|tr0t|dtstd|jdur#t|jd|_|d}|||d}n+t|trJ|jdurAt|jd|_d}|d|}nt|trTd}|}ntd |||vrd||<|r|durv|j dd |D|n|j d d |Dn|dur|j |n|j |dur|j ||q|j |qdS) zFeed tensors to the calibrator.rzgYou need to provide either a dictionary with input names and values in the second arugment in the tupleNr ra"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: {}.TcSg|]}t|jqSr listshape.0sr r r z,Calibrator._feed_tensors..cSr-r r.r1r r r r4r5) isinstancetupledictrrrrr,r/formatrPrepare FeedTensor)r dataset_gen resize_input initializedsampler%r'r r r _feed_tensorsasP        zCalibrator._feed_tensorsTFc CsR||||jt|jt|j|t|jt|j|S)aBCalibrates 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. r@r QuantizeModelnpdtypeas_numpy_dtypenum) rr< input_type output_type allow_floatactivations_type bias_typer=disable_per_channelr r r calibrate_and_quantizes 'z!Calibrator.calibrate_and_quantizecCs6||||jt|jt|j||S)aCalibrates 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. rA)rr<rGrHrIop_output_namer=r r r calibrate_and_quantize_singles !z(Calibrator.calibrate_and_quantize_singlecCs|j|dd|jS)zCalibrates the model with specified generator. Returns: A model with min and max calibration stats. Args: dataset_gen: A generator that generates calibration samples. T)r=)r@r Calibrate)rr<r r r calibrates zCalibrator.calibrate)NN)T)__name__ __module__ __qualname____doc__rr,r@rrOPTIMIZE_TFLITE_MODELr#QUANTIZE_USING_DEPRECATED_QUANTIZERrint8int32rMrO CALIBRATErQr r r r r+s0 % 9  - %r)rUnumpyrC$tensorflow.lite.python.convert_phaserrr"tensorflow.lite.python.interpreterrtensorflow.python.frameworkr"tensorflow.python.util.lazy_loaderrglobalsrrrr r r r s