§
    ¹¯ÅjnÕ  ã                   ó°  — d dl Z d dlZd dlZd dlZd dlZd dlmZ d dlmZ d dl	m
Z
 d dlZd dlZd dlmZmZmZmZ d dlZddlmZmZmZ dej        d	ej        d
ej        fd„Z	 	 d-dej        d	ej        dedz  ded
ej        f
d„Z G d„ d¦  «        Z G d„ d¦  «        Z G d„ de¦  «        Z G d„ de j        ¬¦  «        Z  G d„ d¦  «        Z! G d„ de!¦  «        Z" G d„ de!¦  «        Z# G d„ de#¦  «        Z$ G d „ d!e#¦  «        Z% G d"„ d#e#¦  «        Z& G d$„ d%e j        ¬¦  «        Z' G d&„ d'e'¦  «        Z(dd(ej)        d)di fd*e*e
z  d+ee*         dz  fd,„Z+dS ).é    N)ÚSequence)ÚEnum)ÚPath)Ú
ModelProtoÚTensorProtoÚhelperÚnumpy_helperé   )Ú
apply_plotÚload_model_with_shape_inferÚsmooth_distributionÚpkÚqkÚreturnc                 ó  — t          j        | j        | j        ¬¦  «        }| dd…         t          j        | dd…         |dd…         z  ¦  «        z  |dd…<   | dk    |dk    z  }d||<   | dk    |dk    z  }t           j        || <   |S )z‘
    See https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.rel_entr.html#scipy.special.rel_entr.
    Python implementation.
    ©ÚdtypeNr   )ÚnpÚemptyÚshaper   ÚlogÚinf)r   r   ÚresÚc2Úc1s        úSC:\Projects\Historical\venv\Lib\site-packages\onnxruntime/quantization/calibrate.pyÚrel_entrr      s“   € õ
 Œ(�2”8 2¤8Ð
,Ñ
,Ô
,€CØ���ŒU•R”V˜B˜q˜q˜qœE B q q q¤E™MÑ*Ô*Ñ*€Cˆˆˆ�FØ
�Š'�b˜A’gÑ	€BØ€Cˆ�GØ
ˆqŠ&�R˜!’VÑ	€BÝŒv€Cˆˆ�HØ€Jó    ÚbaseÚaxisc                 óX  — |�|dk    s
J d¦   «         ‚|€
J d¦   «         ‚t          j        | ¦  «                             t           j        ¦  «        } d| z  t          j        | |d¬¦  «        z  } t          j        |¦  «                             t           j        ¦  «        }t          j        | |¦  «        \  } }d|z  t          j        ||d¬¦  «        z  }t          | |¦  «        }t          j        ||¬¦  «        }|�|t          j        |¦  «        z  }|                     | j        ¦  «        S )	zÉ
    Simplifeied version of entropy.
    Source: https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.entropy.html.
    This avoids taking a dependency on scipy just for this function.
    Nr   z0base={base} must be a positive number or `None`.z
qk is Noneç      ð?T)r    Úkeepdims©r    )	r   ÚasarrayÚastypeÚfloat32ÚsumÚbroadcast_arraysr   r   r   )r   r   r   r    ÚvecÚss         r   Úentropyr,   '   s	  € ð ˆ<˜4 !š8˜8˜8Ð%W™8œ8Ð#Øˆ>ˆ>˜<‰>Œ>ˆ>å	Œ�B‰Œ×	Ò	�rœzÑ	*Ô	*€BØ	ˆr‰•B”F˜2 D°4Ð8Ñ8Ô8Ñ	8€Bå	Œ�B‰Œ×	Ò	�rœzÑ	*Ô	*€BÝÔ   RÑ(Ô(�F€BˆØ	ˆr‰•B”F˜2 D°4Ð8Ñ8Ô8Ñ	8€BÝ
�2�rÑ
Ô
€Cå
Œˆs˜ÐÑÔ€AØÐØ	�RŒV�D‰\Œ\ÑˆØ�8Š8�B”HÑÔÐr   c                   óz   — e Zd Z eg d¢¦  «        Z eg d¢¦  «        Zd„ Zed„ ¦   «         Zed„ ¦   «         Z	d„ Z
dS )Ú
TensorData)ÚavgÚstdÚlowestÚhighestÚhistÚ
hist_edgesÚbins)r/   r0   r1   r2   r4   c                 óö  — t          |                     ¦   «         ¦  «        | _        |                     ¦   «         D ]½\  }}|t          j        vr t          d|›dt          j        › d�¦  «        ‚|t          j        v rkt          |d¦  «        s"t          dt          |¦  «        › d|›�¦  «        ‚|j
        t          j        t          j        fvrt          d|j
        › d|›�¦  «        ‚t          | ||¦  «         Œ¾d S )NzUnexpected value z not in ú.r   úUnexpected type z for k=zUnexpected dtype )ÚlistÚkeysÚ_attrsÚitemsr.   Ú_allowedÚ
ValueErrorÚ_floatsÚhasattrÚtyper   r   Úfloat16r'   Úsetattr)ÚselfÚkwargsÚkÚvs       r   Ú__init__zTensorData.__init__G   s  € Ý˜6Ÿ;š;™=œ=Ñ)Ô)ˆŒØ—L’L‘N”Nð 	 ð 	 ‰DˆAˆqØ�
Ô+Ð+Ð+Ý Ð!X°QÐ!XÐ!XÅ*ÔBUÐ!XÐ!XÐ!XÑYÔYÐYØ•JÔ&Ð&Ð&Ý˜q 'Ñ*Ô*ð OÝ$Ð%M½¸Q¹¼Ð%MÐ%MÈÐ%MÐ%MÑNÔNÐNØ”7¥2¤:­r¬zÐ":Ð:Ð:Ý$Ð%N¸¼Ð%NÐ%NÈÐ%NÐ%NÑOÔOÐOÝ�D˜!˜QÑÔÐÐð	 ð 	 r   c                 óž   — t          | d¦  «        rt          | d¦  «        s t          dt          | ¦  «        › d�¦  «        ‚| j        | j        fS )Nr1   r2   z0Attributes 'lowest' and/or 'highest' missing in r7   )r@   ÚAttributeErrorÚdirr1   r2   ©rD   s    r   Úrange_valuezTensorData.range_valueS   sZ   € å�t˜XÑ&Ô&ð 	b­g°d¸IÑ.FÔ.Fð 	bÝ Ð!`ÕTWÐX\ÑT]ÔT]Ð!`Ð!`Ð!`ÑaÔaÐaØ”˜Tœ\Ð*Ð*r   c                 óž   — t          | d¦  «        rt          | d¦  «        s t          dt          | ¦  «        › d�¦  «        ‚| j        | j        fS )Nr/   r0   z)Attributes 'avg' and/or 'std' missing in r7   )r@   rJ   rK   r/   r0   rL   s    r   Úavg_stdzTensorData.avg_stdY   sW   € å�t˜UÑ#Ô#ð 	[­7°4¸Ñ+?Ô+?ð 	[Ý Ð!YÍSÐQUÉYÌYÐ!YÐ!YÐ!YÑZÔZÐZØ”˜$œ(Ð#Ð#r   c                 óL   ‡ — ˆ fd„‰ j         D ¦   «         }‰ j        j        |d<   |S )Nc                 ó2   •— i | ]}|t          ‰|¦  «        “ŒS © )Úgetattr)Ú.0rF   rD   s     €r   ú
<dictcomp>z&TensorData.to_dict.<locals>.<dictcomp>a   s%   ø€ Ð9Ð9Ð9¨�•7˜4 Ñ#Ô#Ð9Ð9Ð9r   ÚCLS)r;   Ú	__class__Ú__name__©rD   Údatas   ` r   Úto_dictzTensorData.to_dict_   s1   ø€ à9Ð9Ð9Ð9¨T¬[Ð9Ñ9Ô9ˆØ”nÔ-ˆˆU‰Øˆr   N)rX   Ú
__module__Ú__qualname__Ú	frozensetr=   r?   rH   ÚpropertyrM   rO   r[   rR   r   r   r.   r.   C   s”   € € € € € ØˆyÐZÐZÐZÑ[Ô[€HØˆiÐIÐIÐIÑJÔJ€Gð
 ð 
 ð 
 ð ð+ð +ñ „Xð+ð
 ð$ð $ñ „Xð$ð
ð ð ð ð r   r.   c                   ó`   — e Zd Zdeeeez  f         fd„Zd„ Zd„ Z	d„ Z
d„ Zd„ Zd„ Zd	„ Zd
„ ZdS )ÚTensorsDatarZ   c           
      óÚ  — || _         i | _        |                     ¦   «         D �]F\  }}t          |t          ¦  «        s t          dt          |¦  «        › d�¦  «        ‚t          |t          ¦  «        r·|t          j	        k    r9t          |¦  «        dk    r&t          |d         |d         ¬¦  «        | j        |<   Œ™t          |¦  «        dk    r4t          |d         |d         |d         |d         ¬	¦  «        | j        |<   Œàt          d
|d›dt          |¦  «        › d|› d�¦  «        ‚t          |t          ¦  «        s t          dt          |¦  «        › d�¦  «        ‚|| j        |<   �ŒHd S )NzKeys must be strings not r7   é   r   r
   ©r1   r2   é   é   )r1   r2   r3   r5   zUnexpected tuple for Úrz	, it has z elements: zValues must be TensorData not )Úcalibration_methodrZ   r<   Ú
isinstanceÚstrÚ	TypeErrorrA   ÚtupleÚCalibrationMethodÚMinMaxÚlenr.   )rD   rh   rZ   rF   rG   s        r   rH   zTensorsData.__init__g   ss  € Ø"4ˆÔØˆŒ	Ø—J’J‘L”Lð 	ñ 	‰DˆAˆqÝ˜a¥Ñ%Ô%ð HÝÐ F½DÀ¹G¼GÐ FÐ FÐ FÑGÔGÐGÝ˜!�UÑ#Ô#ð _Ø%Õ):Ô)AÒAÐAÅcÈ!ÁfÄfÐPQÂkÀkÝ#-°Q°q´TÀ1ÀQÄ4Ð#HÑ#HÔ#H�D”I˜a‘LØÝ�q‘6”6˜Q’;�;Ý#-°Q°q´TÀ1ÀQÄ4ÈaÐPQÌdÐYZÐ[\ÔY]Ð#^Ñ#^Ô#^�D”I˜a‘LØÝÐ ]¸Ð ]Ð ]Ð ]ÅcÈ!ÁfÄfÐ ]Ð ]ÐYZÐ ]Ð ]Ð ]Ñ^Ô^Ð^Ý˜a¥Ñ,Ô,ð MÝÐ KÅÀaÁÄÐ KÐ KÐ KÑLÔLÐLØˆDŒI�a‰L‰Lð	ð 	r   c              #   ó$   K  — | j         E d {V —† d S ©N©rZ   rL   s    r   Ú__iter__zTensorsData.__iter__y   s&   è è € Ø”9ÐÐÐÐÐÐÐÐÐr   c                 ó   — || j         v S rq   rr   ©rD   Úkeys     r   Ú__contains__zTensorsData.__contains__|   s   € Ø�d”iÐÐr   c                 ó   — | j         |         S rq   rr   ru   s     r   Ú__getitem__zTensorsData.__getitem__   s   € ØŒy˜Œ~Ðr   c                 óR   — || j         vrt          d|›d�¦  «        ‚|| j         |<   d S )Nz)Only an existing tensor can be modified, z is not.)rZ   ÚRuntimeError)rD   rv   Úvalues      r   Ú__setitem__zTensorsData.__setitem__‚   s8   € Ø�d”iÐÐÝÐZÈ3ÐZÐZÐZÑ[Ô[Ð[ØˆŒ	�#‰ˆˆr   c                 ó4   — | j                              ¦   «         S rq   )rZ   r:   rL   s    r   r:   zTensorsData.keys‡   s   € ØŒy�~Š~ÑÔÐr   c                 ó4   — | j                              ¦   «         S rq   )rZ   ÚvaluesrL   s    r   r€   zTensorsData.valuesŠ   s   € ØŒy×ÒÑ!Ô!Ð!r   c                 ó4   — | j                              ¦   «         S rq   )rZ   r<   rL   s    r   r<   zTensorsData.items�   s   € ØŒy�ŠÑ Ô Ð r   c                 ó:   — | j         j        | j        | j        dœ}|S )N)rV   rZ   rh   )rW   rX   rZ   rh   rY   s     r   r[   zTensorsData.to_dict�   s*   € ð ”>Ô*Ø”IØ"&Ô"9ð
ð 
ˆð
 ˆr   N)rX   r\   r]   Údictrj   r.   rl   rH   rs   rw   ry   r}   r:   r€   r<   r[   rR   r   r   ra   ra   f   s¶   € € € € € ð°°c¸:ÈÑ;MÐ6MÔ1Nð ð ð ð ð$ð ð ð ð  ð  ðð ð ðð ð ð
 ð  ð  ð"ð "ð "ð!ð !ð !ðð ð ð ð r   ra   c                   ó   — e Zd ZdZdZdZdZdS )rm   r   r
   rc   rf   N)rX   r\   r]   rn   ÚEntropyÚ
PercentileÚDistributionrR   r   r   rm   rm   š   s"   € € € € € Ø€FØ€GØ€JØ€L€L€Lr   rm   c                   ól   — e Zd Zed„ ¦   «         Zej        defd„¦   «         Zd„ Z	d„ Z
d„ Zdedefd	„Zd
S )ÚCalibrationDataReaderc                 óX   — t          |d¦  «        rt          |j        ¦  «        pt          S )NÚget_next)r@   Úcallabler‹   ÚNotImplemented)ÚclsÚsubclasss     r   Ú__subclasshook__z&CalibrationDataReader.__subclasshook__¢   s(   € å˜ *Ñ-Ô-ÐMµ(¸8Ô;LÑ2MÔ2MÐ`ÕR`Ð`r   r   c                 ó   — t           ‚)z9generate the input data dict for ONNXinferenceSession run©ÚNotImplementedErrorrL   s    r   r‹   zCalibrationDataReader.get_next¦   s
   € õ "Ð!r   c                 ó   — | S rq   rR   rL   s    r   rs   zCalibrationDataReader.__iter__«   s   € Øˆr   c                 ó@   — |                       ¦   «         }|€t          ‚|S rq   )r‹   ÚStopIteration)rD   Úresults     r   Ú__next__zCalibrationDataReader.__next__®   s   € Ø—’‘”ˆØˆ>ÝÐØˆr   c                 ó   — t           ‚rq   r’   rL   s    r   Ú__len__zCalibrationDataReader.__len__´   ó   € Ý!Ð!r   Ústart_indexÚ	end_indexc                 ó   — t           ‚rq   r’   )rD   rœ   r�   s      r   Ú	set_rangezCalibrationDataReader.set_range·   r›   r   N)rX   r\   r]   Úclassmethodr�   ÚabcÚabstractmethodrƒ   r‹   rs   r˜   rš   ÚintrŸ   rR   r   r   r‰   r‰   ¡   sª   € € € € € Øðað añ „[ðað 	Ôð"˜$ð "ð "ð "ñ Ôð"ðð ð ðð ð ð"ð "ð "ð" Sð "°Sð "ð "ð "ð "ð "ð "r   r‰   )Ú	metaclassc                   ó„   — e Zd Z	 	 	 	 	 ddeez  dee         dz  fd„Zdgfd„Zd	„ Zd
e	fd„Z
d„ Zd„ Zdefd„Zdefd„ZdS )ÚCalibraterBaseNúaugmented_model.onnxFÚ
model_pathÚop_types_to_calibratec                 óX  — t          |t          ¦  «        r"t          t          |¦  «        ¦  «        | _        n9t          |t          ¦  «        rt          |¦  «        | _        nt          d¦  «        ‚|| _        || _        || _        || _	        || _
        d| _        d| _        dg| _        dS )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.
        :param per_channel: whether to compute ranges per each channel.
        z model_path should be model path.NÚCPUExecutionProvider)ri   rj   r   r   Úmodelr>   r©   Úaugmented_model_pathÚ	symmetricÚuse_external_data_formatÚper_channelÚaugment_modelÚinfer_sessionÚexecution_providers)rD   r¨   r©   r­   r®   r¯   r°   s          r   rH   zCalibraterBase.__init__¼   s«   € õ" �j¥#Ñ&Ô&ð 	AÝ4µT¸*Ñ5EÔ5EÑFÔFˆDŒJˆJÝ˜
