khS dZddlZddlZddlmZmZmZmZmZm Z m Z m Z ddl Z ddl mZddlmcmZddlmZmZddlmZmZmZmZmZmZddlmZddlmZdd l m!Z!dd l"m#Z#m$Z$dd l%m&Z&m'Z'd gZ(ejRe*Z+d e jXde jXfdZ-d e jXde jXfdZ.d e jXde e/e/fde jXfdZ0e!de jXde e/e/fde e/e/fde jXfdZ1GddejdZ3GddejdZ4GddejdZ5GddejdZ6GddejdZ7Gd d ejdZ8dXd"ejdd#e9fd$Z:d%Z;dYd&Ze=d!d*d+d,-d?e=d!d/d01d@e=d!d/d01dAe=d!d*d+d,-dBe=d!d/d01dCe=d!d*d+d,-dDe=d!d/d01dEe=d!d*d+d,-dFe=d!d/d01Z>e'dYdGe?de8fdHZ@e'dYdGe?de8fdIZAe'dYdGe?de8fdJZBe'dYdGe?de8fdKZCe'dYdGe?de8fdLZDe'dYdGe?de8fdMZEe'dYdGe?de8fdNZFe'dYdGe?de8fdOZGe'dYdGe?de8fdPZHe'dYdGe?de8fdQZIe'dYdGe?de8fdRZJe'dYdGe?de8fdSZKe'dYdGe?de8fdTZLe'dYdGe?de8fdUZMe'dYdGe?de8fdVZNe'dYdGe?de8fdWZOy)Za Swin Transformer V2 A PyTorch impl of : `Swin Transformer V2: Scaling Up Capacity and Resolution` - https://arxiv.org/pdf/2111.09883 Code adapted from https://github.com/ChristophReich1996/Swin-Transformer-V2, original copyright/license info below This implementation is experimental and subject to change in manners that will break weight compat: * Size of the pos embed MLP are not spelled out in paper in terms of dim, fixed for all models? vary with num_heads? * currently dim is fixed, I feel it may make sense to scale with num_heads (dim per head) * The specifics of the memory saving 'sequential attention' are not detailed, Christoph Reich has an impl at GitHub link above. It needs further investigation as throughput vs mem tradeoff doesn't appear beneficial. * num_heads per stage is not detailed for Huge and Giant model variants * 'Giant' is 3B params in paper but ~2.6B here despite matching paper dim + block counts * experiments are ongoing wrt to 'main branch' norm layer use and weight init scheme Noteworthy additions over official Swin v1: * MLP relative position embedding is looking promising and adapts to different image/window sizes * This impl has been designed to allow easy change of image size with matching window size changes * Non-square image size and window size are supported Modifications and additions for timm hacked together by / Copyright 2022, Ross Wightman N)AnyCallableDictListOptionalTupleTypeUnionIMAGENET_DEFAULT_MEANIMAGENET_DEFAULT_STD)DropPathMlpClassifierHead to_2tuple_assertndgrid)build_model_with_cfg)feature_take_indices)register_notrace_function) named_apply checkpoint)generate_default_cfgsregister_modelSwinTransformerV2Crxreturnc*|jddddS)z>Permutes a tensor from the shape (B, C, H, W) to (B, H, W, C).rrpermuters ^/var/www/teggl/fontify/venv/lib/python3.12/site-packages/timm/models/swin_transformer_v2_cr.py bchw_to_bhwcr&3 99Q1a  c*|jddddS)z>Permutes a tensor from the shape (B, H, W, C) to (B, C, H, W).rr!rr r"r$s r% bhwc_to_bchwr*8r'r( window_sizec|j\}}}}|j|||dz|d||dz|d|}|jddddddjjd|d|d|}|S)zPartition into non-overlapping windows. Args: x: Input tensor of shape (B, H, W, C). window_size: Window size (height, width). Returns: Windows tensor of shape (num_windows*B, window_size[0], window_size[1], C). rrr!r shapeviewr# contiguous)rr+BHWCwindowss r%window_partitionr9=sJAq!Q q!{1~%{1~qKN7JKXYN\]^Aii1aAq)446;;B AP[\]P^`abG Nr(r8img_sizec|\}}|jd}|jd||dz||dz|d|d|}|jddddddjjd|||}|S)a1Merge windows back to feature map. Args: windows: Windows tensor of shape (num_windows * B, window_size[0], window_size[1], C). window_size: Window size (height, width). img_size: Image size (height, width). Returns: Feature map tensor of shape (B, H, W, C). r/rrr!r r-r.r0)r8r+r:r5r6r7rs r%window_reverser<Ms DAq bA Rk!n,a;q>.A;q>S^_`SacdeA !Q1a#..055b!QBA Hr(ceZdZdZ ddededeeefdedededed d ffd Zdd Z deeefd d fd Z d e jfdZ dde jdee jd e jfdZxZS)WindowMultiHeadAttentionaThis class implements window-based Multi-Head-Attention with log-spaced continuous position bias. Args: dim (int): Number of input features window_size (int): Window size num_heads (int): Number of attention heads drop_attn (float): Dropout rate of attention map drop_proj (float): Dropout rate after projection meta_hidden_dim (int): Number of hidden features in the two layer MLP meta network sequential_attn (bool): If true sequential self-attention is performed dim num_headsr+ drop_attn drop_projmeta_hidden_dimsequential_attnrNc`tt| ||zdk(sJd||_t ||_||_||_tj||dzd|_ tj||_ tj||d|_ tj||_td||tj d|_tj$t'j(d t'j*|z|_|j/y) Nrz`The number of input features (in_features) are not divisible by the number of heads (num_heads).r!T in_features out_featuresbiasr )g?)hidden_featuresrH act_layerdrop )superr>__init__rGrr+r@rDnnLinearqkvDropout attn_dropproj proj_droprReLUmeta_mlp Parametertorchlogones logit_scale"_make_pair_wise_relative_positions) selfr?r@r+rArBrCrD __class__s r%rPz!WindowMultiHeadAttention.__init__ms &68Y!# o n o# #,5k,B'%49937NI.II#CdK I. +"gg   << "uzz)7L2L(MN //1r(c L|jj}tjt tj |j d|tj |j d|djd}|dddddf|dddddfz }|jdddjddj}tj|tjd|jzz}|jd |d y) zMInitialize the pair-wise relative positions to compute the positional biases.r)devicerr?Nr r/?relative_coordinates_logF persistent)r^rcr[stackraranger+flattenr#reshapefloatsignr\absregister_buffer)r`rc coordinatesrelative_coordinatesrfs r%r_z;WindowMultiHeadAttention._make_pair_wise_relative_positionss!!((kk& LL))!,V < LL))!,V <# '!*  +1a:6QaZ9PP3;;Aq!DLLRQRSYY[#(::.B#Ceii &**, ,G.$.  79Q^cdr(cht|}||jk7r||_|jyy)zUpdate window size and regenerate relative position coordinates. Args: window_size: New window size. N)rr+r_)r`r+s r%set_window_sizez(WindowMultiHeadAttention.set_window_sizes5  , $** **D   3 3 5 +r(c|jd|jdz}|j|j}|jddj |j ||}|j d}|S)zCompute the relative positional encodings. Returns: Relative positional encodings of shape (1, num_heads, window_size**2, window_size**2). rr)r+rYrf transposerlr@ unsqueeze)r` window_arearelative_position_biass r%_relative_positional_encodingsz7WindowMultiHeadAttention._relative_positional_encodingss{ &&q)D,<,r>`s "" % 2  2 238_ 2  2  2 2 2  2D e 65c? 6t 6 & &!!Xell-C!u||!r(r>c eZdZdZddddddddddej f deded eeefd eeefd eeefd ed ede de e de de de dedede ejffd Z deeefdeeeefeeefffdZd!de ej de ej fdZdZd eeefd eeefddfdZdZdej dej fd ZxZS)"SwinTransformerV2CrBlocka5This class implements the Swin transformer block. Args: dim (int): Number of input channels num_heads (int): Number of attention heads to be utilized feat_size (Tuple[int, int]): Input resolution window_size (Tuple[int, int]): Window size to be utilized shift_size (int): Shifting size to be used mlp_ratio (int): Ratio of the hidden dimension in the FFN to the input channels proj_drop (float): Dropout in input mapping drop_attn (float): Dropout rate of attention map drop_path (float): Dropout in main path extra_norm (bool): Insert extra norm on 'main' branch if True sequential_attn (bool): If true sequential self-attention is performed norm_layer (Type[nn.Module]): Type of normalization layer to be utilized )rrF@rrJr?r@ feat_sizer+ shift_sizealways_partition dynamic_mask mlp_ratio init_valuesrWrA drop_path extra_normrD norm_layerc6tt| ||_||_t ||_||_||_|j|\|_ |_ |jd|jdz|_ | |_ t|||j| | ||_|||_| dkDr t#| nt%j&|_t+|t-||z| ||_|||_| dkDr t#| nt%j&|_| r||nt%j&|_|j7d|jrdn|j9d |j;y) Nrr)r?r@r+rArBrDrJ) drop_prob)rGrKrMrH attn_maskFrg)rOrrPr?rrtarget_shift_sizerr_calc_window_shiftr+rrxrr>rnorm1rrQIdentity drop_path1rrmlpnorm2 drop_path2norm3rp get_attn_mask init_weights)r`r?r@rr+rrrrrrWrArrrDrras r%rPz!SwinTransformerV2CrBlock.__init__sv$ &68*32;J2G 0(,0,C,CK,P)$/++A.1A1A!1DD,7-((+    _ ;Ds?(Y7PRP[P[P]i0    _ ;Ds?(Y7PRP[P[P])3Z_    %%D4+=+=+?  r(target_window_sizerct|}|j}t|r|ddz|ddzf}|jr||fSt |j |Dcgc]\}}||kr|n|}}}t |j ||Dcgc]\}}}||krdn|}}}}t |t |fScc}}wcc}}}w)Nrr r)rranyrziprtuple)r`rrfwr+srs r%rz+SwinTransformerV2CrBlock._calc_window_shift(s''9: 22  !!3A!6!!;=OPQ=RVW=W X   %'88 869$..J\6]^daAFq)^ ^8;DNNKYj8kllWQ116aq(l l[!5#444_ls B; CNrc t|jr|&tjdg|jd}nJtjd|j d|j ddf|j |j}d}d|jd f|jd |jd f|jd dffD]l}d|jd f|jd |jd f|jd dffD]$}||dd|d|d|d|dddf<|dz }&nt||j}|jd|j}|jd|jdz }|j|dk7tdj|dk(td}|Sd}|S) z7Method generates the attention mask used in shift case.Nrr )dtypercrr/gYrJ)rrr[zerosrr1rrcr+r9r2rxrw masked_fillrm)r`rimg_maskcnthr mask_windowsrs r%rz&SwinTransformerV2CrBlock.get_attn_mask9s t y ;;'>DNN'>A'>? ;;1771:qwwqz1'EQWW]^]e]efC))!,,-&&q))DOOA,>+>?ooa(($/  T--a001**1--0B/BC//!,,d3A '>>$BRBRSTBUU!!!$q4+;+;A+>'>>$BRBRSTBUU HH   # #A1a5'A Bww 2r1%Q(8(89 NN2t'7'7':T=M=Ma=P'PRST  4 /**1-IIyyy; $((T-=-=a-@$BRBRSTBUWXY <)9)9B8 D a!RaRlO & & (  1b"XF;Ar(c`||j|j|j|z}|j\}}}}|j |d|}||j |j |j|z}|j|}|j ||||}|S)zForward pass of Swin Transformer V2 block. Args: x: Input tensor of shape [B, C, H, W]. Returns: Output tensor of shape [B, C, H, W]. r/) rrrr1rlrrrr)r`rr4r5r6r7s r%rz SwinTransformerV2CrBlock.forwards  4+D+DQ+G HI IWW 1a IIaQ   488A; 78 8 JJqM IIaAq !r(r)rrrrrQ LayerNormrrrrmrr ModulerPrr[rrrrrrrrs@r%rrs.+1%*!&"+,"""$$)*,,,!;;;S#X ; sCx ; c3h ;#;;;"%;;;;;"; RYY!;z5 %c3h5 uS#Xc3h/ 05"x 5%,,AW<C  c3h eCQTHo Z^ $(T%,,r(rceZdZdZej fdedeejddffd Z de jde jfdZ xZ S) PatchMergingz|Patch merging layer. This class implements the patch merging as a strided convolution with a normalization before. r?rrNctt| |d|z|_t j d|zd|zd|_y)zInitialize patch merging layer. Args: dim: Number of input channels. norm_layer: Type of normalization layer to be utilized. r-r FrFN)rOrrPnormrQrR reduction)r`r?rras r%rPzPatchMerging.__init__s? lD*,q3w' q3wQWSXYr(rch|j\}}}}ddd|dzd|dzf}tjj||}|j\}}}}|j ||dzd|dzd|j ddddddj d}|j|}|j|}|S)zForward pass of patch merging. Args: x: Input tensor of shape [B, C, H, W]. Returns: Output tensor of shape [B, 2 * C, H // 2, W // 2]. rr rr!r-r.) r1rQrrrlr#rkrr)r`rr4r5r6r7 pad_valuesrs r%rzPatchMerging.forwardsWW 1aAq1uaQ/ MM  a ,WW 1a IIaaAFAq 1 9 9!Q1a K S STU V IIaL NN1 r()rrrrrQrrr rrPr[rrrrs@r%rrsU @B|| ZC ZT"))_ ZPT Z%,,r(rceZdZdZ ddeeeeeffdeeeeeffdededeede d dffd Z deeefd dfd Z d e jd e jfd ZxZS) PatchEmbedz2D Image to Patch Embedding.Nr: patch_sizein_chans embed_dimrstrict_img_sizerct|t|}t|}||_||_|d|dz|d|dzf|_|j d|j dz|_||_tj|||||_ |r|||_ ytj|_ y)a@Initialize patch embedding. Args: img_size: Input image size. patch_size: Patch size. in_chans: Number of input channels. embed_dim: Embedding dimension. norm_layer: Normalization layer. strict_img_size: Enforce strict image size. rr) kernel_sizestrideN) rOrPrr:r grid_size num_patchesrrQConv2drVrr)r`r:rrrrrras r%rPzPatchEmbed.__init__s& X&z*   $"1+A6 zRS}8TU>>!,t~~a/@@.IIh zR\] -7Jy) R[[] r(ct|}||jk7r[||_|d|jdz|d|jdzf|_|jd|jdz|_yy)zVUpdate input image size. Args: img_size: New image size. rrN)rr:rrr)r`r:s r%rzPatchEmbed.set_input_sizesy X& t}} $$DM&qkT__Q-??!PTP_P_`aPbAbcDN#~~a04>>!3DDD  %r(rc|j\}}}}|jr`t||jdk(d|d|jddt||jdk(d|d|jdd|j |}|j |j ddddj dddd}|S) zForward pass of patch embedding. Args: x: Input tensor of shape [B, C, H, W]. Returns: Output tensor of shape [B, C', H', W']. rzInput image height (z) doesn't match model (z).rzInput image width (r r!)r1rrr:rVrr#)r`rr4r7r5r6s r%rzPatchEmbed.forwardsWW 1a    Aq))-A!D[\`\i\ijk\l[mmo+p q Aq))-@CZ[_[h[hij[kZlln+o p IIaL IIaii1a+ , 4 4Q1a @r()r!iNT)rrrrr rrrrrrPrr[rrrrs@r%rrs&5868 -1$(KCsCx01Kc5c?23K K  K !* K"K K> EuS#X E4 E%,,r(rc#8eZdZdZdddddddej dddf dededed ed eeefd eeefd ed ede de e de de de e e e fde ejdededef"fd Z dd eeefd ed e efdZdej$dej$fdZxZS)SwinTransformerV2CrStagea<This class implements a stage of the Swin transformer including multiple layers. Args: embed_dim (int): Number of input channels depth (int): Depth of the stage (number of layers) downscale (bool): If true input is downsampled (see Fig. 3 or V1 paper) feat_size (Tuple[int, int]): input feature map size (H, W) num_heads (int): Number of attention heads to be utilized window_size (int): Window size to be utilized mlp_ratio (int): Ratio of the hidden dimension in the FFN to the input channels proj_drop (float): Dropout in input mapping drop_attn (float): Dropout rate of attention map drop_path (float): Dropout in main path norm_layer (Type[nn.Module]): Type of normalization layer to be utilized. Default: nn.LayerNorm extra_norm_period (int): Insert extra norm layer on main branch every N (period) blocks extra_norm_stage (bool): End each stage with an extra norm layer in main branch sequential_attn (bool): If true sequential self-attention is performed FrrJrrdepth downscaler@rr+rrrrrWrArrextra_norm_periodextra_norm_stagerDc*tt| ||_d|_|r|ddz|ddzfn||_|rt |||_|dz}ntj|_fd}tjtDcgc]e}t|||j |||t|Dcgc]}|dzdk(rdn|dzc}| | | | t| tr| |n| ||||gc}}|_ycc}wcc}}w)NFrr r)rc6|dz}r |zdk(ryr|k(SdS)NrrTF)indexirrrs r% _extra_normz6SwinTransformerV2CrStage.__init__.._extra_normEs0 A Q):%:a%?!11: O*Pcl *9LDO!A I kkmDO = mm$u%&'$# %#..'!1) {![!q(8!qAv"E![\#'##.8D.I)E*y&u- /% &'  "\&'s#D 8D $1D D c|jr|ddz|ddzfn||_|jD]}|j|j|!y)z Updates the resolution to utilize and the window size and so the pair-wise relative positions. Args: window_size (int): New window size feat_size (Tuple[int, int]): New input resolution rr r)rr+N)rrrr)r`rr+rblocks r%rz'SwinTransformerV2CrStage.set_input_size`s]DH>>)A,!+Yq\Q->?W`[[ E  ..' !  r(rrct|}|j|}|jD]A}|jr+tj j s t||}:||}Ct|}|S)zForward pass. Args: x (torch.Tensor): Input tensor of the shape [B, C, H, W] or [B, L, C] Returns: output (torch.Tensor): Output tensor of the shape [B, 2 * C, H // 2, W // 2] ) r&rrrr[jit is_scriptingrr*)r`rr s r%rz SwinTransformerV2CrStage.forwardssk O OOA [[ E&&uyy/E/E/Gua(!H   Or(r)rrrrrQrrrrrmrr rr rrPrr[rrrrs@r%rrsT6"'"'*/2&(ll!"!& %%8 8 8  8  8 c? 8 38_8 8 8 8 e_8 8 8 ee+,8 O8 !8 "#8 $%8 |04 S#X'tn &%,,r(rc1eZdZdZdddddddd d d d d dddddej dddddfdeeefdedeedede de dedededeedfdeedfde dee d e d!e d"e d#e d$e ejd%ed&e d'e d(e d)ed*df0fd+ Z d?deeeefdeeeefdedee d*df d,Zej$j&d@d-Zej$j&dAd.Zej$j'd*ejfd/ZdBded(ee d*dfd0Z dCd1ej0d2eeeeefd3e d4e d5e d6e d*eeej0eej0eej0fffd7Z dDd2eeeefd8e d9e fd:Zd1ej0d*ej0fd;Zd@dZxZ S)Era Swin Transformer V2 A PyTorch impl of : `Swin Transformer V2: Scaling Up Capacity and Resolution` - https://arxiv.org/pdf/2111.09883 Args: img_size: Input resolution. window_size: Window size. If None, grid_size // window_div window_ratio: Window size to patch grid ratio. patch_size: Patch size. in_chans: Number of input channels. depths: Depth of the stage (number of layers). num_heads: Number of attention heads to be utilized. embed_dim: Patch embedding dimension. num_classes: Number of output classes. mlp_ratio: Ratio of the hidden dimension in the FFN to the input channels. drop_rate: Dropout rate. proj_drop_rate: Projection dropout rate. attn_drop_rate: Dropout rate of attention map. drop_path_rate: Stochastic depth rate. norm_layer: Type of normalization layer to be utilized. extra_norm_period: Insert extra norm layer on main branch every N (period) blocks in stage extra_norm_stage: End each stage with an extra norm layer in main branch sequential_attn: If true sequential self-attention is performed. )rrr-NFTr!`r r r r!r rrJravgskipr:rr+ window_ratiorrr num_classesrdepths.r@rr drop_rateproj_drop_rateattn_drop_ratedrop_path_raterrrrD global_poolkwargsrc ftt| t|}||_||_||_t| dt| dz zzx|_ |_ g|_ t|||| |||_ |jj}|"t|Dcgc]}||z c}|_nt||_t#j$d|t'| j)| Dcgc]}|j+}}g}| }d}t-t/| | D]\} \}!} |t1did|d|!d| dk7d|d|z|d|zfd | d |j d |d | d | d| d|d|d|| d|d|xs| dzt| k(d|d|gz }| dk7r |dz}|dz}|xjt3|d|zd| gz c_ t5j6||_t;|j||||_|dk7rt?t@|yycc}wcc}w)Nr r)r:rrrrrrrrrrr@r+rrrrrWrArrrrDrr-zstages.)num_chsrmodule) pool_typerrr)!rOrrPrrrr:rlen num_featureshead_hidden_size feature_infor patch_embedrrr+r[linspacesumsplittolist enumeraterrdictrQrstagesrheadrr)#r`r:rr+rrrrrrrr@rrrrrrrrrrDr  weight_initr!rrrdprr1in_dimin_scale stage_idxrras# r%rPzSwinTransformerV2Cr.__init__s6 !413X& +))1 47 A#f+XY/DZ8Z4[[D1%!!+  $$..  $%KAa<&7%KLD (5D #(>>!^S[#Q#W#WX^#_`aqxxz``-6s697M-N n )I)y / $q.%Q<83Yq\X5MN  $  !,, "2"10$())i.#4"2!Si!mF 5S !0!"&# F&A~! A    $vXX_`i_jVk"l!m m / n0mmV, "    !    &  d + !W&Las  H)5H.cX|2|jj||jj}||tDcgc]}||z c}}t |j D]9\}}dt |dz dz} |jd| z|d| zf||;ycc}w)aUpdates the image resolution, window size and so the pair-wise relative positions. Args: img_size (Optional[Tuple[int, int]]): New input resolution, if None current resolution is used window_size (Optional[int]): New window size, if None based on new_img_size // window_div window_ratio (int): divisor for calculating window size from patch grid size always_partition: always partition / shift windows even if feat size is < window N)r:r rr)rr+r)r*rrrr/r1r~) r`r:r+rrrrrstage stage_scales r%rz"SwinTransformerV2Cr.set_input_sizes      + +X + >((22I  <#;I Fql!2 FGK%dkk2 LE5s519a00K  $Q<;6 ! 8ST'!1 !  !Gs B'c2td|rdSddgS)Nz ^patch_embedz^stages\.(\d+))z^stages\.(\d+).downsample)r)z^stages\.(\d+)\.\w+\.(\d+)N)stemr)r0)r`coarses r% group_matcherz!SwinTransformerV2Cr.group_matchers/ (.$  555  r(c4|jD] }||_ yr)r1r)r`enablers r%set_grad_checkpointingz*SwinTransformerV2Cr.set_grad_checkpointing#s *A#)A  *r(c.|jjS)zMethod returns the classification head of the model. Returns: head (nn.Module): Current classification head )r2fcrs r%get_classifierz"SwinTransformerV2Cr.get_classifier(s yy||r(cJ||_|jj||y)zMethod results the classification head Args: num_classes (int): Number of classes to be predicted global_pool (str): Unused N)rr2reset)r`rr s r%reset_classifierz$SwinTransformerV2Cr.reset_classifier0s'  [1r(rindicesr stop_early output_fmtintermediates_onlycr|dvsJdg}tt|j|\}} |j|}tj j s|s |j} n|jd| dz} t| D]#\} } | |}| |vs|j|%|r|S||fS)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: )NCHWzOutput shape must be NCHW.Nr) rr&r1r*r[r r r/append) r`rrHrrIrJrK intermediates take_indices max_indexr1rr9s r%forward_intermediatesz)SwinTransformerV2Cr.forward_intermediates:s*Y&D(DD& "6s4;;7G"Q i   Q  99 ! ! #:[[F[[)a-0F!&) (HAuaAL $$Q' (  -r( prune_norm prune_headctt|j|\}}|jd|dz|_|r|jdd|S)z@ Prune layers not required for specified intermediates. Nrr)rr&r1rG)r`rHrSrTrPrQs r%prune_intermediate_layersz-SwinTransformerV2Cr.prune_intermediate_layersdsM#7s4;;7G"Q ikk.9q=1   ! !!R (r(cJ|j|}|j|}|Sr)r*r1r`rs r%forward_featuresz$SwinTransformerV2Cr.forward_featuresrs$   Q  KKNr( pre_logitscN|r|j|dS|j|S)NT)r[)r2)r`rr[s r% forward_headz SwinTransformerV2Cr.forward_headws$0:tyyty,L ! Lr(cJ|j|}|j|}|Sr)rZr]rYs r%rzSwinTransformerV2Cr.forwardzs'  ! !! $   a r()NNrNF)Tr)NFFrMF)rFT)!rrrrrQrrrrrrmr rstrrrPrr[r ignorer>rArDrGrr rrRrWrZr]rrrs@r%rrs6%/%)!& $".%3') # # #&(ll!"!& % 1Y,S/Y,Y,c] Y,  Y,  Y,Y,Y,Y,Y,c3hY,c?Y,Y,e_Y,Y, !Y,"#Y,$%Y,&O'Y,()Y,*+Y,,-Y,./Y,23Y,4 5Y,z3759 !/3 uS#X/"%S/2  'tn   < YY   YY** YY 2C2hsm2W[28<$$',( ||( eCcN34(  (  (  ( !%(  tELL!5tELL7I)I#JJ K( X./$# 3S >*   %,,5<< M$M%,,r(rVr$namect|tjrd|vrtjdt |j jddz|j jdzz }tjj|j | |nWd|vr*tjj|j n)tjj|j |j*tjj|jyyt|dr|jyy)NrSg@rr!rr2r)rrQrRrsqrtrmrr1runiform_zeros_xavier_uniform_rIhasattrr)r$rbvals r%rrs&"))$ D=))Bv}}':':1'='BV]]EXEXYZE['[!\\]C GG  V]]SD# 6 t^ GGNN6== ) GG # #FMM 2 ;; " GGNN6;; ' #  ( )r(c|jd|}|jd|}d|vr|Si}|jD]J\}}d|vr*tjd|z }|j dd}|j dd}|||<L|S) zJ convert patch embedding weight from manual patchify + linear proj to convmodel state_dictzhead.fc.weighttaurr^zhead.zhead.fc.)getitemsr[r\replace)rlrkout_dictkvs r%checkpoint_filter_fnrts4J j9J:%H  "1 A: !a% A %/A IIgz *  Or(c tdt|jddD}|jd|}t t ||ft td|d|}|S)Nc3&K|] \}}| ywrr).0rrs r% z1_create_swin_transformer_v2_cr..s\da\sr)rrrr out_indicesT)flatten_sequentialry)pretrained_filter_fn feature_cfg)rr/rnpoprrrtr0)variant pretrainedr!default_out_indicesryrks r%_create_swin_transformer_v2_crrsh\i 8\8Z.[\\**],?@K Wj 1DkJ    E Lr(urlc 2|ddddddttddd |S) zCreate a default configuration dictionary. Args: url: Model weights URL. **kwargs: Additional configuration parameters. Returns: Configuration dictionary. rr!rr)r?bicubicTzpatch_embed.projzhead.fc) rr input_size pool_sizecrop_pct interpolationfixed_input_sizemeanstd first_conv classifierr )rr!s r%_cfgrs9#" %#(    r(zswinv2_cr_tiny_384.untrained)r!rrre)rr)rrrrzswinv2_cr_tiny_224.untrainedrr)rrrzswinv2_cr_tiny_ns_224.sw_in1kztimm/z{https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights-swinv2/swin_v2_cr_tiny_ns_224-ba8166c6.pth) hf_hub_idrrrzswinv2_cr_small_384.untrainedzswinv2_cr_small_224.sw_in1kzyhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights-swinv2/swin_v2_cr_small_224-0813c165.pthzswinv2_cr_small_ns_224.sw_in1kzhttps://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights-swinv2/swin_v2_cr_small_ns_224_iv-2ce90f8e.pthz swinv2_cr_small_ns_256.untrained)r!r)rrzswinv2_cr_base_384.untrainedzswinv2_cr_base_224.untrainedzswinv2_cr_base_ns_224.untrainedzswinv2_cr_large_384.untrainedzswinv2_cr_large_224.untrainedzswinv2_cr_huge_384.untrainedzswinv2_cr_huge_224.untrainedzswinv2_cr_giant_384.untrainedzswinv2_cr_giant_224.untrainedrc Ltddd}tdd|it|fi|S)z,Swin-T V2 CR @ 384x384, trained ImageNet-1k.rrrrrr@r)swinv2_cr_tiny_384r0rrr! model_argss r%rr9 J * t: tY]^hYslrYs ttr(c Ltddd}tdd|it|fi|S)z,Swin-T V2 CR @ 224x224, trained ImageNet-1k.rrrrr)swinv2_cr_tiny_224rrs r%rrrr(c Ntdddd}tdd|it|fi|S)zSwin-T V2 CR @ 224x224, trained ImageNet-1k w/ extra stage norms. ** Experimental, may make default if results are improved. ** rrrTrrr@rr)swinv2_cr_tiny_ns_224rrs r%rr s<   J * wj w\`ak\vou\v wwr(c Ltddd}tdd|it|fi|S)z,Swin-S V2 CR @ 384x384, trained ImageNet-1k.rr r r rrr)swinv2_cr_small_384rrs r%rr9 J * uJ uZ^_iZtmsZt uur(c Ltddd}tdd|it|fi|S),Swin-S V2 CR @ 224x224, trained ImageNet-1k.rrrrr)swinv2_cr_small_224rrs r%rr#rr(c Ntdddd}tdd|it|fi|S)rrrrTrr)swinv2_cr_small_ns_224rrs r%rr.<  J * xz x]abl]wpv]w xxr(c Ntdddd}tdd|it|fi|S)z,Swin-S V2 CR @ 256x256, trained ImageNet-1k.rrrTrr)swinv2_cr_small_ns_256rrs r%rr:rr(c Ltddd}tdd|it|fi|S)z,Swin-B V2 CR @ 384x384, trained ImageNet-1k.rr-rr rr)swinv2_cr_base_384rrs r%rrF9 J * t: tY]^hYslrYs ttr(c Ltddd}tdd|it|fi|S),Swin-B V2 CR @ 224x224, trained ImageNet-1k.rrrrr)swinv2_cr_base_224rrs r%rrQrr(c Ntdddd}tdd|it|fi|S)rrrrTrr)swinv2_cr_base_ns_224rrs r%rr\s<  J * wj w\`ak\vou\v wwr(c Ltddd}tdd|it|fi|S)z,Swin-L V2 CR @ 384x384, trained ImageNet-1k.rrrr0rr)swinv2_cr_large_384rrs r%rrh9!J * uJ uZ^_iZtmsZt uur(c Ltddd}tdd|it|fi|S)z,Swin-L V2 CR @ 224x224, trained ImageNet-1k.rrrrr)swinv2_cr_large_224rrs r%rrsrr(c Ntdddd}tdd|it|fi|S)z,Swin-H V2 CR @ 384x384, trained ImageNet-1k.`r) ,Xrrrr@rr)swinv2_cr_huge_384rrs r%rr~s<" J * t: tY]^hYslrYs ttr(c Ntdddd}tdd|it|fi|S)z,Swin-H V2 CR @ 224x224, trained ImageNet-1k.rr)rrr@rrr)swinv2_cr_huge_224rrs r%rrs<! J * t: tY]^hYslrYs ttr(c Ntdddd}tdd|it|fi|S)z,Swin-G V2 CR @ 384x384, trained ImageNet-1k.r r *r rrrrrrr)swinv2_cr_giant_384rrs r%rr<# J * uJ uZ^_iZtmsZt uur(c Ntdddd}tdd|it|fi|S)z,Swin-G V2 CR @ 224x224, trained ImageNet-1k.rrrrrr)swinv2_cr_giant_224rrs r%rrrr()rVr_)Prloggingrtypingrrrrrrr r r[torch.nnrQtorch.nn.functionalrr timm.datar r timm.layersrrrrrr_builderr _featuresr _features_fxr _manipulaterr _registryrr__all__ getLoggerr_loggerrr&r*rr9r<rr>rrrrrr`rrtrr default_cfgsrrrrrrrrrrrrrrrrrrr(r%rs: JJJ AQQ*+30< ! '  H %!ELL!U\\! !ELL!U\\!  5c? u||   ELL uS#X RWX[]`X`Ra fkfrfr  $vryyvrQryyQh#299#L<<~qryyqhw"))wt#"" cT#s(^4%'&"D =3(%L'&#D =3%8'& $T J 3&0 '&$T =3(&L'&"4 H 3$0'&%d N 3'0'&&' =3&)J''&*#D =3(%L+'&.#D =3%8/'&2&t =3(83'&6$T =3(&L7'&:$T =3&8;'&>#D =3(%L?'&B#D =3%8C'&F$T =3(&LG'&J$T =3&8K'&' Tu4u>Quuu4u>Quu xd xAT x xvDv?RvvvDv?RvvytyBUyyytyBUyyu4u>Quuu4u>QuuxdxATxxvDv?RvvvDv?Rvvu4u>Quuu4u>QuuvDv?RvvvDv?Rvvr(