Ë
    –\;j¹3  ã                   óâ   — 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
mZmZmZmZmZmZmZ  eee j&                  d¬«      Zd	Zd
Zd„ Zd„ Z G d„ d«      Z G d„ d«      Z	 	 	 	 dd„Zy)é    N)Ústaticé   )Ú
get_loggeré   )Ú_channelwise_quant_axis1_opsÚbias_correction_wÚcalculate_quant_cos_errorÚdequant_tensorÚload_variable_dataÚquant_tensorÚset_variable_dataÚstable_sigmoidz&%(asctime)s-%(levelname)s: %(message)s)Úfmtgš™™™™™¹¿gš™™™™™ñ?c                 óª   — t        j                  t         j                  j                  j	                  | «      t
        t        z
  z  t        z   dd¬«      S )Nr   r   )ÚminÚmax)ÚpaddleÚclipÚnnÚ
functionalÚsigmoidÚZETAÚGAMMA©Úalpha_vs    úlG:\00. PROJECTS\API\Inventory\templateJSON\kerjaOCR\Lib\site-packages\paddle/static/quantization/adaround.pyÚcompute_soft_roundingr   ,   s?   € Ü�;‰;Ü�	‰	×Ñ×$Ñ$ WÓ-´¼±Ñ>ÄÑFØØôð ó    c                 ón   — t        j                  t        | «      t        t        z
  z  t        z   dd¬«      S )Nr   r   )Úa_minÚa_max)Únpr   r   r   r   r   s    r   Úcompute_soft_rounding_npr#   4   s-   € Ü�7‰7Ü�wÓ¤4¬%¡<Ñ0´5Ñ8ÀÈôð r   c                   ó&   — e Zd Zdd„Zd„ Zd„ Zd„ Zy)ÚAdaRoundLossc                 ó    — || _         || _        y ©N)Údefault_reg_paramÚdefault_beta_range)ÚselfÚ	reg_paramr)   s      r   Ú__init__zAdaRoundLoss.__init__;   s   € Ø!*ˆÔØ"4ˆÕr   c                 ó®   — t         j                  j                  j                  ||«      }t        j                  t        j
                  |d¬«      «      }|S )Néÿÿÿÿ)Úaxis)r   r   r   Úsquare_error_costÚmeanÚsum)r*   Úada_quantized_outputÚorig_outputÚsquare_costÚ
recon_losss        r   Úcompute_recon_losszAdaRoundLoss.compute_recon_loss?   sB   € Ü—i‘i×*Ñ*×<Ñ<Ø  +ó
ˆô —[‘[¤§¡¨K¸bÔ!AÓBˆ
ØÐr   c                 ó^   ‡ ‡‡— ˆˆˆ fd„}t         j                  j                  |d„ |«      }|S )Nc                  óÈ   •— t        ‰«      } t        j                  t        j                  t        j                  d| z  dz
  «      ‰«       dz   «      }‰j
                  |z  }|S )Nr   r   )r   r   r2   ÚpowÚabsr(   )Úh_vÚreg_termÚ
round_lossr   Úbetar*   s      €€€r   Úround_loss_fnz6AdaRoundLoss.compute_round_loss.<locals>.round_loss_fnG   s\   ø€ ä'¨Ó0ˆCô —z‘zÜ—‘œFŸJ™J q¨3¡w°¡{Ó3°TÓ:Ð:¸QÑ>óˆHð
 ×/Ñ/°(Ñ:ˆJàÐr   c                  ó4   — t        j                  dgdd¬«      S )Nr   Úfloat32ç        )ÚshapeÚdtypeÚ