¥DÑ)Ô)ð 	AÝ4°ZÑ@Ô@ˆDŒJˆJåÐ?Ñ@Ô@Ð@à%:ˆÔ"Ø$8ˆÔ!Ø"ˆŒØ(@ˆÔ%Ø&ˆÔà!ˆÔØ!ˆÔØ$:Ð#;ˆÔ Ð Ð r   r«   c                 ó<   — || _         |                      ¦   «          dS )zz
        reset the execution providers to execute the collect_data. It triggers to re-creating inference session.
        N)r³   Úcreate_inference_session)rD   r³   s     r   Úset_execution_providersz&CalibraterBase.set_execution_providersÞ   s$   € ð $7ˆÔ Ø×%Ò%Ñ'Ô'Ð'Ð'Ð'r   c                 ó¤   — t          j        ¦   «         }t           j        j        |_        t          j        | j        || j        ¬¦  «        | _        dS )z9
        create an OnnxRuntime InferenceSession.
        )Úsess_optionsÚ	providersN)	ÚonnxruntimeÚSessionOptionsÚGraphOptimizationLevelÚORT_DISABLE_ALLÚgraph_optimization_levelÚInferenceSessionr­   r³   r²   )rD   r¸   s     r   rµ   z'CalibraterBase.create_inference_sessionå   sO   € õ #Ô1Ñ3Ô3ˆÝ0;Ô0RÔ0bˆÔ-Ý(Ô9ØÔ%Ø%ØÔ.ð
ñ 
ô 
ˆÔÐÐr   r¬   c                 ó‚  — d„ |j         j        D ¦   «         }|                     d„ |j         j        D ¦   «         ¦  «         |                     d„ |j         j        D ¦   «         ¦  «         d„ |j         j        D ¦   «         }t          ¦   «         }t          j        t          j	        h}|j         j
        D ]‹}| j        r|j        | j        v rtt          j        |j        |j        ¦  «        D ]T}||v rN||         }|j                             d¦  «        r,|j        j        j        |v r||vr|                     |¦  «         ŒUŒŒ||fS )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.
        c                 ó   — i | ]
}|j         |“ŒS rR   ©Úname©rT   Úvis     r   rU   z>CalibraterBase.select_tensors_to_calibrate.<locals>.<dictcomp>ø   s   € ÐDÐDÐD r�r”w ÐDÐDÐDr   c                 ó   — i | ]
}|j         |“ŒS rR   rÂ   )rT   Úots     r   rU   z>CalibraterBase.select_tensors_to_calibrate.<locals>.<dictcomp>ù   s   € ÐEÐEÐE¨B˜BœG RÐEÐEÐEr   c                 ó   — i | ]
}|j         |“ŒS rR   rÂ   )rT   Úits     r   rU   z>CalibraterBase.select_tensors_to_calibrate.<locals>.<dictcomp>ú   s   € ÐDÐDÐD¨B˜BœG RÐDÐDÐDr   c                 ó   — h | ]	}|j         ’Œ
S rR   rÂ   )rT   Úinits     r   ú	<setcomp>z=CalibraterBase.select_tensors_to_calibrate.<locals>.<setcomp>û   s   € ÐEÐEÐE T�t”yÐEÐEÐEr   Útensor_type)ÚgraphÚ
value_infoÚupdateÚoutputÚinputÚinitializerÚsetr   ÚFLOATÚFLOAT16Únoder©   Úop_typeÚ	itertoolsÚchainrA   ÚHasFieldrÍ   Ú	elem_typeÚadd)	rD   r¬   Úvalue_infosrÓ   Útensors_to_calibrateÚtensor_type_to_calibrater×   Útensor_namerÅ   s	            r   Úselect_tensors_to_calibratez*CalibraterBase.select_tensors_to_calibrateñ   s`  € ð EÐD¨U¬[Ô-CÐDÑDÔDˆØ×ÒÐEÐE°%´+Ô2DÐEÑEÔEÑFÔFÐFØ×ÒÐDÐD°%´+Ô2CÐDÑDÔDÑEÔEÐEØEÐE¨U¬[Ô-DÐEÑEÔEˆå"™uœuÐÝ$/Ô$5µ{Ô7JÐ#KÐ à”KÔ$ð 
	Bð 
	BˆDØÔ-ð 	B°´ÀÔA[Ð1[Ð1[Ý#,¤?°4´:¸t¼{Ñ#KÔ#Kð Bð B�KØ" kÐ1Ð1Ø(¨Ô5˜àœG×,Ò,¨]Ñ;Ô;ðBà!#¤Ô!4Ô!>ÐBZÐ!ZÐ!ZØ!,°KÐ!?Ð!?à0×4Ò4°[ÑAÔAÐAøøà# [Ð0Ð0r   c                 ó   — | j         S )zP
        return: augmented onnx model. Call after calling augment_graph
        )r¬   rL   s    r   Úget_augment_modelz CalibraterBase.get_augment_model  s   € ð ŒzÐr   c                 ó   — t           ‚)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
        r’   rL   s    r   Úaugment_graphzCalibraterBase.augment_graph  s
   € õ "Ð!r   Údata_readerc                 ó   — t           ‚)z€
        abstract method: collect the tensors that will be used for range computation. It can be called multiple times.
        r’   )rD   rç   s     r   Úcollect_datazCalibraterBase.collect_data  ó
   € õ "Ð!r   r   c                 ó   — t           ‚)ze
        abstract method: compute data based on the calibration method stored in TensorsData
        r’   rL   s    r   Úcompute_datazCalibraterBase.compute_data"  rê   r   )Nr§   FFF)rX   r\   r]   rj   r   r   rH   r¶   rµ   r   râ   rä   ræ   r‰   ré   ra   rì   rR   r   r   r¦   r¦   »   sú   € € € € € ð 7;Ø3ØØ!&Øð <ð  <à˜$‘Jð <ð  (¨œ}¨tÑ3ð <ð  <ð  <ð  <ðD <RÐ:Rð (ð (ð (ð (ð

ð 

ð 

ð1°ð 1ð 1ð 1ð 1ð:ð ð ð"ð "ð "ð"Ð(=ð "ð "ð "ð "ð"˜kð "ð "ð "ð "ð "ð "r   r¦   c                   ó|   ‡ — e Zd Z	 	 	 	 	 	 	 	 ddeez  dee         dz  fˆ fd„Zd„ Zd	„ Zd
e	fd„Z
d„ Zdefd„Zˆ xZS )ÚMinMaxCalibraterNr§   Fç{®Gáz„?r¨   r©   c
                 ó\  •— t          ¦   «                              ||||||	¬¦  «         g | _        d| _        t	          | j        j        j        ¦  «        | _        d„ | j        j        j        D ¦   «         | _	        || _
        |r|dk     s|dk    rt          d¦  «        ‚|| _        || _        dS )aw  
        :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.
        :param max_intermediate_outputs: maximum number of intermediate outputs before an intermediate range is computed.
        :param per_channel: whether to compute ranges per each channel.
        )r©   r­   r®   r¯   r°   Nc                 ó   — h | ]	}|j         ’Œ
S rR   rÂ   ©rT   rÑ   s     r   rÌ   z,MinMaxCalibrater.__init__.<locals>.<setcomp>L  ó   € Ð&YÐ&YÐ&Y°v v¤{Ð&YÐ&YÐ&Yr   r   r
   z;Invalid averaging constant, which should not be < 0 or > 1.)ÚsuperrH   Úintermediate_outputsÚcalibrate_tensors_rangero   r¬   rÎ   rÑ   Únum_model_outputsÚmodel_original_outputsÚmoving_averager>   Úaveraging_constantÚmax_intermediate_outputs)rD   r¨   r©   r­   r®   r¯   rù   rú   rû   r°   rW   s             €r   rH   zMinMaxCalibrater.__init__*  sÌ   ø€ õ. 	‰Œ×ÒØØ"7Ø!5ØØ%=Ø#ð 	ñ 	
ô 	
ð 	
ð %'ˆÔ!Ø'+ˆÔ$Ý!$ T¤ZÔ%5Ô%<Ñ!=Ô!=ˆÔØ&YÐ&YÀÄÔAQÔAXÐ&YÑ&YÔ&YˆÔ#Ø,ˆÔØð 	\Ð1°AÒ5Ð5Ð9KÈaÒ9OÐ9OÝÐZÑ[Ô[Ð[Ø"4ˆÔØ(@ˆÔ%Ð%Ð%r   c                 óÞ  ‡ ‡‡‡— ‰                       ‰ j        ¦  «        \  }}t          t          j        ¦   «         ¦  «        Št          j        t          j        dgt          j	        ¬¦  «        ‰¦  «        }‰ j        j
        j                             |¦  «         d„ Šˆ fd„Šˆˆˆˆ fd„}|D ]} ||d¦  «          ||d¦  «         Œt          j        ‰ j        ‰ j        ‰ j        ¬¦  «         d	S )
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   c                 ó˜   — |j         D ]0}t          j                             | |j        ¦  «        r	|j        c S Œ1t          d| › d�¦  «        ‚)Nz&Model does not contain a version for 'z'.)Úopset_importÚonnxÚdefsÚhasÚdomainÚversionr{   )rØ   r¬   rÿ   s      r   Úget_op_versionz6MinMaxCalibrater.augment_graph.<locals>.get_op_version^  s]   € Ø %Ô 2ð 0ð 0�Ý”9—=’= ¨,Ô*=Ñ>Ô>ð 0Ø'Ô/Ð/Ð/Ð/ð0åÐSÈÐSÐSÐSÑTÔTÐTr   c                 ó  •‡ — t          ˆ fd„t          ‰j        j        j        ¦  «        D ¦   «         t          ‰j        j        j        ¦  «        ¦  «        }|D ],}‰j        j        j                             ||¦  «         |dz  }Œ-d S )Nc              3   ó4   •K  — | ]\  }}‰|j         v ¯|V — Œd S rq   )rÒ   )rT   ÚiÚxrá   s      €r   ú	<genexpr>zGMinMaxCalibrater.augment_graph.<locals>.insert_nodes.<locals>.<genexpr>f  s5   øè è € ÐZÐZ‘t�q˜!À;ÐRSÔRYÐCYÐCY�ÐCYÐCYÐCYÐCYÐZÐZr   r
   )ÚnextÚ	enumerater¬   rÎ   r×   ro   Úinsert)rá   Ú	new_nodesÚindexr×   rD   s   `   €r   Úinsert_nodesz4MinMaxCalibrater.augment_graph.<locals>.insert_nodesd  s“   øø€ ÝØZÐZÐZÐZ�y¨¬Ô)9Ô)>Ñ?Ô?ÐZÑZÔZÕ\_Ð`dÔ`jÔ`pÔ`uÑ\vÔ\vñô ˆEð "ð ð �Ø”
