
    ^jH                        d Z ddlZddlZddlmZ ddlmZmZmZm	Z	m
Z
mZmZ ddlZddlmZ ddlmZmZ ddlmZmZmZ dd	lmZ dd
lmZ ddlmZmZ ddlmZ dgZ G d dej@                        Z! G d dejD                        Z# G d dejD                        Z$ G d dejD                        Z% G d dejD                        Z&d Z'd(dZ(d)dZ) e e)d       e)d       e)d       e)d       e)dd       e)dd       e)dd       e)dd      d      Z*ed(de&fd        Z+ed(de&fd!       Z,ed(de&fd"       Z-ed(de&fd#       Z.ed(de&fd$       Z/ed(de&fd%       Z0ed(de&fd&       Z1ed(de&fd'       Z2y)*a   Pooling-based Vision Transformer (PiT) in PyTorch

A PyTorch implement of Pooling-based Vision Transformers as described in
'Rethinking Spatial Dimensions of Vision Transformers' - https://arxiv.org/abs/2103.16302

This code was adapted from the original version at https://github.com/naver-ai/pit, original copyright below.

Modifications for timm by / Copyright 2020 Ross Wightman
    N)partial)ListOptionalSequenceTupleUnionTypeAny)nnIMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STD)trunc_normal_	to_2tuplecalculate_drop_path_rates   )build_model_with_cfg)feature_take_indices)register_modelgenerate_default_cfgs)BlockPoolingVisionTransformerc                       e Zd ZdZdeej                  ej                  f   deej                  ej                  f   fdZy)SequentialTuplezI This module exists to work around torchscript typing issues list -> listxreturnc                 $    | D ]
  } ||      } |S N )selfr   modules      Z/var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/timm/models/pit.pyforwardzSequentialTuple.forward#   s     	Fq	A	    N)__name__
__module____qualname____doc__r   torchTensorr#   r   r$   r"   r   r   !   s?    Su||U\\9: uU\\SXS_S_E_?` r$   r   c                        e Zd Z	 	 	 	 	 	 	 ddededededee   dededeee      d	eee	j                        f fd
Zdeej                  ej                  f   deej                  ej                  f   fdZ xZS )Transformerbase_dimdepthheads	mlp_ratiopool	proj_drop	attn_dropdrop_path_prob
norm_layerc                 N   |
|d}t         |           ||z  }|| _        |	r	 |	|fi |nt        j                         | _        t        j                  t        |      D cg c]2  }t        d|||d||||   t        t        j                  d      d|4 c} | _        y c c}w )NdevicedtypeTư>)eps)dim	num_headsr0   qkv_biasr2   r3   	drop_pathr5   r   )super__init__r1   r   Identitynorm
Sequentialranger   r   	LayerNormblocks)r    r-   r.   r/   r0   r1   r2   r3   r4   r5   r8   r9   dd	embed_dimi	__class__s                  r"   rA   zTransformer.__init__*   s     /u$		3=Jy/B/2;;=	mm 5\&#   
###(+"2<<T:
 
&# $ &#s   !7B"r   r   c                    |\  }}|j                   d   }| j                  | j                  ||      \  }}|j                   \  }}}}|j                  d      j                  dd      }t	        j
                  ||fd      }| j                  |      }| j                  |      }|d d d |f   }|d d |d f   }|j                  dd      j                  ||||      }||fS )Nr      )r<   )	shaper1   flatten	transposer)   catrC   rG   reshape)r    r   
