
    
3j                     `   S SK r S SKrS SKJr  S SKJrJr  S SKJr  S SK	J
r
  \R                  " \5      rS\S\\   S\4S	 jr     SS
\R"                  S\\   S\S\\   S\\   S\4S jjrS rSS jr\r         SS
\R"                  S\S\\   S\\   S\S\S\S\S\\   S\4S jjrg)    N)islice)
CollectionOptional)nn)group_parametersnamepatternsreturnc                 .   ^  [        U 4S jU 5       5      $ )zACheck if parameter name matches any pattern (supports wildcards).c              3   R   >#    U  H  n[         R                   " TU5      v   M     g 7fN)fnmatch).0patternr   s     R/home/wildlama/miniconda3/lib/python3.13/site-packages/timm/optim/_param_groups.py	<genexpr>#_matches_pattern.<locals>.<genexpr>   s     FX'wtW--Xs   $')any)r   r	   s   ` r   _matches_patternr      s    FXFFF    modelweight_decayweight_decay_exclude_1dno_weight_decay_listfallback_listfallback_no_weight_decayc                 T   U(       a  [        U5      [        U5      -  nU(       + n/ n/ n/ n	/ n
U R                  5        H  u  pUR                  (       d  M  [        X5      nU(       a-  U(       a  U
R	                  U5        MD  U	R	                  U5        MW  [        X5      nU=(       a'    UR
                  S:*  =(       d    UR                  S5      nU(       d  U(       a-  U(       a  U
R	                  U5        M  U	R	                  U5        M  U(       a  UR	                  U5        M  UR	                  U5        M     / nU	(       a  UR	                  U	SS.5        U(       a  UR	                  XqS.5        U
(       a  UR	                  U
SSS.5        U(       a  UR	                  XSS.5        U$ )N   z.bias        )paramsr   T)r    r   use_fallback)setnamed_parametersrequires_gradr   appendndimendswith)r   r   r   r   r   r   no_wddecaydecay_fallbackno_decayno_decay_fallbackr   paramis_fallbackmatches_patternis_excluded_1dgroupss                    r   param_groups_weight_decayr2      sS     M*S1E-FFEENH--/"" 't;!((/& /tJO4d%**/:cT]][bMcN%,,U3OOE* "))%0LL'9 0< F2>?EF!2BX\]^_cdeMr   c                 >   ^ ^ [        T 5      m [        U U4S jS5      $ )Nc                  .   > [        [        T T5      5      $ r   )tupler   itsizes   r   <lambda>_group.<locals>.<lambda>Q   s    fR./r    )iterr6   s   ``r   _groupr=   O   s    	bB/44r   c                    S n[        U S0 5      R                  SS 5      n/ n/ nU R                  5        H6  u  pxU" Xt5      (       a  UR                  U5        M%  UR                  U5        M8     [	        U5      n	Ub  X* -  * n[        [        XQ5      5      n[	        U5      n
[        U5       VVVs0 s H  u  pU  H  owU_M     M     nnnnUR                  U Vs0 s H  owU
_M     sn5        U$ s  snnnf s  snf )Nc                     U(       d  g[        U[        [        45      (       a*  [        U Vs/ s H  o R	                  U5      PM     sn5      $ U R	                  U5      $ s  snf )NT)
isinstancer5   listr   
startswith)nhphpis      r   _in_head#auto_group_layers.<locals>._in_headU   sN    UDM**R8RcS)R899<<## 9s   Apretrained_cfg
classifier)	getattrgetr#   r%   lenrA   r=   	enumerateupdate)r   layers_per_group
num_groupsrF   head_prefixnames_trunk
names_headrC   _num_trunk_layersnum_trunk_groupsil	layer_maps                 r   auto_group_layersrZ   T   s    $ %!126::<NKKJ&&( ( 8 8
!k>P>PQR>S ) ;'-<=vk<=K;'"+K"8D"8$!!QA!"8ID:>:a)):>? E>s   1C1C8layer_decay	min_scaleno_opt_scaleverbosec
                   ^^^ U(       a  [        U5      [        U5      -  n0 n
0 n[        U S5      (       a  [        X R                  SS9SS9nO[	        U 5      n[        UR                  5       5      S-   nUS-
  m[        UUU4S j[        U5       5       5      nU R                  5        H  u  nnUR                  (       d  M  [        X5      nU(       a  UR                  S::  d  [        X5      (       a  SnS	nOS
nUnUR                  UT5      nUU   nU(       a  UU:  a	  SUl
        M  U(       a  SOSnSUUU4-  nUU;  a"  UUU/ S.U
U'   UU/ S.UU'   U(       a  SUU   S'   U
U   S   R                  U5        UU   S   R                  U5        M     U	(       a*  SSKn[         R#                  SUR%                  U
SS9-  5        [        UR                  5       5      $ )z
Parameter groups for layer-wise lr decay & weight decay
Based on BEiT: https://github.com/microsoft/unilm/blob/master/beit/optim_factory.py#L58
group_matcherF)coarseT)reverser   c              3   H   >#    U  H  n[        TTTU-
  -  5      v   M     g 7fr   )max)r   rW   r[   	layer_maxr\   s     r   r   +param_groups_layer_decay.<locals>.<genexpr>   s&     bPa1I{y1}'EFFPas   "r+   r   r)   	_fallback zlayer_%d_%s%s)lr_scaler   r!   param_names)ri   r   r    r!   rj   r    r   Nzparameter groups: 
%s   )indent)r"   hasattrr   r`   rZ   rd   valuesrA   ranger#   r$   r   r&   rK   r%   json_loggerinfodumps)r   r   r   r   r   r   r[   r\   r]   r^   param_group_namesparam_groupsrY   
num_layerslayer_scalesr   r-   r.   g_decay
this_decaylayer_id
this_scalefallback_suffix
group_namerp   re   s         ``                 @r   param_groups_layer_decayr~   q   s   "  M*S1E-FFLuo&&$U,?,?u,?,MW[\	 &e,	Y%%'(1,JQIbPUV`PabbL--/e"" 't; $

a<LT<h<h GJG%J==y1!(+
J5"'E)4+"$'?'KK
\)& * +!	-j) ' *(L$
 ;?Z(8*%m4;;DAZ *11%8U 0X -

;LUV
0WWX##%&&r   )gh㈵>Tr;   r;   F)   N)	g?r;   r;   FTg      ?r   NF)r   logging	itertoolsr   typingr   r   torchr   timm.modelsr   	getLogger__name__rq   strboolr   Modulefloatr2   r=   rZ   
_layer_mapr~   r;   r   r   <module>r      sX      '  ( 

H
%G3 G*S/ Gd G )-(,02)+).:yy:uo: "&: )o	:
 "#: #':x5
4 

 #02)+).(, (,Q'yyQ'Q' )oQ' "#	Q'
 #'Q' "&Q' Q' Q' uoQ' Q'r   