Ô Ô%×,Ò,¨U°DÑ9Ô9Ð9Ø˜‘
��ðð r   c                 óÈ  •— d}| dz   |z   }|dz   }t           j                             || g|g||¬¦  «        }t           j                             d|‰g|g|¬¦  «        }d„ ‰j        j        j        D ¦   «         }|                     d„ ‰j        j        j        D ¦   «         ¦  «         |                     d	„ ‰j        j        j        D ¦   «         ¦  «         | |v r||          j	        j
        j        }nt          d
| ›d�¦  «        ‚‰j        �rt          ||          j	        j
        j        j        ¦  «        }	dgt#          d|	¦  «        ¢}
 ‰|‰j        ¦  «        dk     r.|j                             t          j        d|
¦  «        ¦  «         n‘t+          t-          j        ¦   «         ¦  «        }t1          j        t5          j        |
t4          j        ¬¦  «        |¦  «        }|j                             |¦  «         ‰j        j        j                             |¦  «          ‰| ||g¦  «         ‰j        j        j                             t          j        ||d g¦  «        ¦  «         d S )Nr
   Ú_Ú_Reshape)r#   rÃ   ÚReshape)ÚinputsÚoutputsrÃ   c                 ó   — i | ]
}|j         |“ŒS rR   rÂ   rÄ   s     r   rU   zNMinMaxCalibrater.augment_graph.<locals>.add_reduce_min_max.<locals>.<dictcomp>  s   € ÐMÐMÐM¨2˜2œ7 BÐMÐMÐMr   c                 ó   — i | ]
}|j         |“ŒS rR   rÂ   )rT   Úos     r   rU   zNMinMaxCalibrater.augment_graph.<locals>.add_reduce_min_max.<locals>.<dictcomp>€  s   € ÐKÐKÐK¨a ¤¨ÐKÐKÐKr   c                 ó   — i | ]
}|j         |“ŒS rR   rÂ   )rT   r  s     r   rU   zNMinMaxCalibrater.augment_graph.<locals>.add_reduce_min_max.<locals>.<dictcomp>�  s   € ÐJÐJÐJ¨a ¤¨ÐJÐJÐJr   z'Unable to guess tensor type for tensor zE, running shape inference before quantization may resolve this issue.r   rc   é   Úaxesr   )r   r   Ú	make_noder¬   rÎ   rÏ   rÐ   rÑ   rÒ   rA   rÍ   rÜ   r>   r°   ro   r   ÚdimÚrangeÚ	attributeÚappendÚmake_attributerj   ÚuuidÚuuid4r	   Ú
from_arrayr   ÚarrayÚint64rÓ   Úmake_tensor_value_info)rá   Úreduce_op_namer#   Úreduce_outputÚintermediate_outputÚreduce_nodeÚreshape_noderÞ   Ú	onnx_typeÚtensor_rankÚreduced_axesÚreduce_axes_nameÚreduce_axesr  r  Úreshape_shape_namerD   s                €€€€r   Úadd_reduce_min_maxz:MinMaxCalibrater.augment_graph.<locals>.add_reduce_min_maxl  s„  ø€ ð ˆHð (¨#Ñ-°Ñ>ˆMØ"/°*Ñ"<ÐÝœ+×/Ò/Ø  Ð0CÐ/DÈxÐ^kð 0ñ ô ˆKõ  œ;×0Ò0ØØ+Ð-?Ð@Ø&˜Ø(ð	 1ñ ô ˆLð NÐM°´Ô1AÔ1LÐMÑMÔMˆKØ×ÒÐKÐK°4´:Ô3CÔ3JÐKÑKÔKÑLÔLÐLØ×ÒÐJÐJ°4´:Ô3CÔ3IÐJÑJÔJÑKÔKÐKØ˜kÐ)Ð)Ø'¨Ô4Ô9ÔEÔO�	�	å ðZ¸kð Zð Zð Zñô ð ð Ôñ 
EÝ! +¨kÔ":Ô"?Ô"KÔ"QÔ"UÑVÔV�Ø !Ð:¥E¨!¨[Ñ$9Ô$9Ð:�à!�> .°$´*Ñ=Ô=ÀÒBÐBØÔ)×0Ò0µÔ1FÀvÈ|Ñ1\Ô1\Ñ]Ô]Ð]Ð]å'*­4¬:©<¬<Ñ'8Ô'8Ð$Ý".Ô"9½"¼(À<ÕWYÔW_Ð:`Ñ:`Ô:`ÐbrÑ"sÔ"s�KØÔ%×,Ò,Ð-=Ñ>Ô>Ð>Ø”JÔ$Ô0×7Ò7¸ÑDÔDÐDàˆL˜ {°LÐ&AÑBÔBÐBØŒJÔÔ#×*Ò*­6Ô+HÈÐXaÐdhÐciÑ+jÔ+jÑkÔkÐkÐkÐkr   Ú	ReduceMinÚ	ReduceMax©Úsave_as_external_dataN)râ   r¬   rj   r#  r$  r	   r%  r   r&  r'  rÎ   rÓ   r!  r   Úsaver­   r¯   )	rD   Útensorsr  Úreshape_shaper4  Útensorr  r  r3  s	   `     @@@r   ræ   zMinMaxCalibrater.augment_graphS  sA  øøøø€ ð ×5Ò5°d´jÑAÔA‰
ˆ�Ý ¥¤¡¤Ñ.Ô.ÐÝ$Ô/µ´¸"¸ÅRÄXÐ0NÑ0NÔ0NÐPbÑcÔcˆØŒ
ÔÔ$×+Ò+¨MÑ:Ô:Ð:ð	Uð 	Uð 	Uð	ð 	ð 	ð 	ð 	ð,	lð ,	lð ,	lð ,	lð ,	lð ,	lð ,	lð ,	lð\ ð 	4ð 	4ˆFØÐ˜v {Ñ3Ô3Ð3ØÐ˜v {Ñ3Ô3Ð3Ð3åŒ	ØŒJØÔ%Ø"&Ô"?ð	
ñ 	
ô 	
ð 	
ð 	
ð 	
r   c                 ó   — g | _         d S rq   ©rõ   rL   s    r   Úclear_collected_dataz%MinMaxCalibrater.clear_collected_data¤  ó   € Ø$&ˆÔ!Ð!Ð!r   rç   c           
      óŒ  ‡ — 	 |                      ¦   «         }|snŸ‰ j                             ˆ fd„t          ‰ j                             ¦   «         ‰ j                             d |¦  «        d¬¦  «        D ¦   «         ¦  «         ‰ j        �1t          ‰ j        ¦  «        ‰ j        k    r‰  	                    ¦   «          Œ¶t          ‰ j        ¦  «        dk    r‰ j
        €t          d¦  «        ‚‰                      ¦   «         }t          |t          ¦  «        s t          dt!          |¦  «        › d�¦  «        ‚‰  	                    ¦   «          d S )	NTc                 ó:   •— g | ]\  }}|j         ‰j        vr|nd ‘ŒS rq   )rÃ   rø   )rT   Úsess_or|   rD   s      €r   ú