cls_tokenstoken_lengthBCHWs           r"   r#   zTransformer.forwardL   s    :!''*99  IIa4MAzWW
1aIIaL""1a(IIz1o1-IIaLKKNq-<-'(
aKK1%%aAq1*}r$   )N        rY   NNNN)r%   r&   r'   intfloatr   r
   r   r	   r   ModulerA   r   r)   r*   r#   __classcell__rK   s   @r"   r,   r,   )   s     #'!!4848 $ $  $ 	 $
  $ 3- $  $  $ %T%[1 $ !bii1 $Du||U\\9: uU\\SXS_S_E_?` r$   r,   c            	       v     e Zd Z	 	 	 ddedededef fdZdeej                  ej                  f   fdZ	 xZ
S )	Pooling
in_featureout_featurestridepadding_modec           	          ||d}t         |           t        j                  ||f|dz   |dz  |||d|| _        t        j
                  ||fi || _        y )Nr7   r   rM   )kernel_sizepaddingrc   rd   groups)r@   rA   r   Conv2dconvLinearfc)	r    ra   rb   rc   rd   r8   r9   rH   rK   s	           r"   rA   zPooling.__init__a   sq     /II	
 
aK%	
 	
	 ))J:r:r$   r   c                 N    | j                  |      }| j                  |      }||fS r   )rj   rl   )r    r   	cls_tokens      r"   r#   zPooling.forwardy   s'    IIaLGGI&	)|r$   )zerosNN)r%   r&   r'   rZ   strrA   r   r)   r*   r#   r]   r^   s   @r"   r`   r`   `   sW     !(;; ; 	;
 ;0uU\\5<<-G'H r$   r`   c                   L     e Zd Z	 	 	 	 	 	 d	dedededededef fdZd Z xZS )
ConvEmbeddingin_channelsout_channelsimg_size
patch_sizerc   rg   c	                    ||d}	t         
|           |}t        |      | _        t        |      | _        t        j                  | j                  d   d|z  z   | j                  d   z
  |z  dz         | _        t        j                  | j                  d   d|z  z   | j                  d   z
  |z  dz         | _        | j                  | j                  f| _	        t        j                  ||f|||dd|	| _        y )Nr7   r   rM   r   T)rf   rc   rg   bias)r@   rA   r   ru   rv   mathfloorheightwidth	grid_sizer   ri   rj   )r    rs   rt   ru   rv   rc   rg   r8   r9   rH   rK   s             r"   rA   zConvEmbedding.__init__   s     /!(+#J/jj$--"2Q["@4??STCU"UY_!_bc!cdZZq!1AK!?$//RSBT!TX^ ^ab bc
++tzz2II
 #
 
	r$   c                 (    | j                  |      }|S r   )rj   r    r   s     r"   r#   zConvEmbedding.forward   s    IIaLr$   )         r   NN)r%   r&   r'   rZ   rA   r#   r]   r^   s   @r"   rr   rr      s[    
   

 
 	

 
 
 
<r$   rr   c            #           e Zd ZdZ	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d)dededededee   dee   dee   d	ed
ededededededededef" fdZ	d Z
ej                  j                  d        Zej                  j                  d*d       Zej                  j                  d*d       Zdej$                  fdZd+d
edee   fdZ	 	 	 	 	 d,dej,                  deeeee   f      dededed edeeej,                     eej,                  eej,                     f   f   fd!Z	 	 	 d-deeee   f   d"ed#efd$Zd% Zd.d&edej,                  fd'Zd( Z xZS )/r   z Pooling-based Vision Transformer

    A PyTorch implement of 'Rethinking Spatial Dimensions of Vision Transformers'
        - https://arxiv.org/abs/2103.16302
    ru   rv   rc   	stem_type	base_dimsr.   r/   r0   num_classesin_chansglobal_pool	distilled	drop_ratepos_drop_drateproj_drop_rateattn_drop_ratedrop_path_ratec           
         t         |           ||d}|dv sJ || _        || _        |d   |d   z  }|	| _        |
| _        || _        |rdnd| _        g | _        t        |
||||fi || _
        t        j                  t        j                  d|| j                  j                  | j                  j                   fi |      | _        t        j                  t        j                  d| j                  |fi |      | _        t        j&                  |      | _        g }t+        ||d      }|}t-        t/        |            D ]w  }d }||   ||   z  }|dkD  rt1        ||fd	di|}|t3        ||   ||   ||   |f|||||   d