fill_value)r   Úfull© r   r   Ú<lambda>z1AdaRoundLoss.compute_round_loss.<locals>.<lambda>X   s   € ”F—K‘K q c°ÀsÕKr   )r   r   Úcond)r*   r   Ú
warm_startr?   r@   r>   s   `` `  r   Úcompute_round_losszAdaRoundLoss.compute_round_lossF   s-   ú€ ö	ô —Y‘Y—^‘^ØÙKØó
ˆ
ð Ðr   c                 ó®   — | j                   \  }}||z  }||z
  ||z
  z  }|d||z
  z  dt        j                  |t        j                  z  «      z   z  z   }|S )Ng      à?r   )r)   r"   ÚcosÚpi)	r*   Úmax_iterÚcur_iterrK   Ú
start_betaÚend_betaÚwarm_start_end_iterÚrel_iterr?   s	            r   Úcompute_betazAdaRoundLoss.compute_beta^   sx   € à#×6Ñ6Ñˆ
�Hð )¨8Ñ3Ðð Ð2Ñ2ØÐ*Ñ*ñ
ˆð ˜# ¨hÑ!6Ñ7Ø”—‘�x¤"§%¡%Ñ'Ó(Ñ(ñ
ñ 
ˆð ˆr   N)g{®Gáz„?)é   r   )Ú__name__Ú
__module__Ú__qualname__r,   r7   rL   rV   rH   r   r   r%   r%   :   s   „ ó5òòó0r   r%   c                   óH   — e Zd Z	 	 	 	 	 d
d„Zd„ Zd„ Zd„ Zd„ Zd„ Zd„ Z	d	„ Z
y)ÚAdaRoundNc                 ó"  — || _         || _        d| _        d| _        d| _        t        «       | _        || _        || _        || _	        d| _
        |t        v rd| _
        || _        |dz   | _        | j                  |j                  «       ||«       y )Ngš™™™™™¹?é   rC   r   r   ú.alpha)Úis_trainÚnum_iterationsrK   Úweight_bitsÚoffsetr%   Úadaround_lossÚori_weight_tensorÚscaleÚscopeÚ
quant_axisr   Úweight_var_nameÚ
alpha_nameÚinitialize_alphaÚcopy)r*   rf   Úweight_tensorrg   ri   Úweight_op_typer`   ra   s           r   r,   zAdaRound.__init__q   s�   € ð !ˆŒØ,ˆÔØˆŒØˆÔØˆŒÜ)›^ˆÔØ!.ˆÔØˆŒ
ØˆŒ
ØˆŒØÔ9Ñ9ØˆDŒOØ.ˆÔØ)¨HÑ4ˆŒØ×Ñ˜m×0Ñ0Ó2°E¸?ÕKr   c                 óh  — t        ||| j                  ¬«      }t        j                  |«      }||z
  }t        j                  t
        t        z
  |t        z
  z  dz
  «       }t        j                  |j                  d|dz   t        j                  j                  j                  |«      ¬«      | _        y)zN
        Initializes alpha parameter, same shape as the weight tensor
        ©rh   r   rB   r_   )rD   rE   ÚnameÚdefault_initializerN)r   rh   r"   ÚfloorÚlogr   r   r   Úcreate_parameterrD   r   ÚinitializerÚAssignr   )r*   Útensorrf   Úvar_nameÚtensor_scaleÚtensor_floorÚalphas          r   rk   zAdaRound.initialize_alpha‹   s�   € ô $ F¨E¸d¿o¹oÔNˆÜ—x‘x Ó-ˆØ Ñ,ˆÜ—‘œ¤™¨&´5©.Ñ9¸AÑ=Ó>Ð>ˆÜ×.Ñ.Ø—+‘+ØØ˜HÑ$Ü &§	¡	× 5Ñ 5× <Ñ <¸UÓ Cô	
ˆ�r   c                 óŒ   — t        | j                  || j                  |«       |j                  |||gd| j                  ¬«      }|S )NT©ÚprogramÚfeedÚ
fetch_listÚreturn_numpyrg   )r   rg   ri   Úrun)r*   r   ÚplaceÚexeÚdataÚfp32_fetch_listÚweight_tensor_dequantÚadaround_out_tensors           r   Ú)_calculate_output_with_adarounded_weightsz2AdaRound._calculate_output_with_adarounded_weightsš   sT   € ô 	Ø�J‰J˜˜t×3Ñ3Ð5Jô	
ð "Ÿg™gØØØ'Ð(ØØ—*‘*ð &ó 
Ðð #Ð"r   c                 ó&  — t        | j                  | j                  «      }t        |«      }t	        | j
                  j                  «       | j                  | j                  ¬«      }t        j                  |«      }t        j                  ||«      }|S ©Nrp   )r   rg   rj   r#   r   re   rl   rf   rh   r"   rs   Úadd)r*   Únp_alphaÚh_alpharz   rm   Úweight_tensor_quants         r   Ú_calculate_quant_weightz AdaRound._calculate_quant_weightª   su   € Ü% d§j¡j°$·/±/ÓBˆÜ*¨8Ó4ˆô $Ø×"Ñ"×'Ñ'Ó)Ø�J‰JØ—‘ô
ˆô Ÿ™ Ó.ˆô !Ÿf™f ]°GÓ<ÐØ"Ð"r   c                 ó„   — | j                  «       }t        || j                  z   | j                  | j                  ¬«      }|S rŒ   )r‘   r
   rc   rf   rh   )r*   r�   rˆ   s      r   Ú_calculate_adarounded_weightsz&AdaRound._calculate_adarounded_weights»   s@   € Ø"×:Ñ:Ó<Ðô !/Ø $§+¡+Ñ-Ø�J‰JØ—‘ô!