<listcomp>z1MinMaxCalibrater.collect_data.<locals>.<listcomp>­  sA   ø€ ð ð ð á%˜ ð $œ[°Ô0KÐKÐK�E�EÐQUðð ð r   F©Ústrictr   úNo data is collected.z+compute_data must return a TensorsData not r7   )r‹   rõ   r!  Úzipr²   Úget_outputsÚrunrû   ro   r?  rö   r>   rì   ri   ra   rk   rA   )rD   rç   r  Úts   `   r   ré   zMinMaxCalibrater.collect_data§  sf  ø€ ð	,Ø ×)Ò)Ñ+Ô+ˆFØð ØØÔ%×,Ò,ðð ð ð å),ØÔ*×6Ò6Ñ8Ô8¸$Ô:L×:PÒ:PÐQUÐW]Ñ:^Ô:^Ðglð*ñ *ô *ðñ ô ñô ð ð Ô-Ð9Ý˜Ô1Ñ2Ô2°dÔ6SÒSÐSà×)Ò)Ñ+Ô+Ð+ð!	,õ$ ˆtÔ(Ñ)Ô)¨QÒ.Ð.°4Ô3OÐ3WÝÐ4Ñ5Ô5Ð5à×ÒÑÔˆÝ˜!�[Ñ)Ô)ð 	VÝÐTÍ$ÈqÉ'Ì'ÐTÐTÐTÑUÔUÐUØ×!Ò!Ñ#Ô#Ð#Ð#Ð#r   c                 ót  — |s|S |                      ¦   «         D �]\  }}t          |t          ¦  «        r|j        d         }|j        d         }n|\  }}t          ||         t          ¦  «        r'||         j        d         }||         j        d         }n||         \  }}| j        r!|| j        ||z
  z  z   }	|| j        ||z
  z  z   }
n t          ||¦  «        }	t          ||¦  «        }
t          |t          ¦  «        st          ||         t          ¦  «        rt          |	|
¬¦  «        ||<   �Œ|	|
f||<   �Œ|S )Nr   r
   rd   )r<   ri   r.   rM   rù   rú   ÚminÚmax)rD   Ú	old_rangeÚ	new_rangerv   r|   Úold_minÚold_maxÚnew_minÚnew_maxÚ	min_valueÚ	max_values              r   Úmerge_rangezMinMaxCalibrater.merge_rangeÂ  s]  € Øð 	ØÐà#Ÿ/š/Ñ+Ô+ð 	8ñ 	8‰JˆC�å˜%¥Ñ,Ô,ð )ØÔ+¨AÔ.�ØÔ+¨AÔ.��à#(Ñ �˜å˜) Cœ.­*Ñ5Ô5ð 2Ø# Cœ.Ô4°QÔ7�Ø# Cœ.Ô4°QÔ7��à#,¨S¤>Ñ �˜àÔ"ð 2Ø# dÔ&=ÀÈ7ÑARÑ&SÑS�	Ø# dÔ&=ÀÈ7ÑARÑ&SÑS�	�	å ¨Ñ1Ô1�	Ý ¨Ñ1Ô1�	õ ˜%¥Ñ,Ô,ð 8µ
¸9ÀS¼>Í:Ñ0VÔ0Vð 8Ý!+°9ÀiÐ!PÑ!PÔ!P�	˜#‘‘à"+¨YÐ!7�	˜#‘‘àÐr   r   c           
      ó.  ‡ ‡‡‡— t          ‰ j        ¦  «        dk    r‰ j        S ˆ fd„t          t          ‰ j        d         ¦  «        ¦  «        D ¦   «         Šˆfd„‰ j        D ¦   «         }i Š|D ]E}|                     ¦   «         D ].\  }}‰                     |g ¦  «                             |¦  «         Œ/ŒF‰‰ j        d…         Šˆfd„t          dt          ‰¦  «        d¦  «        D ¦   «         }ˆˆ fd„‰D ¦   «         }g }t          dt          ‰¦  «        d¦  «        D �]}‰ j        rHt          j
        |‰|                  d¬¦  «        }	t          j
        |‰|d	z                     d¬¦  «        }
nGt          j        |‰|                  d¬¦  «        }	t          j        |‰|d	z                     d¬¦  «        }
‰ j        rUt          j        t          j        |	¦  «        t          j        |
¦  «        gd¬¦  «        }|                     | |f¦  «         Œõ|                     |	|
f¦  «         �Œt          t           j        t%          t'          ||d
¬¦  «        ¦  «        ¦  «        }‰ j        r!‰                      ‰ j        |¦  «        ‰ _        n|‰ _        ‰ j        S )zŒ
        Compute the min-max range of tensor
        :return: dictionary mapping: {added node names: (ReduceMin, ReduceMax) pairs }
        r   c                 óX   •— g | ]&}‰j                              ¦   «         |         j        ‘Œ'S rR   )r²   rI  rÃ   )rT   r  rD   s     €r   rD  z1MinMaxCalibrater.compute_data.<locals>.<listcomp>ì  s0   ø€ ÐsÐsÐsÀQ˜Ô*×6Ò6Ñ8Ô8¸Ô;Ô@ÐsÐsÐsr   c           	      óN   •— g | ]!}t          t          ‰|d ¬¦  «        ¦  «        ‘Œ"S ©FrE  ©rƒ   rH  ©rT   r+  Úoutput_namess     €r   rD  z1MinMaxCalibrater.compute_data.<locals>.<listcomp>í  óA   ø€ ð 
ð 
ð 
à#õ •�\Ð#6¸uÐEÑEÔEÑFÔFð
ð 
ð 
r   Nc                 óR   •— g | ]#}‰|                               d ¦  «        d         ‘Œ$S )r  r   )Ú
rpartition)rT   r  Úadded_output_namess     €r   rD  z1MinMaxCalibrater.compute_data.<locals>.<listcomp>÷  s?   ø€ ð "
ð "
ð "
Ø9:Ð˜qÔ!×,Ò,¨SÑ1Ô1°!Ô4ð"
ð "
ð "
r   rc   c                 ó4   •— i | ]}|‰j         v¯|‰|         “ŒS rR   )rø   )rT   r  Úmerged_output_dictrD   s     €€r   rU   z1MinMaxCalibrater.compute_data.<locals>.<dictcomp>û  s5   ø€ ð $
ð $
ð $
Ø)*ÀAÈTÔMhÐDhÐDhˆAÐ! !Ô$ÐDhÐDhÐDhr   r$   r
   FrE  )ro   rõ   rö   r  r<   Ú
setdefaultr!  r÷   rù   r   ÚnanmeanÚnanminÚnanmaxr®   Úabsra   rm   rn   rƒ   rH  rW  )rD   Úoutput_dicts_listÚdrF   rG   Úcalibrate_tensor_namesÚmerged_added_output_dictÚpairsr  Úmin_value_arrayÚmax_value_arrayÚmax_absolute_valueÚnew_calibrate_tensors_rangerb  rd  r^  s   `            @@@r   rì   zMinMaxCalibrater.compute_dataã  s  øøøø€ õ ˆtÔ(Ñ)Ô)¨QÒ.Ð.ØÔ/Ð/àsÐsÐsÐsÍ%ÕPSÐTXÔTmÐnoÔTpÑPqÔPqÑJrÔJrÐsÑsÔsˆð
ð 
ð 
ð 
à'+Ô'@ð
ñ 
ô 
Ðð
  ÐØ"ð 	?ð 	?ˆAØŸš™	œ	ð ?ð ?‘��1Ø"×-Ò-¨a°Ñ4Ô4×;Ò;¸AÑ>Ô>Ð>Ð>ð?à)¨$Ô*@Ð*BÐ*BÔCÐð"