|gz  }|}| xj                  t5        ||dz
  d|z  z  d|       gz  c_        y t7        | | _        t        j:                  |d   |d   z  fddi|| _        |x| _        x| _         | _!        t        j&                  |      | _"        |	dkD  r!t        jF                  | jB                  |	fi |nt        jH                         | _%        d | _&        |rI|	dkD  r+t        jF                  | jB                  | j                  fi |nt        jH                         | _&        d| _'        tQ        | j"                  d       tQ        | j$                  d       | jS                  | jT                         y )Nr7   )tokenr   rM   r   )pT)	stagewiserc   )r1   r2   r3   r4   transformers.)num_chs	reductionr!   r;   r:   Fg{Gz?)std)+r@   rA   r   r/   r   r   r   
num_tokensfeature_inforr   patch_embedr   	Parameterr)   randnr{   r|   	pos_embedrn   Dropoutpos_dropr   rE   lenr`   r,   dictr   transformersrF   rC   num_featureshead_hidden_sizerI   	head_droprk   rB   head	head_distdistilled_trainingr   apply_init_weights)r    ru   rv   rc   r   r   r.   r/   r0   r   r   r   r   r   r   r   r   r   r8   r9   rH   rI   r   dprprev_dimrJ   r1   rK   s                              r"   rA   z!PoolingVisionTransformer.__init__   s   , 	/j((("
aL58+	& &(!a(9h
TZa^`aekk!Y@P@P@W@WY]YiYiYoYo&vsu&vwekk!T__i&VSU&VW

^4'Ns5z" 	uAD!!uQx/I1u  	 [!aa	

 (("1v
 
 
 
L !H$xFQJRSUVRVCVanopnq_r"s!tt-	u0 ,\:LL2r!:KKK	ENNND1DN I.DORSOBIIdnnk@R@Y[YdYdYf	R]`aRaRYYt~~t7G7GN2NgigrgrgtDN"'dnn#.dnn#.

4%%&r$   c                     t        |t        j                        rUt        j                  j	                  |j
                  d       t        j                  j	                  |j                  d       y y )Nr   g      ?)
isinstancer   rF   init	constant_rx   weight)r    ms     r"   r   z&PoolingVisionTransformer._init_weights   sE    a&GGaffa(GGahh, 'r$   c                 
    ddhS )Nr   rn   r   r    s    r"   no_weight_decayz(PoolingVisionTransformer.no_weight_decay  s    [))r$   c                     || _         y r   )r   r    enables     r"   set_distilled_trainingz/PoolingVisionTransformer.set_distilled_training  s
    "(r$   c                     |rJ d       y )Nz$gradient checkpointing not supportedr   r   s     r"   set_grad_checkpointingz/PoolingVisionTransformer.set_grad_checkpointing
  s    AAAz6r$   r   c                 b    | j                   | j                  | j                   fS | j                  S r   )r   r   r   s    r"   get_classifierz'PoolingVisionTransformer.get_classifier  s)    >>%99dnn,,99r$   c                 6   || _         ||| _        t        | j                  d      r | j                  j                  j
                  nd }t        | j                  d      r | j                  j                  j                  nd }|dkD  r#t        j                  | j                  |||      nt        j                         | _        | j                  L|dkD  r-t        j                  | j                  | j                   ||      nt        j                         | _        y y )Nr   r   r7   )r   r   hasattrr   r   r8   r9   r   rk   rI   rB   r   )r    r   r   r8   r9   s        r"   reset_classifierz)PoolingVisionTransformer.reset_classifier  s    &"*D,3DIIx,H!!((d*1$))X*F		  &&DZehiZiBIIdnnk&PUVoqozozo|	>>%hsvwhwRYYt~~t7G7GPV^cd}  ~I  ~I  ~KDN &r$   r   indicesrC   
stop_early
output_fmtintermediates_onlyc                 h   |dv sJ d       g }t        t        | j                        |      \  }}	| j                  |      }| j	                  || j
                  z         }| j                  j                  |j                  d   dd      }
t        | j                        dz
  }t        j                  j                         s|s| j                  }n| j                  d|	dz    }t        |      D ](  \  }} |||
f      \  }}
||v s|j                  |       * |r|S |k(  r| j                  |
      }