Ðð
 %Ð$r   c                 ó&   — | j                  «       }|S r'   )r‘   )r*   r�   s     r   Úupdate_final_weightszAdaRound.update_final_weightsÆ   s   € Ø"×:Ñ:Ó<ÐØ"Ð"r   c                 ó¢   — | j                   j                  | j                  ||«      }| j                   j                  ||«      }||z   }|||dœ}|S )N)Úlossr>   r6   )rd   rL   r   r7   )	r*   r?   rK   r‰   Úorig_out_tensorr>   r6   r—   Úlossess	            r   Úget_losszAdaRound.get_lossÊ   sc   € Ø×'Ñ'×:Ñ:Ø�L‰L˜* dó
ˆ
ð ×'Ñ'×:Ñ:Ø ó
ˆ
ð ˜JÑ&ˆàØ$Ø$ñ
ˆð
 ˆr   c                 ó¤   — || j                   | j                  z  k  }| j                  j                  | j                   || j                  «      }||fS r'   )ra   rK   rd   rV   )r*   Úcur_iterationrK   r?   s       r   Úupdate_beta_warmzAdaRound.update_beta_warmÙ   sO   € Ø" T×%8Ñ%8¸4¿?¹?Ñ%JÑJˆ
Ø×!Ñ!×.Ñ.Ø×Ñ °·±ó
ˆð �ZÐÐr   )NNNTéè  )rX   rY   rZ   r,   rk   rŠ   r‘   r“   r•   rš   r�   rH   r   r   r\   r\   p   s<   „ ð
 ØØØØóLò4