ð "
ð "
ð "
Ý>CÀAÅsÐK]ÑG^ÔG^Ð`aÑ>bÔ>bð"
ñ "
ô "
Ðð$
ð $
ð $
ð $
ð $
Ø.@ð$
ñ $
ô $
Ð ð ˆÝ�q�#Ð0Ñ1Ô1°1Ñ5Ô5ð 	Añ 	AˆAØÔ"ð iÝ"$¤*Ð-EÐFXÐYZÔF[Ô-\ÐcdÐ"eÑ"eÔ"e�Ý"$¤*Ð-EÐFXÐYZÐ]^ÑY^ÔF_Ô-`ÐghÐ"iÑ"iÔ"i��å"$¤)Ð,DÐEWÐXYÔEZÔ,[ÐbcÐ"dÑ"dÔ"d�Ý"$¤)Ð,DÐEWÐXYÐ\]ÑX]ÔE^Ô,_ÐfgÐ"hÑ"hÔ"h�àŒ~ð AÝ%'¤Yµ´°Ñ0GÔ0GÍÌÐP_ÑI`ÔI`Ð/aÐhiÐ%jÑ%jÔ%jÐ"Ø—’Ð1Ð1Ð3EÐFÑGÔGÐGÐGà—’˜o¨Ð?Ñ@Ô@Ð@Ñ@å&1ÝÔ$¥d­3Ð/EÀuÐUZÐ+[Ñ+[Ô+[Ñ&\Ô&\ñ'
ô '
Ð#ð Ô'ð 	GØ+/×+;Ò+;¸DÔ<XÐZuÑ+vÔ+vˆDÔ(Ð(à+FˆDÔ(àÔ+Ð+r   )Nr§   FFFrï   NF)rX   r\   r]   rj   r   r   rH   ræ   r?  r‰   ré   rW  ra   rì   Ú__classcell__©rW   s   @r   rî   rî   )  sî   ø€ € € € € ð 7;Ø3ØØ!&ØØØ!%Øð'Að 'Aà˜$‘Jð'Að  (¨œ}¨tÑ3ð'Að 'Að 'Að 'Að 'Að 'AðRO
ð O
ð O
ðb'ð 'ð 'ð$Ð(=ð $ð $ð $ð $ð6ð ð ðB3,˜kð 3,ð 3,ð 3,ð 3,ð 3,ð 3,ð 3,ð 3,r   rî   c                   óx   ‡ — e Zd Z	 	 	 	 	 	 	 	 	 dd	eez  d
ee         dz  fˆ fd„Zd„ Zd„ Zde	fd„Z
defd„Zˆ xZS )ÚHistogramCalibraterNr§   FÚ
percentileé€   é   ç-²�ïÿX@Úsamer¨   r©   c                 óX  •— t          ¦   «                              |||||¬¦  «         g | _        d| _        t	          | j        j        j        ¦  «        | _        d„ | j        j        j        D ¦   «         | _	        d| _
        || _        || _        || _        |	| _        d| _        |
| _        dS )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`
        )r©   r­   r®   r¯   Nc                 ó   — h | ]	}|j         ’Œ
S rR   rÂ   rò   s     r   rÌ   z/HistogramCalibrater.__init__.<locals>.<setcomp>=  ró   r   )rô   rH   rõ   rö   ro   r¬   rÎ   rÑ   r÷   rø   Ú	collectorÚmethodÚnum_binsÚnum_quantized_binsrw  rß   Úscenario)rD   r¨   r©   r­   r¯   r  r®   r€  r�  rw  r‚  rW   s              €r   rH   zHistogramCalibrater.__init__  s´   ø€ õ2 	‰Œ×ÒØØ"7Ø!5ØØ%=ð 	ñ 	
ô 	
ð 	
ð %'ˆÔ!Ø'+ˆÔ$Ý!$ T¤ZÔ%5Ô%<Ñ!=Ô!=ˆÔØ&YÐ&YÀÄÔAQÔAXÐ&YÑ&YÔ&YˆÔ#ØˆŒØˆŒØ ˆŒØ"4ˆÔØ$ˆŒØ$(ˆÔ!Ø ˆŒˆˆr   c                 ó  — |                       | j        ¦  «        \  | _        }| j        D ]5}|| j        vr*| j        j        j                             ||         ¦  «         Œ6t          j        | j        | j	        | j
        ¬¦  «         dS )zƒ
        make all quantization_candidates op type nodes as part of the graph output.
        :return: augmented ONNX model
        r7  N)râ   r¬   rß   rø   rÎ   rÑ   r!  r   r9  r­   r¯   )rD   rÞ   r<  s      r   ræ   z!HistogramCalibrater.augment_graphF  s›   € ð
 26×1QÒ1QÐRVÔR\Ñ1]Ô1]Ñ.ˆÔ! ;ØÔ/ð 	Dð 	DˆFØ˜TÔ8Ð8Ð8Ø”
Ô Ô'×.Ò.¨{¸6Ô/BÑCÔCÐCøåŒ	ØŒJØÔ%Ø"&Ô"?ð	
ñ 	
ô 	
ð 	
ð 	
ð 	
r   c                 ó   — g | _         d S rq   r>  rL   s    r   r?  z(HistogramCalibrater.clear_collected_dataV  r@  r   rç   c                 óð  ‡ ‡‡— d„ ‰ j                              ¦   «         D ¦   «         }d„ ‰ j                              ¦   «         D ¦   «         Š	 |                     ¦   «         }|sn”‰ j                              d|¦  «        }g }t          |¦  «        D ]L\  }}‰|         |v r(|                     t          j        |¦  «        ¦  «         Œ7|                     |¦  «         ŒM‰ j                             |¦  «         Œ«t          ‰ j        ¦  «        dk    rt          d¦  «        ‚ˆfd„‰ j        D ¦   «         }i Š|D ]E}	|	                     ¦   «         D ].\  }
}‰                     |
g ¦  «                             |¦  «         Œ/ŒFˆˆ fd„‰D ¦   «         }‰ j        s8t          ‰ j        ‰ j        ‰ j        ‰ j        ‰ j        ‰ j        ¬	¦  «        ‰ _        ‰ j                             |¦  «         ‰                      ¦   «          dS )
zy
        Entropy Calibrator collects operators' tensors as well as generates tensor histogram for each operator.
        c                 ó   — h | ]	}|j         ’Œ
S rR   rÂ   ©rT   Únode_args     r   rÌ   z3HistogramCalibrater.collect_data.<locals>.<setcomp>]  s   € ÐYÐYÐY¨X˜8œ=ÐYÐYÐYr   c                 ó   — g | ]	}|j         ‘Œ
S rR   rÂ   r‡  s     r   rD  z4HistogramCalibrater.collect_data.<locals>.<listcomp>^  s   € ÐWÐWÐW¨(˜œÐWÐWÐWr   TNr   rG  c           	      óN   •— g | ]!}t          t          ‰|d ¬¦  «        ¦  «        ‘Œ"S r[  r\  r]  s     €r   rD  z4HistogramCalibrater.collect_data.<locals>.<listcomp>t  r_  r   c                 ó4   •— i | ]}|‰j         v ¯|‰|         “ŒS rR   )rß   )rT   r  Úmerged_dictrD   s     €€r   rU   z4HistogramCalibrater.collect_data.<locals>.<dictcomp>~  s,   ø€ ÐfÐfÐf°1ÀqÈDÔLeÐGeÐGe˜Q ¨A¤ÐGeÐGeÐGer   )r  r®   r€  r�  rw  r‚  )r²   Ú
get_inputsrI  r‹   rJ  r  r!  Úcopyrõ   ro   r>   r<   re  r~  ÚHistogramCollectorr  r®   r€  r�  rw  r‚  Úcollectr?  )rD   rç   Úinput_names_setr  r  Úfixed_outputsÚoutput_indexrÑ   rj  rk  rF   rG   Úclean_merged_dictrŒ  r^  s   `            @@r   ré   z HistogramCalibrater.collect_dataY  sL  øøø€ ð ZÐY¸Ô9K×9VÒ9VÑ9XÔ9XÐYÑYÔYˆØWÐW°dÔ6H×6TÒ6TÑ6VÔ6VÐWÑWÔWˆð	<Ø ×)Ò)Ñ+Ô+ˆFØð ØØÔ(×,Ò,¨T°6Ñ:Ô:ˆGð ˆMÝ(1°'Ñ(:Ô(:ð 1ð 1Ñ$�˜fØ Ô-°Ð@Ð@Ø!×(Ò(­¬°6Ñ):Ô):Ñ;Ô;Ð;Ð;à!×(Ò(¨Ñ0Ô0Ð0Ð0àÔ%×,Ò,¨]Ñ;Ô;Ð;ð	<õ" ˆtÔ(Ñ)Ô)¨QÒ.Ð.ÝÐ4Ñ5Ô5Ð5ð
ð 
ð 
ð 
à'+Ô'@ð
ñ 
ô 
Ðð
 ˆØ"ð 	8ð 	8ˆAØŸš™	œ	ð 8ð 8‘��1Ø×&Ò& q¨"Ñ-Ô-×4Ò4°QÑ7Ô7Ð7Ð7ð8ð gÐfÐfÐfÐf¸ÐfÑfÔfÐàŒ~ð 	Ý/Ø”{Øœ.ØœØ#'Ô#:Øœ?Øœðñ ô ˆDŒNð 	Œ×ÒÐ0Ñ1Ô1Ð1à×!Ò!Ñ#Ô#Ð#Ð#Ð#r   r   c                 óˆ  — | j         st          d¦  «        ‚t          | t          ¦  «        rt          j        }ndt          | t          ¦  «        rt          j        }nBt          | t          ¦  «        rt          j	        }n t          dt          | ¦  «        › d�¦  «        ‚t          || 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.zUnknown calibrater z". This method must be overwritten.)r~  r>   ri   ÚEntropyCalibraterrm   r…   ÚPercentileCalibraterr†   ÚDistributionCalibraterr‡   rk   rA   ra   Úcompute_collection_result)rD   Úcals     r   rì   z HistogramCalibrater.compute_data�  s¶   € ð
 Œ~ð 	ZÝÐXÑYÔYÐYå�dÕ-Ñ.Ô.ð 	bÝ#Ô+ˆCˆCÝ˜Õ2Ñ3Ô3ð 	bÝ#Ô.ˆCˆCÝ˜Õ4Ñ5Ô5ð 	bÝ#Ô0ˆCˆCåÐ`µ$°t±*´*Ð`Ð`Ð`ÑaÔaÐaÝ˜3 ¤× HÒ HÑ JÔ JÑKÔKÐKr   )	Nr§   Frw  Frx  ry  rz  r{  )rX   r\   r]   rj   r   r   rH   ræ   r?  r‰   ré   ra   rì   rs  rt  s   @r   rv  rv    sÞ   ø€ € € € € ð 7;Ø3Ø!&ØØØØØØð*!ð *!à˜$‘Jð*!ð  (¨œ}¨tÑ3ð*!ð *!ð *!ð *!ð *!ð *!ðX
ð 
ð 
ð 'ð 'ð 'ð2$Ð(=ð 2$ð 2$ð 2$ð 2$ðhL˜kð Lð Lð Lð Lð Lð Lð Lð Lr   rv  c                   óP   ‡ — e Zd Z	 	 	 	 	 	 	 d	deez  dee         dz  fˆ fd„Zˆ xZS )
r–  Nr§   Fr,   rx  r¨   r©   c	           
      óZ   •— t          ¦   «                              ||||||||¬¦  «         dS )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.
        )r  r®   r€  r�  N©rô   rH   )
rD   r¨   r©   r­   r¯   r  r®   r€  r�  rW   s
            €r   rH   zEntropyCalibrater.__init__¡  sH   ø€ õ* 	‰Œ×ÒØØ!Ø Ø$ØØØØ1ð 	ñ 		
ô 		
ð 		
ð 		
ð 		
r   )Nr§   Fr,   Frx  rx  ©rX   r\   r]   rj   r   r   rH   rs  rt  s   @r   r–  r–     sx   ø€ € € € € ð 7;Ø3Ø!&ØØØØð
ð 
à˜$‘Jð
ð  (¨œ}¨tÑ3ð
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
r   r–  c                   óP   ‡ — e Zd Z	 	 	 	 	 	 	 d
deez  dee         dz  fˆ fd	„Zˆ xZS )r—  Nr§   Frw  ry  rz  r¨   r©   c	           
      óZ   •— t          ¦   «                              ||||||||¬¦  «         dS )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.
        )r  r®   r€  rw  Nr�  )
rD   r¨   r©   r­   r¯   r  r®   r€  rw  rW   s
            €r   rH   zPercentileCalibrater.__init__Ã  sH   ø€ õ* 	‰Œ×ÒØØ!Ø Ø$ØØØØ!ð 	ñ 		
ô 		
ð 		
ð 		
ð 		
r   )Nr§   Frw  Fry  rz  rž  rt  s   @r   r—  r—  Â  sx   ø€ € € € € ð 7;Ø3Ø!&ØØØØð
ð 
à˜$‘Jð
ð  (¨œ}¨tÑ3ð
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
r   r—  c                   óN   ‡ — e Zd Z	 	 	 	 	 	 d
deez  dee         dz  fˆ fd	„Zˆ xZS )r˜  Nr§   FÚdistributionrx  r{  r¨   r©   c           	      óX   •— t          ¦   «                              |||||||¬¦  «         dS )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
        )r  r€  r‚  Nr�  )	rD   r¨   r©   r­   r¯   r  r€  r‚  rW   s	           €r   rH   zDistributionCalibrater.__init__å  sE   ø€ õ. 	‰Œ×ÒØØ!Ø Ø$ØØØð 	ñ 	
ô 	
ð 	
ð 	
ð 	
r   )Nr§   Fr¢  rx  r{  rž  rt  s   @r   r˜  r˜  ä  su   ø€ € € € € ð 7;Ø3Ø!&ØØØð
ð 
à˜$‘Jð
ð  (¨œ}¨tÑ3ð
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
ð 
r   r˜  c                   óR   — e Zd ZdZej        d„ ¦   «         Zej        d„ ¦   «         ZdS )ÚCalibrationDataCollectorzL
    Base class for collecting data for calibration-based quantization.
    c                 ó   — t           ‚)z‹
        Generate informative data based on given data.
            name_to_arr : dict
                tensor name to NDArray data
        r’   ©rD   Úname_to_arrs     r   r�  z CalibrationDataCollector.collect  s
   € õ "Ð!r   c                 ó   — t           ‚)z?
        Get the optimal result among collection data.
        r’   rL   s    r   r™  z2CalibrationDataCollector.compute_collection_result  s
   € õ
 "Ð!r   N)rX   r\   r]   Ú__doc__r¡   r¢   r�  r™  rR   r   r   r¥  r¥    s\   € € € € € ðð ð 	Ôð"ð "ñ Ôð"ð 	Ôð"ð "ñ Ôð"ð "ð "r   r¥  c                   ól   — e Zd ZdZd„ Zd„ Zd„ Zd„ Zd„ Zd„ Z	d„ Z