|
|fS )a   Forward features that returns intermediates.

        Args:
            x: Input image tensor
            indices: Take last n blocks if int, all if None, select matching indices if sequence
            norm: Apply norm layer to compatible intermediates
            stop_early: Stop iterating over blocks when last desired intermediate hit
            output_fmt: Shape of intermediate feature outputs
            intermediates_only: Only return intermediate features
        Returns:

        )NCHWzOutput shape must be NCHW.r   r   r   N)r   r   r   r   r   r   rn   expandrN   r)   jitis_scripting	enumerateappendrC   )r    r   r   rC   r   r   r   intermediatestake_indices	max_indexrS   last_idxstagesfeat_idxstages                  r"   forward_intermediatesz.PoolingVisionTransformer.forward_intermediates  s6   * Y&D(DD&"6s4;L;L7Mw"Wi QMM!dnn,-^^**1771:r2>
t(()A-99!!#:&&F&&~	A6F(0 	(OHe!1j/2MAz<'$$Q'	(
   x:.J=((r$   
prune_norm
prune_headc                     t        t        | j                        |      \  }}| j                  d|dz    | _        |rt        j                         | _        |r| j                  dd       |S )z@ Prune layers not required for specified intermediates.
        Nr   r    )r   r   r   r   rB   rC   r   )r    r   r   r   r   r   s         r"   prune_intermediate_layersz2PoolingVisionTransformer.prune_intermediate_layersO  sb     #7s4;L;L7Mw"Wi --ny1}=DI!!!R(r$   c                    | j                  |      }| j                  || j                  z         }| j                  j	                  |j
                  d   dd      }| j                  ||f      \  }}| j                  |      }|S )Nr   r   )r   r   r   rn   r   rN   r   rC   )r    r   rS   s      r"   forward_featuresz)PoolingVisionTransformer.forward_features_  su    QMM!dnn,-^^**1771:r2>
))1j/::YYz*
r$   
pre_logitsc                    | j                   | j                  dk(  sJ |d d df   |d d df   }}| j                  |      }| j                  |      }|s"| j                  |      }| j                  |      }| j                  r.| j
                  r"t        j                  j                         s||fS ||z   dz  S | j                  dk(  r	|d d df   }| j                  |      }|s| j                  |      }|S )Nr   r   r   rM   )	r   r   r   r   r   trainingr)   r   r   )r    r   r   x_dists       r"   forward_headz%PoolingVisionTransformer.forward_headg  s    >>%##w...!Q$1a4vAq!A^^F+FIIaL/&&4==AWAWAY&y  F
a''7*adGq!AIIaLHr$   c                 J    | j                  |      }| j                  |      }|S r   )r   r   r   s     r"   r#   z PoolingVisionTransformer.forward~  s'    !!!$a r$   )r   r   r   overlap0   r   r   rM         rM   r   r   r        r   FrY   rY   rY   rY   rY   NN)Tr   )NFFr   F)r   FTF) r%   r&   r'   r(   rZ   rp   r   r[   boolrA   r   r)   r   ignorer   r   r   r   r\   r   r   r   r*   r   r   r   r   r   r   r   r#   r]   r^   s   @r"   r   r      s      &'3#,#, #&#!$&$&$&$&)R'R' R' 	R'
 R'  }R' C=R' C=R' R' R' R' R' R' R' "R'  "!R'" "#R'$ "%R'h-
 YY* * YY) ) YYB B		 KC Khsm K 8<$$',/)||/) eCcN34/) 	/)
 /) /) !%/) 
tELL!5tELL7I)I#JJ	K/)f ./$#	3S	>*  	 $ 5<< .r$   c                     i }t        j                  d      }| j                         D ]  \  }}|j                  d |      }|||<    |S )z preprocess checkpoints zpools\.(\d)\.c                 D    dt        | j                  d            dz    dS )Nr   r   z.pool.)rZ   group)exps    r"   <lambda>z&checkpoint_filter_fn.<locals>.<lambda>  s"    }S15F5J4K6%R r$   )recompileitemssub)
state_dictmodelout_dictp_blockskvs         r"   checkpoint_filter_fnr     sV    Hzz*+H  " 1
 LLRTUV Or$   c                     t        t        d            }|j                  d|      }t        t        | |ft
        t        d|      d|}|S )Nr   out_indiceshook)feature_clsr   )pretrained_filter_fnfeature_cfg)tuplerE   popr   r   r   r   )variant
