
    khR                         d dl Z d dlmZ d dlmZmZ d dlmZ d dlm	Z	  e j                  e      Z	 	 ddej                  dedee   fd	Zd
 ZddZeZ	 	 	 	 	 	 ddej                  dedee   dededee   defdZy)    N)islice)
CollectionOptional)nn)group_parametersmodelweight_decayno_weight_decay_listc                    t        |      }g }g }| j                         D ]Y  \  }}|j                  s|j                  dk  s|j	                  d      s||v r|j                  |       I|j                  |       [ |dd||dgS )N   z.bias        )paramsr	   )setnamed_parametersrequires_gradndimendswithappend)r   r	   r
   decayno_decaynameparams          T/var/www/teggl/fontify/venv/lib/python3.12/site-packages/timm/optim/_param_groups.pyparam_groups_weight_decayr      s    
 34EH--/  e""::?dmmG4@T8TOOE"LL  R0,79 9    c                 <     t                t         fdd      S )Nc                  .    t        t                     S N)tupler   itsizes   r   <lambda>z_group.<locals>.<lambda>%   s    fR./ r    )iterr    s   ``r   _groupr&   #   s    	bB/44r   c                    d }t        | di       j                  dd       }g }g }| j                         D ]0  \  }} |||      r|j                  |      n|j                  |       2 t	        |      }	||	| z   }t        t        ||            }t	        |      }
t        |      D ci c]  \  }}|D ]  }||  }}}}|j                  |D ci c]  }||
 c}       |S c c}}}w c c}w )Nc                     |syt        |t        t        f      r(t        |D cg c]  }| j	                  |       c}      S | j	                  |      S c c}w )NT)
isinstancer   listany
startswith)nhphpis      r   _in_headz#auto_group_layers.<locals>._in_head)   sI    UDM*R8cS)899<<## 9s   Apretrained_cfg
classifier)	getattrgetr   r   lenr*   r&   	enumerateupdate)r   layers_per_group
num_groupsr0   head_prefixnames_trunk
names_headr-   _num_trunk_layersnum_trunk_groupsil	layer_maps                 r   auto_group_layersrC   (   s   $ %!126::<NKKJ&&( T1 (K 8
!k>P>PQR>ST ;'-*<=vk+;<=K;'"+K"8DD$!Q!DQADDID:>aa))>? E>s   +C%
C,weight_decay_exclude_1dlayer_decayend_layer_decayverbosec                    t        |      }i }i }t        | d      rt        | | j                  d      d      }	nt	        |       }	t        |	j                               dz   }
|
dz
  t        fdt        |
      D              }| j                         D ]  \  }}|j                  s|r|j                  dk  s||v rd}d	}nd
}|}|	j                  |      }d||fz  }||vr||   }||g d||<   ||g d||<   ||   d   j                  |       ||   d   j                  |        |r-ddl}t        j!                  d|j#                  |d      z         t        |j                               S )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   .   K   | ]  }|z
  z    y wr   r$   ).0r@   rE   	layer_maxs     r   	<genexpr>z+param_groups_layer_decay.<locals>.<genexpr>^   s     R1	A6Rs   r   r   r   zlayer_%d_%s)lr_scaler	   param_names)rP   r	   r   rQ   r   r   Nzparameter groups: 
%s   )indent)r   hasattrr   rI   rC   maxvaluesr*   ranger   r   r   r4   r   json_loggerinfodumps)r   r	   r
   rD   rE   rF   rG   param_group_namesparam_groupsrB   
num_layerslayer_scalesr   r   g_decay
this_decaylayer_id
group_name
this_scalerX   rN   s       `               @r   param_groups_layer_decayre   E   s    34Luo&$UE,?,?u,?,MW[\	 &e,	Y%%'(1,JQIRj@QRRL--/ 9e"" $

aDDX<X GJG%J==y1"h%88
\)%h/J& *!-j) ' *(L$ 	*%m4;;DAZ *11%8;9> -

;LUV
0WWX##%&&r   )gh㈵>r$   )   N)g?r$   Tg      ?NF)logging	itertoolsr   typingr   r   torchr   timm.modelsr   	getLogger__name__rY   Modulefloatstrr   r&   rC   
_layer_mapboolre   r$   r   r   <module>rs      s      '  ( '

H
%
 #029yy99 )o9,5
4 

 #02(, +/>'yy>'>' )o>' "&	>'
 >' "%>' >'r   