d	„ Zd
„ Zedd„¦   «         Zd„ Zd„ ZdS )r�  a`  
    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
    c                 óh   — i | _         || _        || _        || _        || _        || _        || _        d S rq   )Úhistogram_dictr  r®   r€  r�  rw  r‚  )rD   r  r®   r€  r�  rw  r‚  s          r   rH   zHistogramCollector.__init__&  s9   € Ø ˆÔØˆŒØ"ˆŒØ ˆŒØ"4ˆÔØ$ˆŒØ ˆŒˆˆr   c                 ó   — | j         S rq   )r­  rL   s    r   Úget_histogram_dictz%HistogramCollector.get_histogram_dict/  s   € ØÔ"Ð"r   c                 óò   — t          d¦  «         | j        dv r|                      |¦  «        S | j        dk    r1| j        r|                      |¦  «        S |                      |¦  «        S t          d¦  «        ‚)Nz/Collecting tensor data and making histogram ...>   r,   r¢  rw  úDOnly 'entropy', 'percentile' or 'distribution' methods are supported)Úprintr  Úcollect_valuer®   Úcollect_absolute_valuer>   r§  s     r   r�  zHistogramCollector.collect2  sƒ   € ÝÐ?Ñ@Ô@Ð@ð Œ;Ð5Ð5Ð5Ø×%Ò% kÑ2Ô2Ð2ØŒ[˜LÒ(Ð(ØŒ~ð 7Ø×2Ò2°;Ñ?Ô?Ð?à×)Ò)¨+Ñ6Ô6Ð6åÐcÑdÔdÐdr   c                 óò  — |                      ¦   «         D �]`\  }}t          |t          ¦  «        r€|D ]9}t          |t          j        ¦  «        sJ dt          |¦  «        › d|›�¦   «         ‚Œ:d„ |D ¦   «         }t          |¦  «        dk    sJ d|› d|›�¦   «         ‚t          j        |¦  «        }n>t          |t          j        ¦  «        s"t          dt          |¦  «        › d|›�¦  «        ‚|}| 	                    ¦   «         }|j
        dk    r)t          j        |¦  «        }t          j        |¦  «        }n6t          j        d|j        ¬¦  «        }t          j        d|j        ¬¦  «        }t          j        |¦  «        }|| j        vrgt          j        || j        ¬¦  «        \  }	}
|
                     |j        ¦  «        }
|j        t          j        k    s
J d	¦   «         ‚|	|
||f| j        |<   �ŒÛ| j        |         }|d
         }|d         }t+          |d¦  «        sJ dt          |¦  «        › �¦   «         ‚t+          |d¦  «        sJ dt          |¦  «        › �¦   «         ‚|d         }|d         }t          j        |¦  «        }||d         k    rI|d         |d         z
  }t          j        |d         |z   ||z   |¦  «        }t          j        ||f¦  «        }t          j        ||¬¦  «        \  }	}
|
                     |j        ¦  «        }
|	dt          |¦  «        …xx         |z  cc<   |j        t          j        k    s
J d	¦   «         ‚|	|
t1          ||¦  «        t3          ||¦  «        f| j        |<   �ŒbdS )z5
        Collect histogram on absolute value
        r8   z for tensor=c                 ó   — h | ]	}|j         ’Œ
S rR   r   )rT   Úas     r   rÌ   z<HistogramCollector.collect_absolute_value.<locals>.<setcomp>I  s   € Ð4Ð4Ð4 a˜!œ'Ð4Ð4Ð4r   r
   z6The calibration expects only one element type but got r   r   )r5   zMonly float32 or float16 is supported, every constant must be explicitly typedrc   rf   r   z'old_min should be a numpy array but is rý   N)r<   ri   r9   r   ÚndarrayrA   ro   r%   r>   ÚflattenÚsizerg  rh  r&  r   Úabsoluter­  Ú	histogramr€  r&   Úfloat64r@   ÚarangeÚhstackrM  rN  )rD   r¨  r<  Údata_arrÚarrÚdtypesÚdata_arr_nprU  rV  r3   r4   Úold_histogramrQ  rR  Úold_histÚold_hist_edgesÚ	temp_amaxÚwidthÚnew_bin_edgess                      r   r´  z)HistogramCollector.collect_absolute_valueA  sË  € ð !,× 1Ò 1Ñ 3Ô 3ð 4	sñ 4	sÑˆF�HÝ˜(¥DÑ)Ô)ð 'Ø#ð mð m�CÝ% c­2¬:Ñ6Ô6ÐlÐlÐ8lÍ4ÐPSÉ9Ì9Ð8lÐ8lÐbhÐ8lÐ8lÑlÔlÐ6ÐlØ4Ð4¨8Ð4Ñ4Ô4�Ý˜6‘{”{ aÒ'Ð'Ð'ØkÈVÐkÐkÐagÐkÐkñ (Ô'Ð'õ !œj¨Ñ2Ô2��Ý ­"¬*Ñ5Ô5ð 'Ý Ð!ZµD¸±N´NÐ!ZÐ!ZÐPVÐ!ZÐ!ZÑ[Ô[Ð[à&�Ø%×-Ò-Ñ/Ô/ˆKØÔ !Ò#Ð#ÝœI kÑ2Ô2�	ÝœI kÑ2Ô2�	�	åœH Q¨kÔ.?Ð@Ñ@Ô@�	ÝœH Q¨kÔ.?Ð@Ñ@Ô@�	åœ+ kÑ2Ô2ˆKà˜TÔ0Ð0Ð0å#%¤<°À$Ä-Ð#PÑ#PÔ#PÑ ��jØ'×.Ò.¨{Ô/@ÑAÔA�
Ø"Ô(­B¬JÒ6Ð6Ð6Øcñ 7Ô6Ð6ð 04°ZÀÈIÐ.V�Ô# FÑ+Ñ+à $Ô 3°FÔ ;�Ø'¨Ô*�Ø'¨Ô*�Ý˜w¨Ñ0Ô0ÐkÐkÐ2kÕ\`ÐahÑ\iÔ\iÐ2kÐ2kÑkÔkÐ0Ý˜w¨Ñ0Ô0ÐkÐkÐ2kÕ\`ÐahÑ\iÔ\iÐ2kÐ2kÑkÔkÐ0Ø(¨Ô+�Ø!.¨qÔ!1�ÝœI kÑ2Ô2�	Ø˜~¨bÔ1Ò1Ð1à*¨1Ô-°¸qÔ0AÑA�Eå$&¤I¨n¸RÔ.@À5Ñ.HÈ)ÐV[ÑJ[Ð]bÑ$cÔ$c�MÝ%'¤Y°ÀÐ/NÑ%OÔ%O�NÝ#%¤<°À.Ð#QÑ#QÔ#QÑ ��jØ'×.Ò.¨{Ô/@ÑAÔA�
Ø�_•s˜8‘}”}�_Ð%Ð%Ô%¨Ñ1Ð%Ð%Ñ%Ø"Ô(­B¬JÒ6Ð6Ð6Øcñ 7Ô6Ð6ð 04°ZÅÀWÈiÑAXÔAXÕZ]Ð^eÐgpÑZqÔZqÐ.r�Ô# FÑ+Ñ+ði4	sð 4	sr   c           	      ó¼  — |                      ¦   «         D �]E\  }}t          j        |¦  «        }|                     ¦   «         }|j        dk    r)t          j        |¦  «        }t          j        |¦  «        }n6t          j        d|j        ¬¦  «        }t          j        d|j        ¬¦  «        }t          j        t          t          |¦  «        t          |¦  «        ¦  «        |j        ¬¦  «        }|| j        v r0| j        |         }|                      |||||¦  «        | j        |<   �Œt          j        || j        | |f¬¦  «        \  }}	||	|||f| j        |<   �ŒGdS )z1
        Collect histogram on real value
        r   r   ©r  N)r<   r   r%   r¹  rº  rg  rh  r&  r   rN  ri  r­  Úmerge_histogramr¼  r€  )
rD   r¨  r<  rÀ  rU  rV  Ú	thresholdrÄ  r3   r4   s
             r   r³  z HistogramCollector.collect_value{  sc  € ð !,× 1Ò 1Ñ 3Ô 3ð 	ñ 	ÑˆF�HÝ”z (Ñ+Ô+ˆHØ×'Ò'Ñ)Ô)ˆHàŒ}˜qÒ Ð ÝœI hÑ/Ô/�	ÝœI hÑ/Ô/�	�	åœH Q¨h¬nÐ=Ñ=Ô=�	ÝœH Q¨h¬nÐ=Ñ=Ô=�	åœ¥¥S¨¡^¤^µS¸±^´^Ñ!DÔ!DÈHÌNÐ[Ñ[Ô[ˆIà˜Ô,Ð,Ð,Ø $Ô 3°FÔ ;�Ø.2×.BÒ.BØ! 8¨Y¸	À9ñ/ô /�Ô# FÑ+Ñ+õ $&¤<°¸$¼-ÐQZÐPZÐ\eÐOfÐ#gÑ#gÔ#gÑ ��jàØØØØð/�Ô# FÑ+Ñ+ð)	ð 	r   c                 óR  — |\  }}}}	}
||
k    rPt          j        |t          |¦  «        |
 |
f¬¦  «        \  }}||z   |t          ||¦  «        t	          |	|¦  «        |
fS |
dk    r0t          j        |t          |¦  «        | |f¬¦  «        \  }}||z  }nqt          |¦  «        }d|
z  |z  }t          ||
z
  |z  dz   ¦  «        }|d|z  z   }||z  |
z   }t          j        ||| |f¬¦  «        \  }}||||z
  …xx         |z  cc<   ||t          ||¦  «        t	          |	|¦  «        |fS )NrË  r   rc   r
   )r   r¼  ro   rM  rN  r£   )rD   rÄ  rÀ  rS  rT  Únew_thresholdrÅ  rÆ  rQ  rR  Úold_thresholdÚnew_histr  r3   r4   Úold_num_binsÚ
old_strideÚhalf_increased_binsÚnew_num_binss                      r   rÌ  z"HistogramCollector.merge_histogram›  sŒ  € ØFSÑCˆ�> 7¨G°]à˜MÒ)Ð)Ýœ, xµ°X±´ÈÀ~ÐWdÐFeÐfÑfÔf‰KˆH�aà˜8Ñ#ØÝ�G˜WÑ%Ô%Ý�G˜WÑ%Ô%Øðð ð  Ò!Ð!Ý#%¤<°½#¸h¹-¼-ÐQ^ÐP^Ð`mÐOnÐ#oÑ#oÔ#oÑ ��jØ˜Ñ ��å" 8™}œ}�Ø Ñ.°Ñ=�
Ý&)¨=¸=Ñ+HÈZÑ*WÐZ[Ñ*[Ñ&\Ô&\Ð#Ø+¨aÐ2EÑ.EÑE�Ø 3°jÑ @À=Ñ P�Ý#%¤<°¸,ÐP]È~Ð_lÐNmÐ#nÑ#nÔ#nÑ ��jØÐ(¨<Ð:MÑ+MÐMÐNÐNÔNÐRZÑZÐNÐNÑNàØÝ�G˜WÑ%Ô%Ý�G˜WÑ%Ô%Øðð r   c                 óf  — | j         rt          | j         ¦  «        dk    rt          d¦  «        ‚t          d| j        ›d�¦  «         | j        dk    r|                      ¦   «         S | j        dk    r|                      ¦   «         S | j        dk    r|                      ¦   «         S t          d¦  «        ‚)	Nr   z=Histogram has not been collected. Please run collect() first.z0Finding optimal threshold for each tensor using z algorithm ...r,   rw  r¢  r±  )r­  ro   r>   r²  r  Úcompute_entropyÚcompute_percentileÚcompute_distributionrL   s    r   r™  z,HistogramCollector.compute_collection_result»  sº   € ØÔ"ð 	^¥c¨$Ô*=Ñ&>Ô&>À!Ò&CÐ&CÝÐ\Ñ]Ô]Ð]ÝÐ^ÀÄÐ^Ð^Ð^Ñ_Ô_Ð_àŒ;˜)Ò#Ð#Ø×'Ò'Ñ)Ô)Ð)ØŒ[˜LÒ(Ð(Ø×*Ò*Ñ,Ô,Ð,ØŒ[˜NÒ*Ð*Ø×,Ò,Ñ.Ô.Ð.åÐcÑdÔdÐdr   c                 óÄ  — | j         dk     s| j         dk    rt          d¦  «        ‚| j        }| j         }i }t          dt	          |¦  «        › �¦  «         t          d| j        › �¦  «         t          dd|z
  › d|› d	�¦  «         |                     ¦   «         D �]Å\  }}|d         }|d
         }|                     ¦   «         }t          j	        ||z  ¦  «        }	| j
        r_t          j        |	|dz  ¦  «        }
