U h@sddlZddlZddlZddlZddlmZddlmZddlm Z m Z m Z m Z m Z ddlZddlZddlmZmZmZmZddlZddlmZmZmZGdd d ZGd d d ZGd d d eZGdddejdZGdddZGdddeZ GdddeZ!Gddde!Z"Gddde!Z#Gddde!Z$GdddejdZ%Gdd d e%Z&dd!ej'd"ife e(efe e e(d#d$d%Z)dS)&N)Enum)Path)DictOptionalSequenceTupleUnion) ModelProto TensorProtohelper numpy_helper) apply_plotload_model_with_shape_infersmooth_distributionc@sBeZdZedddddddgZdd Zed d Zed d ZdS) TensorDataavgstdlowesthighesthist hist_edgesbinscKsD|D]6\}}|tjkr2td|dtjdt|||qdS)NzUnexpected value z not in .)itemsr_allowed ValueErrorsetattr)selfkwargskvr"F/tmp/pip-unpacked-wheel-socb9apf/onnxruntime/quantization/calibrate.py__init__s zTensorData.__init__cCs4t|drt|ds(tdt|d|j|jfS)Nrrz0Attributes 'lowest' and/or 'highest' missing in r)hasattrAttributeErrordirrrrr"r"r# range_value!szTensorData.range_valuecCs4t|drt|ds(tdt|d|j|jfS)Nrrz)Attributes 'avg' and/or 'std' missing in r)r%r&r'rrr(r"r"r#avg_std'szTensorData.avg_stdN) __name__ __module__ __qualname__ frozensetrr$propertyr)r*r"r"r"r#rs  rc@sReZdZeeeeeffdddZddZ ddZ dd Z d d Z d d Z dS) TensorsDatadatac Cs||_i|_|D]\}}t|ts:tdt|dt|tr|tj krvt |dkrvt |d|dd|j|<qt |dkrt |d|d|d|dd |j|<qtd |d d t |d |dt|t stdt|d||j|<qdS)NzKeys must be strings not rrr )rr)rrrrzUnexpected tuple for rz , it has z elements: zValues must be TensorData not ) calibration_methodr2r isinstancestr TypeErrortypetupleCalibrationMethodMinMaxlenr)rr7r2r r!r"r"r#r$/s    &" zTensorsData.__init__ccs|jEdHdSNr1r(r"r"r#__iter__AszTensorsData.__iter__cCs ||jkSr@r1rkeyr"r"r# __contains__DszTensorsData.__contains__cCs |j|Sr@r1rBr"r"r# __getitem__GszTensorsData.__getitem__cCs(||jkrtd|d||j|<dS)Nz)Only an existing tensor can be modified, z is not.)r2 RuntimeError)rrCvaluer"r"r# __setitem__Js zTensorsData.__setitem__cCs |jSr@)r2valuesr(r"r"r#rIOszTensorsData.valuesN)r+r,r-rr9rrrr$rArDrErHrIr"r"r"r#r0.s r0c@seZdZdZdZdZdZdS)r=rr r3r5N)r+r,r-r>Entropy Percentile Distributionr"r"r"r#r=Ssr=c@s<eZdZeddZejedddZddZ dd Z d S) CalibrationDataReadercCst|drt|jptS)Nget_next)r%callablerNNotImplemented)clssubclassr"r"r#__subclasshook__[sz&CalibrationDataReader.__subclasshook__returncCstdS)z9generate the input data dict for ONNXinferenceSession runNNotImplementedErrorr(r"r"r#rN_szCalibrationDataReader.get_nextcCs|Sr@r"r(r"r"r#rAdszCalibrationDataReader.__iter__cCs|}|dkrt|Sr@)rN StopIteration)rresultr"r"r#__next__gszCalibrationDataReader.__next__N) r+r,r- classmethodrSabcabstractmethoddictrNrArZr"r"r"r#rMZs  rM) metaclassc@s~eZdZdeeefeeedddZdgfdd Z d d Z e d d dZ ddZ ddZedddZedddZdS)CalibraterBaseNaugmented_model.onnxF model_pathop_types_to_calibratecCsht|trtt||_nt|tr0t||_ntd||_||_||_||_ d|_ d|_ dg|_ dS)a :param model_path: ONNX model to calibrate. It should be a model file path :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param symmetric: make range of tensor symmetric (central point is 0). :param use_external_data_format: use external data format to store model which size is >= 2Gb z model_path should be model path.NCPUExecutionProvider) r8r9rrmodelrrdaugmented_model_path symmetricuse_external_data_formatZ augment_model infer_sessionexecution_providers)rrcrdrgrhrir"r"r#r$os   zCalibraterBase.__init__recCs||_|dS)zz reset the execution providers to execute the collect_data. It triggers to re-creating inference session. N)rkcreate_inference_session)rrkr"r"r#set_execution_providerssz&CalibraterBase.set_execution_providerscCs,t}tjj|_tj|j||jd|_dS)z9 create an OnnxRuntime InferenceSession. ) sess_optionsZ providersN) onnxruntimeZSessionOptionsZGraphOptimizationLevelZORT_DISABLE_ALLZgraph_optimization_levelZInferenceSessionrgrkrj)rrnr"r"r#rls z'CalibraterBase.create_inference_sessionrfc Csdd|jjD}|dd|jjD|dd|jjDdd|jjD}t}tjh}|jj D]h}|j r|j |j krjt |j|jD]@}||kr||}|jdr|jjj|kr||kr||qqj||fS)z select input/output tensors of candidate nodes to calibrate. returns: tensors (set): set of tensor name. value_infos (dict): tensor name to value info. cSsi|] }|j|qSr"name).0vir"r"r# sz>CalibraterBase.select_tensors_to_calibrate..cSsi|] }|j|qSr"rq)rsZotr"r"r#ruscSsi|] }|j|qSr"rq)rsitr"r"r#ruscSsh|] }|jqSr"rq)rsinitr"r"r# sz=CalibraterBase.select_tensors_to_calibrate.. tensor_type)graphZ value_infoupdateoutputinput initializersetr FLOATnoderdZop_type itertoolschainr;ZHasFieldryZ elem_typeadd) rrf value_infosr~tensors_to_calibrateZtensor_type_to_calibrater tensor_namertr"r"r#select_tensors_to_calibrates&   z*CalibraterBase.select_tensors_to_calibratecCs|jS)zP return: augmented onnx model. Call after calling augment_graph rpr(r"r"r#get_augment_modelsz CalibraterBase.get_augment_modelcCstdS)z abstract method: augment the input model to prepare for collecting data. It will: 1. augment the model to be able to collect desired statistics data 2. save augmented model to augmented_model_paths NrVr(r"r"r# augment_graphszCalibraterBase.augment_graph data_readercCstdS)z abstract method: collect the tensors that will be used for range computation. It can be called multiple times. NrV)rrr"r"r# collect_dataszCalibraterBase.collect_datarTcCstdS)ze abstract method: compute data based on the calibration method stored in TensorsData NrVr(r"r"r# compute_dataszCalibraterBase.compute_data)NraFF)r+r,r-rr9rrrr$rmrlr rrrrMrr0rr"r"r"r#r`ns   r`csjeZdZdeeefeeedfdd Zdd Z d d Z e d d dZ ddZ edddZZS)MinMaxCalibraterNraF{Gz?rbcsttj|||||dg|_d|_t|jjj|_dd|jjjD|_ ||_ |rj|dksb|dkrjt d||_ dS)a :param model_path: ONNX model to calibrate. It is a model path :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param symmetric: make range of tensor symmetric (central point is 0). :param use_external_data_format: use external data format to store model which size is >= 2Gb :param moving_average: compute the moving average of the minimum and maximum values instead of the global minimum and maximum. :param averaging_constant: constant smoothing factor to use when computing the moving average. rdrgrhriNcSsh|] }|jqSr"rqrsr|r"r"r#rxsz,MinMaxCalibrater.__init__..rr z;Invalid averaging constant, which should not be < 0 or > 1.) superr$intermediate_outputscalibrate_tensors_ranger?rfrzr|num_model_outputsmodel_original_outputsmoving_averageraveraging_constant)rrcrdrgrhrirr __class__r"r#r$szMinMaxCalibrater.__init__csj\}}ttttjdgtj d}jj j |fdd}|D]}||d||dqXt jjjjddS) z Adds ReduceMin and ReduceMax nodes to all quantization_candidates op type nodes in model and ensures their outputs are stored as part of the graph output :return: augmented ONNX model r Zdtypecsd}|d|}|d}tjj||g|g||d}tjjd|g|g|d}jjj||gjjjt |t j dgdS)Nr _Z_Reshape)keepdimsrrZReshape)inputsoutputsrr) onnxr Z make_noderfrzrextendr|appendZmake_tensor_value_infor r)rZreduce_op_namerZ reduce_outputintermediate_outputZ reduce_nodeZ reshape_nodeZreshape_shape_namerr"r#add_reduce_min_maxs$ z:MinMaxCalibrater.augment_graph..add_reduce_min_maxZ ReduceMinZ ReduceMaxZsave_as_external_dataN)rrfr9uuidZuuid4r Z from_arraynparrayint64rzr~rrsavergri)rZtensorsrZ reshape_shapertensorr"rr#rs   zMinMaxCalibrater.augment_graphcCs g|_dSr@rr(r"r"r#clear_collected_data(sz%MinMaxCalibrater.clear_collected_datarcCsn|}|sq&|j|jd|qt|jdkrRsz1MinMaxCalibrater.compute_data..csg|]}tt|qSr"r^ziprsr output_namesr"r#rSsNcsg|]}|ddqS)rr) rpartitionr)added_output_namesr"r#r\sr3cs i|]}|jkr||qSr")rr)merged_output_dictrr"r#ru`s z1MinMaxCalibrater.compute_data..)Zaxisr )r?rrranger setdefaultrrrrZmeanrrr8intsizefloatrhabsr<r0r=r>r^rr)routput_dicts_listdr r!Zcalibrate_tensor_namesZmerged_added_output_dictpairsrrrZmin_value_arrayZmax_value_arrayZmax_absolute_valueZnew_calibrate_tensors_ranger")rrrrr#rIsN    zMinMaxCalibrater.compute_data)NraFFFr)r+r,r-rr9rrrr$rrrMrrr0r __classcell__r"r"rr#rs  #+rc sbeZdZdeeefeeed fd d Zd d Z ddZ e dddZ e dddZZS)HistogramCalibraterNraF percentile-X@samerbc svtj|||||dg|_d|_t|jjj|_dd|jjjD|_ d|_ ||_ ||_ ||_ | |_d|_| |_dS)a= :param model_path: ONNX model to calibrate. It is a model path. :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param use_external_data_format: use external data format to store model which size is >= 2Gb :param method: A string. One of ['entropy', 'percentile']. :param symmetric: make range of tensor symmetric (central point is 0). :param num_bins: number of bins to create a new histogram for collecting tensor values. :param num_quantized_bins: number of quantized bins. Default 128. :param percentile: A float number between [0, 100]. Default 99.99. :param scenario: see :class:`DistributionCalibrater` rNcSsh|] }|jqSr"rqrr"r"r#rxsz/HistogramCalibrater.__init__..)rr$rrr?rfrzr|rr collectormethodnum_binsnum_quantized_binsrrscenario) rrcrdrgrirrhrrrrrr"r#r$s$zHistogramCalibrater.__init__cCsV||j\|_}|jD]"}||jkr|jjj||qtj|j|j |j ddS)z make all quantization_candidates op type nodes as part of the graph output. :return: augmented ONNX model rN) rrfrrrzr|rrrrgri)rrrr"r"r#rs  z!HistogramCalibrater.augment_graphcCs g|_dSr@rr(r"r"r#rsz(HistogramCalibrater.clear_collected_datarcs|}|sq&jjd|qtjdkr.csg|]}tt|qSr"rrrr"r#rscs i|]}|jkr||qSr")rr) merged_dictrr"r#rus z4HistogramCalibrater.collect_data..)rrhrrrr)rNrrrjrr?rrrrrHistogramCollectorrrhrrrrcollectr)rrrrrr r!Zclean_merged_dictr")rrrr#rs4   z HistogramCalibrater.collect_datarTcCsh|jstdt|tr tj}n8t|tr2tj}n&t|trDtj }nt dt |dt ||j S)z Compute the min-max range of tensor :return: dictionary mapping: {tensor name: (min value, max value)} z9No collector created and can't generate calibration data.zUnknown calibrater z". This method must be overwritten.)rrr8EntropyCalibraterr=rJPercentileCalibraterrKDistributionCalibraterrLr:r;r0compute_collection_result)rcalr"r"r#rs   z HistogramCalibrater.compute_data) NraFrFrrrr)r+r,r-rr9rrrr$rrrMrr0rrr"r"rr#rs"  ,&rcs6eZdZd eeefeeedfdd ZZ S) rNraFentropyrrbc s tj||||||||ddS)a :param model_path: ONNX model to calibrate. It is a model path :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param use_external_data_format: use external data format to store model which size is >= 2Gb :param method: A string. One of ['entropy', 'percentile', 'distribution']. :param symmetric: make range of tensor symmetric (central point is 0). :param num_bins: number of bins to create a new histogram for collecting tensor values. :param num_quantized_bins: number of quantized bins. Default 128. )rrhrrNrr$) rrcrdrgrirrhrrrr"r#r$szEntropyCalibrater.__init__)NraFrFrr r+r,r-rr9rrrr$rr"r"rr#rs  rcs6eZdZd eeefeeedfdd ZZ S) rNraFrrrrbc s tj||||||||ddS)a :param model_path: ONNX model to calibrate. It is a model path :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param use_external_data_format: use external data format to store model which size is >= 2Gb :param method: A string. One of ['entropy', 'percentile', 'distribution']. :param symmetric: make range of tensor symmetric (central point is 0). :param num_quantized_bins: number of quantized bins. Default 128. :param percentile: A float number between [0, 100]. Default 99.99. )rrhrrNr) rrcrdrgrirrhrrrr"r#r$szPercentileCalibrater.__init__)NraFrFrrrr"r"rr#rs  rcs6eZdZd eeefeeedfdd ZZ S) rNraF distributionrrrbc stj|||||||ddS)a :param model_path: ONNX model to calibrate. It is a model path :param op_types_to_calibrate: operator types to calibrate. By default, calibrate all the float32/float16 tensors. :param augmented_model_path: save augmented model to this path. :param use_external_data_format: use external data format to store model which size is >= 2Gb :param method: A string. One of ['entropy', 'percentile', 'distribution']. :param symmetric: make range of tensor symmetric (central point is 0). :param num_bins: number of bins to create a new histogram for collecting tensor values. :param scenario: for float 8 only, if `scenario="same"`, the algorithm weights and float 8 follow the same distribution, if `scenario="p3"`, it assumes the weights follow a gaussian law and float 8 ~ X^3 where X is a gaussian law )rrrNr)rrcrdrgrirrrrr"r#r$@szDistributionCalibrater.__init__)NraFrrrrr"r"rr#r?s  rc@s,eZdZdZejddZejddZdS)CalibrationDataCollectorzL Base class for collecting data for calibration-based quantization. cCstdS)z Generate informative data based on given data. name_to_arr : dict tensor name to NDArray data NrVr name_to_arrr"r"r#rgsz CalibrationDataCollector.collectcCstdS)z? Get the optimal result among collection data. NrVr(r"r"r#rpsz2CalibrationDataCollector.compute_collection_resultN)r+r,r-__doc__r\r]rrr"r"r"r#rbs  rc@sveZdZdZddZddZddZdd Zd d Zd d Z ddZ ddZ ddZ e dddZddZddZdS)ra` Collecting histogram for each tensor. Percentile and Entropy method are supported. ref: https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py ref: https://docs.nvidia.com/deeplearning/tensorrt/pytorch-quantization-toolkit/docs/_modules/ pytorch_quantization/calib/histogram.html cCs.i|_||_||_||_||_||_||_dSr@)histogram_dictrrhrrrr)rrrhrrrrr"r"r#r$szHistogramCollector.__init__cCs|jSr@)rr(r"r"r#get_histogram_dictsz%HistogramCollector.get_histogram_dictcCsNtd|jdkr||S|jdkrB|jr6||S||SntddS)Nz/Collecting tensor data and making histogram ...>rrrDOnly 'entropy', 'percentile' or 'distribution' methods are supported)printr collect_valuerhcollect_absolute_valuerrr"r"r#rs     zHistogramCollector.collectcCsP|D]@\}}t|}|}|jdkrDt|}t|}nd}d}t|}||jkrtj ||j d\}}||||f|j|<q|j|}|d} |d} |d} |d} t|} | | dkr| d| d}t | d|| ||}t | |f} tj || d\}}|dt | | 7<||t| |t| |f|j|<qdS)z5 Collect histogram on absolute value r)rr3r5r N)rrasarrayflattenrrrabsoluter histogramrZarangeZhstackr?)rrrdata_arrrrrr old_histogramold_minold_maxold_histold_hist_edgesZ temp_amaxwidthZ new_bin_edgesr"r"r#rs2        z)HistogramCollector.collect_absolute_valuec Cs|D]\}}t|}|}|jdkrBt|}t|}nd}d}tt|t|}||jkr|j|}| ||||||j|<qtj ||j | |fd\}} || |||f|j|<qdS)z1 Collect histogram on real value rrN) rrrrrrrrrmerge_histogramrr) rrrrrr thresholdrrrr"r"r#rs2       z HistogramCollector.collect_valuecCs|\}}}} } || krRtj|t|| | fd\} } | ||t||t| || fS| dkrtj|t|| |fd\} }| |7} nrt|}d| |}t|| |d}|d|}||| }tj||| |fd\} }| ||||7<| |t||t| ||fSdS)Nrrr3r )rrr?rrr)rrrZnew_minZnew_maxZ new_thresholdrrrrZ old_thresholdZnew_histrrrZ old_num_binsZ old_strideZhalf_increased_binsZ new_num_binsr"r"r#rs2    z"HistogramCollector.merge_histogramcCsp|jrt|jdkrtdtd|jd|jdkr@|S|jdkrR|S|jdkrd|StddS) Nrz=Histogram has not been collected. Please run collect() first.z0Finding optimal threshold for each tensor using z algorithm ...rrrr)rr?rrrcompute_entropycompute_percentilecompute_distributionr(r"r"r#rs   z,HistogramCollector.compute_collection_resultcCs|jdks|jdkrtd|j}|j}i}tdt|td|jtdd|d|d |D]8\}}|d}|d }|}t ||} |j rt | |d} t ||  t || f||<nDd|d } t | d | } t | | } t || t || f||<|d } |d}||d| krP| ||d f||<||d |krv||d|f||<|||dd ||<t jdddkrpt||qp|S)Nrdz.rrinfr5)copyZ scipy.statsrrrzerosrrdeepcopyr Zastyperrr8ZndarrayZargmin)rrrrrrrrZzero_bin_indexZnum_half_quantized_binZ kl_divergenceZ thresholdsr start_indexZ end_indexZsliced_distributionpZleft_outliers_countZright_outliers_countZnonzerosZquantized_binsZnum_merged_binsindexstartendqZnormZmin_kl_divergence_idxrrrr"r"r#rsd             z(HistogramCollector.get_entropy_thresholdN)r )r+r,r-rr$rrrrrrrr staticmethodrrrr"r"r"r#rxs %  . rraF)rfrdc Csd}|tjkrdd|krdn|d}d|kr.dn|d}d|krBdn|d} t||||||| d}n |tjkrd|krzdn|d} d |krdn|d } d|krdn|d}t|||||| | d }n|tjkr$d|krd n|d} d |krd n|d } d|krdn|d}t|||||| | d}nL|tjkrpd|kr>d n|d} d|krTdn|d} t||||| | d}|r| | |St d|dS)NrhFrrr)rirhrrrrr)rirhrrrrrT)rirhrrrr)rirrzUnsupported calibration method ) r=r>rrJrrKrrLrrrlr)rfrdrgZcalibrate_methodriZ extra_optionsZ calibratorrhrrrrrrr"r"r#create_calibratorsp     r")*r\rr renumrpathlibrtypingrrrrrZnumpyrrr r r r roZ quant_utilsrrrrr0r=ABCMetarMr`rrrrrrrr>r9r"r"r"r"r#sF  %k*y""#t