pretrainedkwargsdefault_out_indicesr   r   s         r"   _create_pitr	    sY    a/**],?@K   2VE E Lr$   c                 4    | ddd dddt         t        dddd	|S )
Nr   )r   r   r   g?bicubicTzpatch_embed.convr   z
apache-2.0)urlr   
input_size	pool_sizecrop_pctinterpolationfixed_input_sizemeanr   
first_conv
classifierlicenser   )r  r  s     r"   _cfgr    s5    =t%.B(  r$   ztimm/)	hf_hub_id)r   r   )r  r  )zpit_ti_224.in1kzpit_xs_224.in1kzpit_s_224.in1kzpit_b_224.in1kzpit_ti_distilled_224.in1kzpit_xs_distilled_224.in1kzpit_s_distilled_224.in1kzpit_b_distilled_224.in1kr   c           	      ^    t        ddg dg dg dd      }t        d| fi t        |fi |S )	N      @   r  r  r   r   r   r   r   r   r   rv   rc   r   r.   r/   r0   	pit_b_224r   r	  r  r  
model_argss      r"   r   r     <    J {JM$z2LV2LMMr$   c           	      ^    t        ddg dg dg dd      }t        d| fi t        |fi |S )	Nr   r   r   r   r   r      r   r  	pit_s_224r!  r"  s      r"   r(  r(    r$  r$   c           	      ^    t        ddg dg dg dd      }t        d| fi t        |fi |S )	Nr   r   r   r   r   r   r  
pit_xs_224r!  r"  s      r"   r*  r*    <    J |ZN4
3Mf3MNNr$   c           	      ^    t        ddg dg dg dd      }t        d| fi t        |fi |S )	Nr   r       r.  r.  r   r   r   r  
pit_ti_224r!  r"  s      r"   r/  r/    r+  r$   c           	      `    t        ddg dg dg ddd      }t        d	| fi t        |fi |S )
Nr  r  r  r  r  r   Trv   rc   r   r.   r/   r0   r   pit_b_distilled_224r!  r"  s      r"   r2  r2    @    J ,jWD<Vv<VWWr$   c           	      `    t        ddg dg dg ddd      }t        d	| fi t        |fi |S )
Nr   r   r   r   r&  r   Tr1  pit_s_distilled_224r!  r"  s      r"   r5  r5    r3  r$   c           	      `    t        ddg dg dg ddd      }t        d	| fi t        |fi |S )
Nr   r   r   r   r   r   Tr1  pit_xs_distilled_224r!  r"  s      r"   r7  r7    A    J -zXT*=WPV=WXXr$   c           	      `    t        ddg dg dg ddd      }t        d	| fi t        |fi |S )
Nr   r   r-  r   r   r   Tr1  pit_ti_distilled_224r!  r"  s      r"   r:  r:     r8  r$   r   )r   )3r(   ry   r   	functoolsr   typingr   r   r   r   r   r	   r
   r)   r   	timm.datar   r   timm.layersr   r   r   _builderr   	_featuresr   	_registryr   r   vision_transformerr   __all__rD   r   r\   r,   r`   rr   r   r   r	  r  default_cfgsr   r(  r*  r/  r2  r5  r7  r:  r   r$   r"   <module>rE     s    	  D D D   A K K * + < % &
&bmm 4")) 4nbii >!BII !H^ryy ^B	 %g.g.W-W-!%("* "&("* !%(!* !%(!*& * 	N-E 	N 	N 	N-E 	N 	N 	O.F 	O 	O 	O.F 	O 	O 
X7O 
X 
X 
X7O 
X 
X 
Y8P 
Y 
Y 
Y8P 
Y 
Yr$   