t          j        ||
         |j        ¬¦  «         t          j        ||
         |j        ¬¦  «        f||<   nzd|z
  dz  }t          j        |	d|z
  ¦  «        }
t          j        |	|¦  «        }t          j        ||         |j        ¬¦  «        t          j        ||
         |j        ¬¦  «        f||<   |d         }|d         }||         d         |k     r|||         d
         f||<   ||         d
         |k    r||         d         |f||<   g ||         ¢|d d…         ¢R ||<   t          j                             dd¦  «        dv rt#          ||¦  «         �ŒÇ|S )Nr   éd   z<Invalid percentile. Must be in range 0 <= percentile <= 100.úNumber of tensors : úNumber of histogram bins : zPercentile : (g      Y@ú,ú)r
   r   g      i@r"   rc   rf   ÚQUANTIZATION_DEBUGÚ0©r
   Ú1)rw  r>   r­  r²  ro   r€  r<   r(   r   Úcumsumr®   Úsearchsortedr&  r   ÚosÚenvironÚgetr   )rD   r­  rw  Úthresholds_dictr<  r¼  r3   r4   ÚtotalÚcdfÚ	idx_rightÚpercent_to_cut_one_sideÚidx_leftrU  rV  s                  r   rØ  z%HistogramCollector.compute_percentileÉ  s™  € ØŒ?˜QÒÐ $¤/°CÒ"7Ð"7ÝÐ[Ñ\Ô\Ð\àÔ,ˆØ”_ˆ
àˆåÐ:¥S¨Ñ%8Ô%8Ð:Ð:Ñ;Ô;Ð;ÝÐ;¨D¬MÐ;Ð;Ñ<Ô<Ð<ÝÐA˜u zÑ1ÐAÐA°JÐAÐAÐAÑBÔBÐBà!/×!5Ò!5Ñ!7Ô!7ð 	-ñ 	-ÑˆF�IØ˜Q”<ˆDØ" 1œˆJØ—H’H‘J”JˆEÝ”)˜D 5™LÑ)Ô)ˆCØŒ~ð ÝœO¨C°¸eÑ1CÑDÔD�	õ ”X˜j¨Ô3¸:Ô;KÐLÑLÔLÐLÝ”H˜Z¨	Ô2¸*Ô:JÐKÑKÔKð+� Ñ'Ð'ð
 ,1°:Ñ+=ÀÑ*FÐ'ÝœO¨C°Ð7NÑ1NÑOÔO�	Ýœ?¨3Ð0GÑHÔH�å”H˜Z¨Ô1¸Ô9IÐJÑJÔJÝ”H˜Z¨	Ô2¸*Ô:JÐKÑKÔKð+� Ñ'ð " !œˆIØ! !œˆIØ˜vÔ& qÔ)¨IÒ5Ð5Ø+4°oÀfÔ6MÈaÔ6PÐ*Q� Ñ'Ø˜vÔ& qÔ)¨IÒ5Ð5Ø+:¸6Ô+BÀ1Ô+EÀyÐ*Q� Ñ'Ø&K¨¸Ô(?Ð&KÀ$ÀrÈÀrÄ(Ð&KÐ&KˆO˜FÑ#åŒz�~Š~Ð2°CÑ8Ô8¸HÐDÐDÝ˜4 Ñ,Ô,Ð,ùàÐr   c                 óÌ  — | j         }| j        }i }t          dt          |¦  «        › �¦  «         t          d| j        › d�¦  «         t          d| j        › �¦  «         |                     ¦   «         D ]p\  }}|                      ||¦  «        }|||<   g |¢|d d…         ¢R ||<   t          j         	                    dd¦  «        dv rt          |d	         |d
         ¦  «         Œq|S )NrÜ  rÝ  z: (The number may increase depends on the data it collects)zNumber of quantized bins : rc   rà  rá  râ  r   r
   )r­  r�  r²  ro   r€  r<   Úget_entropy_thresholdræ  rç  rè  r   )rD   r­  r�  ré  r<  r¼  Úoptimal_thresholds          r   r×  z"HistogramCollector.compute_entropy÷  s  € ØÔ,ˆØ!Ô4ÐàˆåÐ:¥S¨Ñ%8Ô%8Ð:Ð:Ñ;Ô;Ð;ÝÐu¨D¬MÐuÐuÐuÑvÔvÐvÝÐE¨DÔ,CÐEÐEÑFÔFÐFà!/×!5Ò!5Ñ!7Ô!7ð 	7ð 	7ÑˆF�IØ $× :Ò :¸9ÐFXÑ YÔ YÐØ&7ˆO˜FÑ#Ø&JÐ(9Ð&J¸IÀbÀqÀb¼MÐ&JÐ&JˆO˜FÑ#õ Œz�~Š~Ð2°CÑ8Ô8¸HÐDÐDÝ˜9 Qœ<¨°1¬Ñ6Ô6Ð6øàÐr   r
   c                 ó"  — |dk    rt          d|› d�¦  «        ‚|d d…         |dd …         z   dz  }|dk    rš| |z                       ¦   «         |                      ¦   «         z  }| |dz  z                       ¦   «         |                      ¦   «         z  |dz  z
  dz  }t          j        ||j        ¬¦  «        t          j        ||j        ¬¦  «        fS t          |¦  «        |k    r³t          |¦  «        dz  dk    r�| ||z  z                       ¦   «         |                      ¦   «         z  }| ||z  |z
  dz  z                       ¦   «         |                      ¦   «         z  dz  }t          j        ||j        ¬¦  «        t          j        ||j        ¬¦  «        fS t          j        |¦  «        |z  }d|t          j        |¦  «        <   d|t          j        |¦  «        <   t          j        |¦  «        |z  |z  }| |z                       ¦   «         |                      ¦   «         z  }| |dz  z                       ¦   «         |                      ¦   «         z  |dz  z
  dz  }t          j        ||j        ¬¦  «        t          j        ||j        ¬¦  «        fS )	Nr   zpower=z <= 0 is invalid.rý   r
   g      à?rc   r   )	r>   r(   r   r&  r   r£   ri  ÚisnanÚisinf)r3   r4   Úpowerr€   r/   r0   Úfacts          r   Ú_avg_stdzHistogramCollector._avg_std  sW  € à�AŠ:ˆ:ÝÐ> eÐ>Ð>Ð>Ñ?Ô?Ð?Ø˜S˜b˜S”/ J¨q¨r¨r¤NÑ2°cÑ9ˆØ�AŠ:ˆ:Ø˜&‘=×%Ò%Ñ'Ô'¨$¯(ª(©*¬*Ñ4ˆCØ˜6 1™9Ñ$×)Ò)Ñ+Ô+¨d¯hªh©j¬jÑ8¸3À¹6ÑAÀcÑIˆCÝ”8˜C zÔ'7Ð8Ñ8Ô8½"¼(À3ÈjÔN^Ð:_Ñ:_Ô:_Ð_Ð_Ýˆu‰:Œ:˜ÒÐ¥3 u¡:¤:°¡>°QÒ#6Ð#6Ø˜& %™-Ñ'×,Ò,Ñ.Ô.°·²±´Ñ;ˆCØ˜F E™M¨CÑ/°AÑ5Ñ5×:Ò:Ñ<Ô<¸t¿xºx¹z¼zÑIÈcÑQˆCÝ”8˜C zÔ'7Ð8Ñ8Ô8½"¼(À3ÈjÔN^Ð:_Ñ:_Ô:_Ð_Ð_åŒv�f‰~Œ~ Ñ&ˆØ ˆ�RŒX�d‰^Œ^ÑØ ˆ�RŒX�d‰^Œ^ÑÝ”˜‘” 5Ñ(¨4Ñ/ˆØ�f‰}×!Ò!Ñ#Ô# d§h¢h¡j¤jÑ0ˆØ�v˜q‘yÑ ×%Ò%Ñ'Ô'¨$¯(ª(©*¬*Ñ4°s¸A±vÑ=À#ÑEˆÝŒx˜ :Ô#3Ð4Ñ4Ô4µb´h¸sÈ*ÔJZÐ6[Ñ6[Ô6[Ð[Ð[r   c           
      óŽ  — | j         dk     rt          d¦  «        ‚| j        }i }t          dt	          |¦  «        › �¦  «         t          d| j         › �¦  «         t          d| j        ›d�¦  «         |                     ¦   «         D �]=\  }}|d         }|d         }|j        t          j	        k    sJ ‚| j        d	k    r|  
                    ||d¬
¦  «        \  }}n6| j        dk    r|  
                    ||d¬
¦  «        \  }}nt          d¦  «        ‚|j        t          j	        k    sJ ‚|j        t          j	        k    sJ ‚|j        t          j	        k    sJ ‚t          |||||                     ¦   «         |                     ¦   «         ¬¦  «        ||<   t          j                             dd¦  «        dv rt#          ||¦  «         �Œ?|S )Ni   z3Invalid num_bins. Must be in range 512 <= num_bins.rÜ  rÝ  zScenario : rß  r   r
   r{  )rõ  Úp3gUUUUUUÕ?z,Invalid scenario. Must be in {'same', 'p3'}.)r/   r0   r3   r4   r1   r2   rà  rá  râ  )r€  r>   r­  r²  ro   r‚  r<   r   r   r½  r÷  r.   rM  rN  ræ  rç  rè  r   )	rD   r­  ré  r<  r¼  r3   r4   Úavg_coefÚstd_coefs	            r   rÙ  z'HistogramCollector.compute_distribution"  sà  € ØŒ=˜3ÒÐÝÐRÑSÔSÐSàÔ,ˆØˆåÐ:¥S¨Ñ%8Ô%8Ð:Ð:Ñ;Ô;Ð;ÝÐ;¨D¬MÐ;Ð;Ñ<Ô<Ð<ÝÐ.˜DœMÐ.Ð.Ð.Ñ/Ô/Ð/à!/×!5Ò!5Ñ!7Ô!7ð 	-ñ 	-ÑˆF�IØ˜Q”<ˆDØ" 1œˆJàÔ#¥r¤zÒ1Ð1Ð1Ð1ØŒ} Ò&Ð&Ø%)§]¢]°4¸È1 ]Ñ%MÔ%MÑ"�˜(˜(Ø” $Ò&Ð&Ø%)§]¢]°4¸È9 ]Ñ%UÔ%UÑ"�˜(˜(å Ð!OÑPÔPÐPØ”>¥R¤ZÒ/Ð/Ð/Ð/Ø”>¥R¤ZÒ/Ð/Ð/Ð/ØÔ#¥r¤zÒ1Ð1Ð1Ð1Ý&0ØØØØ%Ø!—~’~Ñ'Ô'Ø"ŸšÑ(Ô(ð'ñ 'ô 'ˆO˜FÑ#õ Œz�~Š~Ð2°CÑ8Ô8¸HÐDÐDÝ˜4 Ñ,Ô,Ð,ùàÐr   c           	      óˆ  ‡— |d         }|d         }|j         }|dz  }|dz  }|d         j        Št          j        ||z
  dz   ¦  «        }ˆfd„t	          |j         ¦  «        D ¦   «         }	t	          ||dz   d¦  «        D �]:}