ò#ò #ò"	%ò#òó r   r\   c                 óÆ  — |d   j                   }i }|j                  «       D �]÷  \  }}t        j                  d|› �«       ||   }t	        ||«      }||   }d }|j                  «       j                  D ]E  }|j                  dk(  sŒ|j                  ||«       |j                  «       j                  |«      }|}ŒG t        j                  «       }d|_        t        j                  «       }t        j                  «       }t        j                  ||«      5  t        j                   j"                  j%                  «       5  t'        ||||||	¬«      }t        j(                  dd|j*                  z   d¬	«      }t        j(                  d
d|j*                  z   d¬	«      }t        j(                  dddgd¬	«      }t        j(                  dddgd¬	«      }|j-                  ||||«      }t        j.                  j1                  |
¬«      }|d   } |j3                  | «       d d d «       d d d «       |j5                  |«       t7        j6                  «       }!|!}"t9         | «       «      D �]x  \  }#}$|!}"t7        j6                  «       }!|j5                  ||$|gd|¬«      }%j;                  «       }&|j=                  ||||$||&«      }'t?        |%d   |'d   «      }(|r|(dkD  rt        j                  d«        nò|jA                  |#«      \  })}*|%d   |'d   |)|*dœ}+|j5                  ||+jC                  «       D �,cg c]  },|,j                   ‘Œ c},d¬«      }-t        j                  djE                  |#|
tG        jH                  |-d   «      tG        jH                  |-d   «      tG        jH                  |-d   «      |!|"z
  «      «       tJ        jL                  jO                  «        |#|	k(  s�Œy n jQ                  «       ||<   |r*tS        |||   ||jT                  |jV                  ¬«      ||<   ~�Œú |jY                  «       D ]  }t[        |||||   «       Œ y # 1 sw Y   �Œ/xY w# 1 sw Y   �Œ4xY wc c},w )Nr   zStart adaround op: Úfetchr   )rg   ri   rn   ra   r˜   )r.   rB   )rq   rD   rE   r‰   r?   r.   rK   Úbool)Úlearning_rater—   Tr~   g®Gáz®ï?z)The cosine error is small, skip training.)r˜   r‰   r?   rK   )r€   r�   r‚   zUIter {:d}, lr {:.5f}, loss {:.5f}, loss_round {:.5f}, loss_recon {:.5f}, time {:.5f}sr   )rb   ).rq   ÚitemsÚ_loggerÚinfor   Úglobal_blockÚopsÚtypeÚ_rename_inputÚvarr   ÚExecutionStrategyÚnum_iteration_per_drop_scopeÚProgramÚprogram_guardr   ÚutilsÚunique_nameÚguardr\   r†   rD   rš   Ú	optimizerÚAdamÚminimizerƒ   ÚtimeÚ	enumerater“   rŠ   r	   r�   ÚvaluesÚformatr"   r1   ÚsysÚstdoutÚflushr•   r   rh   rb   Úkeysr   ).Údata_loaderÚfp32_programr�   r…   rg   r„   Úquantized_op_pairsÚweight_op_pairsÚ
scale_dictra   ÚlrÚbias_correctionÚ	fast_modeÚfetch_op_nameÚfinal_weight_tensor_quant_dictri   Úquant_op_out_namern   Úweight_var_tensorrf   r‡   Ú_opÚexec_strategyÚstartup_programÚtrain_programÚadaroundr˜   r‰   Úbeta_tensorÚwarm_start_tensorÚtrain_fetches_lossr²   r—   Ú
start_timeÚprev_start_timeÚir†   Únp_orig_out_tensorÚadaround_weight_tensor_dequantÚnp_adaround_out_tensorÚ	cos_errorr?   rK   Ú	feed_dictÚvÚouts.                                                 r   Úrun_adaroundrÛ   á   sy  € ð ˜q‘M×&Ñ&€MØ%'Ð"Ø.@×.FÑ.F×.HÑ*ˆÐ*Ü�‰Ð*¨?Ð*;Ð<Ô=Ø(¨Ñ9ˆä.¨u°oÓFÐØ˜?Ñ+ˆØˆØ×,Ñ,Ó.×2Ô2ˆCØ�x‰x˜7Ó"Ø×!Ñ! -Ð1BÔCØ".×";Ñ";Ó"=×"AÑ"AØ%ó#�ð !2‘ð 3ô ×0Ñ0Ó2ˆØ56ˆÔ2Ü Ÿ.™.Ó*ˆÜŸ™Ó(ˆÜ×!Ñ! -°ÕAÜ—‘×)Ñ)×/Ñ/Õ1ä#ØØ%ØØ$3Ø#1Ø#1ô�ô #)§+¡+Ø*Ø /×"7Ñ"7Ñ7Ø#ô#�ô
 '-§k¡kØ.Ø /×"7Ñ"7Ñ7Ø#ô'Ð#ô
 %Ÿk™kØ¨¨A w°iô�ô %+§K¡KØ%¨b°!¨W¸Fô%Ð!ð &.×%6Ñ%6ØØ%Ø'Ø#ó	&Ð"ô #×,Ñ,×1Ñ1ÀÐ1ÓC�	Ø)¨&Ñ1�Ø×"Ñ" 4Ô(÷G 2÷ BðJ 	�‰�Ô ä—Y‘Y“[ˆ
Ø$ˆÜ ¡£×/‰GˆAˆtØ(ˆOÜŸ™›ˆJà!$§¡Ø$ØØ+Ð,Ø!Øð ")ó "Ðð ×6Ñ6Ó8ð +ð ×BÑBØ ØØØØ#Ø2óð #ô 2Ø" 1Ñ%Ð'=¸aÑ'@óˆIñ ˜Y¨Ò-Ü—‘ÐHÔIÙØ'×8Ñ8¸Ó;ÑˆD�*à#5°aÑ#8Ø'=¸aÑ'@ØØ(ñ	ˆIð —'‘'ØØØ,>×,EÑ,EÔ,GÓHÑ,G q˜AŸF›FÐ,GÑHØ!ð	 ó ˆCô �L‰LØg×nÑnØØÜ—G‘G˜C ™F“OÜ—G‘G˜C ™F“OÜ—G‘G˜C ™F“OØ Ñ0óô	ô �J‰J×ÑÔØ�NÔ"Ùðu 0ðz ×)Ñ)Ó+ð 	'Øñ	
ñ Ü>OØ!Ø.¨Ñ?ØØ×#Ñ#Ø$×0Ñ0ô?Ð*¨?Ñ;ò ðK /IðP .×2Ñ2Ö4ˆÜØØØØ*¨?Ñ;õ		
ñ 5÷g 2Ñ1ú÷ BÑAüòh Is+   Ä)QÄ=CQÈQÌ(QÑQÑ	QÑQ	)rž   gü©ñÒMbP?FT)Úloggingr¹   rµ   Únumpyr"   r   r   Ú
log_helperr   r¯   r   r   r	   r
   r   r   r   r   rX   ÚINFOr¤   r   r   r   r#   r%   r\   rÛ   rH   r   r   Ú<module>rà      s…   ðó Û 
Û ã ã Ý å #÷	÷ 	ó 	ñ Øˆg�l‰lÐ Hô€ð 	€Ø
€òò÷3ñ 3÷ln ñ n ðv ØØØô_
r   