||
z
  }t          ||
z   dz   |¦  «        }||         ||         f|	|
|z
  <   t          j        |||…         ¦  «        }|                     ¦   «         }t          |d|…         ¦  «        }t          ||d…         ¦  «        }|dxx         |z  cc<   |dxx         |z  cc<   |dk     	                    t          j
        ¦  «        }t          j        |t          j
        ¬¦  «        }|j         |z  }t	          |¦  «        D ]&}||z  }||z   }t          |||…         ¦  «        ||<   Œ'|dxx         t          |||z  d…         ¦  «        z  cc<   t          j        |j         t          j
        ¬¦  «        }t	          |¦  «        D ]9}||z  }||z   }t          |||…         ¦  «        }|dk    r||         |z  |||…<   Œ:t          |¦  «        }t          |¦  «        }|�|€!t          j        t          j        ‰¬¦  «        }n$t          j        t          ||¦  «        ‰¬¦  «        }|||
|z
  <   �Œ<t          j        |¦  «        }|	|         }|d         }|d         }|d         |k     r
||d         f}|d         |k    r
|d         |f}t!          |d         d	¦  «        sJ ‚t!          |d         d	¦  «        sJ ‚|S )
aF  Given a dataset, find the optimal threshold for quantizing it.
        The reference distribution is `q`, and the candidate distribution is `p`.
        `q` is a truncated version of the original distribution.
        Ref: http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf
        r   r
   rc   c                 óh   •— g | ].}t          j        d ‰¬¦  «        t          j        d ‰¬¦  «        f‘Œ/S )r   r   )r   r&  )rT   r  r   s     €r   rD  z<HistogramCollector.get_entropy_threshold.<locals>.<listcomp>X  s<   ø€ ÐnÐnÐnÈq•r”x ¨Ð/Ñ/Ô/µ´¸!À5Ð1IÑ1IÔ1IÐJÐnÐnÐnr   Nrý   r   rf   r   )rº  r   r   Úzerosr  rM  rŽ  Údeepcopyr(   r&   r'  r   r&  r   r,   Úargminr@   )rD   r¼  r�  r3   r4   r€  Úzero_bin_indexÚnum_half_quantized_binÚkl_divergenceÚ
thresholdsr  rœ   r�   Úsliced_distributionÚpÚleft_outliers_countÚright_outliers_countÚnonzerosÚquantized_binsÚnum_merged_binsr  ÚstartÚendÚqÚnormÚdivÚmin_kl_divergence_idxrñ  rU  rV  r   s                                 @r   rð  z(HistogramCollector.get_entropy_thresholdJ  sÈ  ø€ ð ˜Œ|ˆØ˜q”\ˆ
Ø”9ˆØ! Q™ˆØ!3°qÑ!8Ðà˜!”Ô"ˆÝœ Ð2HÑ!HÈ1Ñ!LÑMÔMˆØnÐnÐnÐnÕTYÐZgÔZlÑTmÔTmÐnÑnÔnˆ
õ  Ð-¨~ÀÑ/AÀ1ÑEÔEð .	<ñ .	<ˆAØ(¨1Ñ,ˆKÝ˜N¨QÑ.°Ñ2°HÑ=Ô=ˆIà6@ÀÔ6MÈzÐZcÔOdÐ5eˆJ�qÐ1Ñ1Ñ2å"&¤-°°[ÀÐ5JÔ0KÑ"LÔ"LÐð $×(Ò(Ñ*Ô*ˆAÝ"% d¨<¨K¨<Ô&8Ñ"9Ô"9ÐÝ#& t¨I¨J¨JÔ'7Ñ#8Ô#8Ð ØˆaˆDˆDŒDÐ'Ñ'ˆDˆD‰DØˆbˆEˆEŒEÐ)Ñ)ˆEˆE‰Eð ˜Qš—’¥r¤xÑ0Ô0ˆHõ  œXÐ&8ÅÄÐIÑIÔIˆNØ1Ô6Ð:LÑLˆOõ Ð1Ñ2Ô2ð Lð L�Ø Ñ/�Ø˜oÑ-�Ý(+Ð,?ÀÀcÀ	Ô,JÑ(KÔ(K�˜uÑ%Ð%Ø˜2ÐÐÔ¥#Ð&9Ð:LÈÑ:^Ð:`Ð:`Ô&aÑ"bÔ"bÑbÐÐÑõ ”˜œ¥r¤xÐ0Ñ0Ô0ˆAÝÐ1Ñ2Ô2ð @ð @�Ø Ñ/�Ø˜oÑ-�å˜8 E¨# IÔ.Ñ/Ô/�Ø˜1’9�9Ø#1°%Ô#8¸4Ñ#?�A�e˜C�i‘Løå# AÑ&Ô&ˆAÝ# AÑ&Ô&ˆAØˆy˜A˜IÝ”h�rœv¨UÐ3Ñ3Ô3��å”h�w q¨!™}œ}°EÐ:Ñ:Ô:�Ø8;ˆM˜!Ð4Ñ4Ñ5Ñ5å "¤	¨-Ñ 8Ô 8ÐØ&Ð'<Ô=ÐØ˜a”Lˆ	Ø˜a”Lˆ	Ø˜QÔ )Ò+Ð+Ø!*Ð,=¸aÔ,@Ð AÐØ˜QÔ )Ò+Ð+Ø!2°1Ô!5°yÐ AÐÝÐ(¨Ô+¨WÑ5Ô5Ð5Ð5Ð5ÝÐ(¨Ô+¨WÑ5Ô5Ð5Ð5Ð5Ø Ð r   N)r
   )rX   r\   r]   rª  rH   r¯  r�  r´  r³  rÌ  r™  rØ  r×  Ústaticmethodr÷  rÙ  rð  rR   r   r   r�  r�    sú   € € € € € ðð ð!ð !ð !ð#ð #ð #ðeð eð eð8sð 8sð 8sðtð ð ð@ð ð ð@eð eð eð,ð ,ð ,ð\ð ð ð* ð\ð \ð \ñ „\ð\ð*&ð &ð &ðPX!ð X!ð X!ð X!ð X!r   r�  r§   Fr¬   r©   c                 ó
  — d }|t           j        k    rˆ|                     dd¦  «        }|                     dd¦  «        }	|                     dd¦  «        }
|                     dd ¦  «        }|                     dd¦  «        }t          | |||||	|
||¬¦	  «	        }�n#|t           j        k    rY|                     d	d
¦  «        }|                     dd
¦  «        }|                     dd¦  «        }t          | ||||||¬¦  «        }nº|t           j        k    rY|                     d	d¦  «        }|                     dd¦  «        }|                     dd¦  «        }t          | ||||||¬¦  «        }nQ|t           j        k    rA|                     d	d¦  «        }|                     dd¦  «        }t          | |||||¬¦  «        }|r3| 
                    ¦   «          |r||_        |                     ¦   «          |S t          d|› �¦  «        ‚)Nr®   Frù   rú   rï   rû   r°   )r¯   r®   rù   rú   rû   r°   r€  rx  r�  )r¯   r®   r€  r�  ry  rw  rz  T)r¯   r®   r€  rw  r‚  r{  )r¯   r€  r‚  zUnsupported calibration method )rm   rn   rè  rî   r…   r–  r†   r—  r‡   r˜  ræ   r³   rµ   r>   )r¬   r©   r­   Úcalibrate_methodr¯   r¹   Úextra_optionsÚ
calibratorr®   rù   rú   rû   r°   r€  r�  rw  r‚  s                    r   Úcreate_calibratorr  ¥  sm  € ð €JØÕ,Ô3Ò3Ð3à!×%Ò% k°5Ñ9Ô9ˆ	Ø&×*Ò*Ð+;¸UÑCÔCˆØ*×.Ò.Ð/CÀTÑJÔJÐØ#0×#4Ò#4Ð5OÐQUÑ#VÔ#VÐ Ø#×'Ò'¨°uÑ=Ô=ˆÝ%ØØ!Ø Ø%=ØØ)Ø1Ø%=Ø#ð

ñ 

ô 

ˆ
‰
ð 
Õ.Ô6Ò	6Ð	6à ×$Ò$ Z°Ñ5Ô5ˆØ*×.Ò.Ð/CÀSÑIÔIÐØ!×%Ò% k°5Ñ9Ô9ˆ	Ý&ØØ!Ø Ø%=ØØØ1ð
ñ 
ô 
ˆ
ˆ
ð 
Õ.Ô9Ò	9Ð	9à ×$Ò$ Z°Ñ6Ô6ˆØ"×&Ò& |°VÑ<Ô<ˆ
Ø!×%Ò% k°4Ñ8Ô8ˆ	Ý)ØØ!Ø Ø%=ØØØ!ð
ñ 
ô 
ˆ
ˆ
ð 
Õ.Ô;Ò	;Ð	;à ×$Ò$ Z°Ñ6Ô6ˆØ ×$Ò$ Z°Ñ8Ô8ˆå+ØØ!Ø Ø%=ØØð
ñ 
ô 
ˆ
ð ð Ø× Ò Ñ"Ô"Ð"Øð 	7Ø-6ˆJÔ*Ø×+Ò+Ñ-Ô-Ð-ØÐå
ÐIÐ7GÐIÐIÑ
JÔ
JÐJr   )Nr   ),r¡   rŽ  rÙ   ræ  r#  Úcollections.abcr   Úenumr   Úpathlibr   Únumpyr   r   r   r   r   r	   rº   Úquant_utilsr   r   r   r¸  r   Úfloatr£   r,   r.   ra   rm   ÚABCMetar‰   r¦   rî   rv  r–  r—  r˜  r¥  r�  rn   rj   r  rR   r   r   ú<module>r     sö  ðð €
€
€
Ø €€€Ø Ð Ð Ð Ø 	€	€	€	Ø €€€Ø $Ð $Ð $Ð $Ð $Ð $Ø Ð Ð Ð Ð Ð Ø Ð Ð Ð Ð Ð à Ð Ð Ð Ø €€€Ø >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >Ð >à Ð Ð Ð à UÐ UÐ UÐ UÐ UÐ UÐ UÐ UÐ UÐ Uð�”ð  ¤ð °´
ð ð ð ð ð" Øð	ð Ø
Œ
ðà
Œ
ðð �$‰,ðð ð	ð
 „Zðð ð ð ð8 ð  ð  ð  ð  ñ  ô  ð  ðF1ð 1ð 1ð 1ð 1ñ 1ô 1ð 1ðhð ð ð ð ˜ñ ô ð ð"ð "ð "ð "ð " c¤kð "ñ "ô "ð "ð4k"ð k"ð k"ð k"ð k"ñ k"ô k"ð k"ð\m,ð m,ð m,ð m,ð m,�~ñ m,ô m,ð m,ð`DLð DLð DLð DLð DL˜.ñ DLô DLð DLðN
ð 
ð 
ð 
ð 
Ð+ñ 
ô 
ð 
ðD
ð 
ð 
ð 
ð 
Ð.ñ 
ô 
ð 
ðD 
ð  
ð  
ð  
ð  
Ð0ñ  
ô  
ð  
ðF"ð "ð "ð "ð "¨¬ð "ñ "ô "ð "ð,E!ð E!ð E!ð E!ð E!Ð1ñ E!ô E!ð E!ðT 37Ø/Ø&Ô-Ø"ØØðNKð NKØ�‰:ðNKà# Cœ=¨4Ñ/ðNKð NKð NKð NKð NKð NKr   