o *j@s ddlZddlZddlZddlZddlmZddlmmZ ddl m m Z ddl mZmZmZGdddejZddejfddZd d Zd d ZGd ddejZGdddejZGdddejZGdddejZGdddejZGdddejZdS)N)DropPath to_2tuple trunc_normal_cs4eZdZdZddejdffdd ZddZZS)Mlpz Multilayer perceptron.NcsNt|p|}|p |}t|||_||_t|||_t||_dSN) super__init__nnLinearfc1actfc2Dropoutdrop)self in_featureshidden_featuresZ out_features act_layerr __class__c/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/models/cv/vidt/backbone.pyr s z Mlp.__init__cCs6||}||}||}||}||}|Sr)r r rr)rxrrrforward!s     z Mlp.forward) __name__ __module__ __qualname____doc__r GELUr r __classcell__rrrrrsri'c Cs|d}|}|jdtjd}|jdtjd}d}||ddddddf||}||ddddddf||}tj|tj|jd} |d| d|} |dddddddf| } |dddddddf| } tj| dddddddddf| dddddddddffd d d } tj| dddddddddf| dddddddddffd d d } tj | | fd d } | S) aD Masked Sinusoidal Positional Encoding Args: x: [PATCH] tokens mask: the padding mask for [PATCH] tokens num_pos_feats: the size of channel dimension temperature: the temperature value scale: the normalization scale Returns: pos: Sinusoidal positional encodings r!)dtypegư>N)r#devicerdim) ZcumsumtorchZfloat32aranger%stacksincosflattencat) rmaskZ num_pos_featsZ temperaturescaleZnot_maskZy_embedZx_embedepsZdim_tZpos_xZpos_yposrrrmasked_sin_pos_encoding*s2&&  JJr5cCsR|j\}}}}|||||||||}|ddddddd|||}|S)z Args: x: (B, H, W, C) window_size (int): window size Returns: windows: (num_windows*B, window_size, window_size, C) rr"r)r!r&r$)shapeviewpermute contiguous)r window_sizeBHWCwindowsrrrwindow_partitionWs rAcCsbt|jd||||}||||||||d}|dddddd|||d}|S)z Args: windows: (num_windows*B, window_size, window_size, C) window_size (int): Window size H (int): Height of image W (int): Width of image Returns: x: (B, H, W, C) rr$r"r)r!r&r6)intr7r8r9r:)r@r;r=r>r<rrrrwindow_reversehs  $rCcs:eZdZdZ    d fdd Z   d dd ZZS) ReconfiguredAttentionModulea Window based multi-head self attention (W-MSA) module with relative position bias -> extended with RAM. It supports both of shifted and non-shifted window. !!!!!!!!!!! IMPORTANT !!!!!!!!!!! The original attention module in Swin is replaced with the reconfigured attention module in Section 3. All the Args are shared, so only the forward function is modified. See https://arxiv.org/pdf/2110.03921.pdf !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Args: dim (int): Number of input channels. window_size (tuple[int]): The height and width of the window. num_heads (int): Number of attention heads. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 proj_drop (float, optional): Dropout ratio of output. Default: 0.0 TNrcst||_||_||_||}|p|d|_tt d|ddd|dd||_ t |jd} t |jd} t t | | g} t| d} | dddddf| dddddf} | ddd} | dddddf|jdd7<| dddddf|jdd7<| dddddfd|jdd9<| d}|d|tj||d|d|_t||_t|||_t||_t|j d d tjdd |_dS) Ngr!rr"r$relative_position_indexr)Zbias{Gz?Zstdr')rr r(r; num_headsr2r Parameterr*zerosrelative_position_bias_tabler+r,Zmeshgridr/r9r:sumZregister_bufferr qkvr attn_dropproj proj_droprZSoftmaxsoftmax)rr(r;rIqkv_biasqk_scalerOrQZhead_dimZcoords_hZcoords_wZcoordsZcoords_flattenZrelative_coordsrErrrr sX "  (,   z$ReconfiguredAttentionModule.__init__Fc* Cs|jd|jdks J|jd}||}|sS|j\}} } } | | } ||| | }tj||gdd}||} | ddd| ddf| dd| dddf}}nh|dj\}} } } | | } |dj\}}}}||}|d|| | }|d||| }tj|||gdd}||} | ddd| ddf| dd| | |ddf| dd| |dddf}}}||| | d}t||}|jd}||||d|j| |j}| ddddd}|d|d|d}}}||j }|| d d}|j |j d|jd|jd|jd|jdd}| ddd}||d}|durY|jd}|||||j||}|dd} || }|d|j||}||}||}|| dd|||| }!||dd|j| |j}| ddddd}|d|d|d}"}#}$|r||||d|j| |j}|ddddddddddddf dddddd }%|%d||j||d}%|%d|%d}&}'tj|&|#gddtj|'|$gdd}#}$|"|j }"|"|# d d}(|dur|(|}(||(}(||(}(|(|$ dd|d| })t|!|| | }!tj|!|| | | |)gdd}||}||}|ddd| | ddf|| | | }!|dd| | dddf})|!|)fS) al Forward function. RAM module receives [Patch] and [DET] tokens and returns their calibrated ones Args: x: [PATCH] tokens det: [DET] tokens mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None -> mask for shifted window attention "additional inputs for RAM" cross_attn: whether to use cross-attention [det x patch] (for selective cross-attention) cross_attn_mask: mask for cross-attention Returns: patch_x: the calibrated [PATCH] tokens det_x: the calibrated [DET] tokens rr"r'Nr$r)r!r&r6)r;r7r8r*r0rNrAZreshaperIr9r2 transposerLrEr: unsqueezerRrOrCrPrQ)*rrdetr1 cross_attncross_attn_maskr;Zlocal_map_sizer<r=r>r?NZfull_qkvZ patch_qkvZdet_qkv_Zori_HZori_WZori_N shifted_xZcross_xZcross_patch_qkvZB_Z _patch_qkvZpatch_qZpatch_kZpatch_vZ patch_attnZrelative_position_biasZnWZtmp0Ztmp1Zpatch_xZdet_qZdet_kZdet_vZpatch_kvZ cross_patch_kZ cross_patch_vZdet_attnZdet_xrrrrs  8 R                    *z#ReconfiguredAttentionModule.forward)TNrr)NFNrrrrr rr rrrrrDzs2rDc sBeZdZdZddddddddejejf fdd Zd d ZZ S) SwinTransformerBlocka] Swin Transformer Block. Args: dim (int): Number of input channels. num_heads (int): Number of attention heads. window_size (int): Window size. shift_size (int): Shift size for SW-MSA. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. drop (float, optional): Dropout rate. Default: 0.0 attn_drop (float, optional): Attention dropout rate. Default: 0.0 drop_path (float, optional): Stochastic depth rate. Default: 0.0 act_layer (nn.Module, optional): Activation layer. Default: nn.GELU norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm r@TNrc  st||_||_||_||_||_d|jkr#|jks(JdJd| ||_t|t |j|||| |d|_ | dkrDt | nt |_| ||_t||} t|| | |d|_d|_d|_dS)Nrz shift_size must in 0-window_size)r;rIrSrTrOrQr)rrrr)rr r(rIr; shift_size mlp_rationorm1rDrattnrr ZIdentity drop_pathnorm2rBrmlpr=r>)rr(rIr;rbrcrSrTrrOrfr norm_layerZmlp_hidden_dimrrrr WsB (     zSwinTransformerBlock.__init__c Cs"|j\}}}|j|j} } || | |jksJd|} ||}|ddd| | ddf|dd| | dddf}} ||| | |}|} d}}|j| |j|j}|j| |j|j}t|dd||||f}|j\}}}}|\}}| |}|j dkrt j ||j |j fdd}|}n|}d}|r| |}| |} ||f}n| |} |}|j ||| ||d\}} |j dkrt j ||j |j fdd}n|}|dks|dkr|ddd| d| ddf}||| | |}t j|| gdd}| ||}|||||}|S) a Forward function. Args: x: Input feature, tensor size (B, H*W + DET, C). i.e., binded [PATCH, DET] tokens H, W: Spatial resolution of the input feature. mask_matrix: Attention mask for cyclic shift. "additional inputs' pos: (patch_pos, det_pos) cross_attn: whether to use cross attn [det x [det + patch]] cross_attn_mask: attention mask for cross-attention Returns: x: calibrated & binded [PATCH, DET] tokens input feature has wrong sizeNr)r"r!)Zshiftsdims)r1rXrYrZr"r')r7r=r> det_token_numrdr8r;Fpaddet_pos_linearrbr*Zrollrer:r0rfrhrg)rrZ mask_matrixr4rYrZr<Lr?r=r>ZshortcutrXZorig_xZpad_lZpad_tZpad_rZpad_br\HpWp patch_posdet_posr] attn_maskZ cross_patchrrrrsb  >     $zSwinTransformerBlock.forward) rrrrr r LayerNormr rr rrrrr_Es,r_cs0eZdZdZejdffdd ZddZZS) PatchMergingz Patch Merging Layer Args: dim (int): Number of input channels. norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm Tcsdt||_|rd|nd}tjd||dd|_|d||_tj||dd|_|||_dS)Nr!r&FrF) rr r(r r reductionnormZ expansionrg)rr(riexpandZ expand_dimrrrr s zPatchMerging.__init__c Cs|j\}}}||||jksJd|ddd||ddf|dd||dddf}}|||||}|ddkpE|ddk}|rXt|ddd|dd|df}|ddddddddddf} |ddddddddddf} |ddddddddddf} |ddddddddddf} t| | | | gd}||dd|}|ddd}tj||gdd}||}| |}|S) aV Forward function. Args: x: Input feature, tensor size (B, H*W, C), i.e., binded [PATCH, DET] tokens H, W: Spatial resolution of the input feature. Returns: x: merged [PATCH, DET] tokens; only [PATCH] tokens are reduced in spatial dim, while [DET] tokens is fix-scale rjNr!r"rr$r&r') r7rlr8rmrnr*r0repeatrzry) rrr=r>r<rpr?rXZ pad_inputZx0x1Zx2Zx3rrrrs$ > $$$$  zPatchMerging.forward rrrrr rvr rr rrrrrwsrwc sDeZdZdZdddddddejdddf fdd Zd d d ZZS) BasicLayera A basic Swin Transformer layer for one stage. Args: dim (int): Number of feature channels depth (int): Depths of this stage. num_heads (int): Number of attention head. window_size (int): Local window size. Default: 7. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. drop (float, optional): Dropout rate. Default: 0.0 attn_drop (float, optional): Attention dropout rate. Default: 0.0 drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. r`raTNrFc st |_ d|_||_|_||_t f ddt |D|_ | dur=| | d|_ dSd|_ dS)Nr!csPg|]$}t |ddkrdn dttr!|nd qS)r!r) r(rIr;rbrcrSrTrrOrfri)r_ isinstancelist.0i rOr(rrfrcrirIrTrSr;rr Gs$  z'BasicLayer.__init__..)r(rir{) rr r;rbdepthr(use_checkpointr ModuleListrangeblocks downsample)rr(rrIr;rcrSrTrrOrfrirlastrrrrr /s     zBasicLayer.__init__c Cs|jd}tt||j|j}tt||j|j} tjd|| df|jd} td|j t|j |j t|j df} td|j t|j |j t|j df} d} | D]}| D]}| | dd||ddf<| d7} q_q[t | |j}| d|j|j}| d| d}| |dktd |dktd}|r|jdd\}}||kr||kstj|d||fd tjd}t|||j}|}| |dktd |dktd}| ||| d d}tj|d|jfdd }nd}d}||f}t|jD]8\}}|||_|_|rd }|}|}nd }d}d|f}|jr7tj||||||d }q ||||||d }q |jdurd||||}|dd|dd}}||||||fS||||||fS)aD Forward function. Args: x: Input feature, tensor size (B, H*W, C). H, W: Spatial resolution of the input feature. det_pos: pos encoding for det token input_mask: padding mask for inputs cross_attn: whether to use cross attn [det x [det + patch]] rr")r%Nr$r!gYrsize)valueTF)r4rYrZ)r7rBnpceilr;r*rKr%slicerbrAr8rWZ masked_fillfloatrm interpolatetoboolr5r(rnrl enumeraterr=r>r checkpointr)rrr=r>rt input_maskrYr<rqrrZimg_maskZh_slicesZw_slicesZcnthwZ mask_windowsruZ_HZ_WrsrZr4Zn_blkZblkZ _cross_attnZ_cross_attn_mask_posZx_downWhWwrrrr^s                zBasicLayer.forward)Fr~rrrrrs/rcs2eZdZdZ    d fdd Zdd ZZS) PatchEmbedaE Image to Patch Embedding Args: patch_size (int): Patch token size. Default: 4. in_chans (int): Number of input image channels. Default: 3. embed_dim (int): Number of linear projection output channels. Default: 96. norm_layer (nn.Module, optional): Normalization layer. Default: None r&r)`NcsXtt|}||_||_||_tj||||d|_|dur'|||_ dSd|_ dS)N)Z kernel_sizeZstride) rr r patch_sizein_chans embed_dimr ZConv2drPrz)rrrrrirrrr s  zPatchEmbed.__init__c Cs|\}}}}||jddkr#t|d|jd||jdf}||jddkr@t|ddd|jd||jdf}||}|jdurp|d|d}}|ddd}||}|ddd|j ||}|S)zForward function.r"rNr!r)r$) rrrmrnrPrzr/rVr8r)rrr\r=r>rrrrrrs$   zPatchEmbed.forward)r&r)rNr^rrrrrs rcseZdZdZddddgdgddd d d d d d ejdd gdddffdd ZddZej j ddZ dddgfddZ ddZ d!fdd Zdd ZZS)"SwinTransformera Swin Transformer backbone. A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows` - https://arxiv.org/pdf/2103.14030 Args: pretrain_img_size (int): Input image size for training the pretrained model, used in absolute position embedding. Default 224. patch_size (int | tuple(int)): Patch size. Default: 4. in_chans (int): Number of input image channels. Default: 3. embed_dim (int): Number of linear projection output channels. Default: 96. depths (tuple[int]): Depths of each Swin Transformer stage. num_heads (tuple[int]): Number of attention head of each stage. window_size (int): Window size. Default: 7. mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4. qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. drop_rate (float): Dropout rate. attn_drop_rate (float): Attention dropout rate. Default: 0. drop_path_rate (float): Stochastic depth rate. Default: 0.2. norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. ape (bool): If True, add absolute position embedding to the patch embedding. Default: False. patch_norm (bool): If True, add normalization after patch embedding. Default: True. out_indices (Sequence[int]): Output from which stages. frozen_stages (int): Stages to be frozen (stop grad and set eval mode). -1 means not freezing any args. use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. r&r)r)r!r!r!)r)r r`raTNrg?F)r"r!r)r$cst||_t||_|_||_||_||_||_ t |||jr%|ndd|_ |jr\t |}t |}|d|d|d|dg}t td|d|d|_t|jddt j| d|_ddtd| t|D}t |_t|jD]F}ttd |||||||| | | | |t|d|t|d|d|||jkrtnd||jdkrdnd |d }|j|q{fd dt|jD}||_|D]}|||}d |}|||q| dS)N)rrrrirr"rGrH)pcSsg|]}|qSr)item)rrrrrrRsz,SwinTransformer.__init__..r!T)r(rrIr;rcrSrTrrOrfrirrrcsg|] }td|qS)r!)rBrrrrrmsrz)!rr pretrain_img_sizelen num_layersrape patch_norm out_indices frozen_stagesr patch_embedrr rJr*rKabsolute_pos_embedrrpos_dropZlinspacerMrlayersrrrBrwappend num_featuresZ add_module_freeze_stages)rrrrrZdepthsrIr;rcrSrTZ drop_rateZattn_drop_rateZdrop_path_raterirrrrrpatches_resolutionZdprZi_layerlayerrZ layer_namerrrr st    &    zSwinTransformer.__init__cCs|jdkr|j|jD]}d|_q|jdkr!|jr!d|j_|jdkrI|jtd|jdD]}|j |}||D]}d|_qBq3dSdS)NrFr"r!) rreval parametersZ requires_gradrrrrr)rparamrmrrrrzs      zSwinTransformer._freeze_stagescCsddhS)N det_pos_embed det_tokenrrrrrno_weight_decayszSwinTransformer.no_weight_decaydrxcs |_|_ttd|jd_tjdd_|_ td||}t|dd}tj|_ fddt t jdD_ |dkrOj j |_dt jt j_jD]}||_|jd urp||j_|jD] }||_t||j|_qsqb|d krd jd _d Sd S) a* A function to add necessary (leanable) variables to Swin Transformer for object detection Args: method: vidt or vidt_wo_neck det_token_num: the number of object to detect, i.e., number of object queries pos_dim: the channel dimension of positional encodings for [DET] and [PATCH] tokens cross_indices: the indices where to use the [DET X PATCH] cross-attention there are four possible stages in [0, 1, 2, 3]. 3 indicates Stage 4 in the ViDT paper. r"rrGrHcsg|] }j|dqS)r")rrrrrrs z0SwinTransformer.finetune_det..Zvidtr!NZ vidt_wo_neckr$)methodrlr rJr*rKrrrpos_dimrrrZ num_channelsr cross_indicesr mask_divisorrrr r(ro)rrrlrrrrblockrrr finetune_dets<     zSwinTransformer.finetune_detc Cs|jd|jd|jd}}}||}|d|d}}|ddd}||}|j|dd}|j}t j |d ||j ||j fd tjd}g} t|jD]l} |j| } | |jvrhdnd } tj||gdd }| |||||| d \} }}}}}|ddd|j ddf|dd|j dddf}}| dkr| ddd|j ddf|||ddddd}| |qZ| ||||ddddd| dd|j dddfddd}|ddd}| \}}}}||||||fS) a Forward function. Args: x: input rgb images mask: input padding masks [0: rgb values, 1: padded values] Returns: patch_outs: multi-scale [PATCH] tokens (four scales are used) these tokens are the first input of the neck decoder det_tgt: final [DET] tokens obtained at the last stage this tokens are the second input of the neck decoder det_pos: the learnable pos encoding for [DET] tokens. these encodings are used to generate reference points in deformable attention rr!r)r"r$NrTFr')rrtrY)r7rrr/rVrrr{rrmrrrrr*rrrrrr0rlr8r9r)rrr1r<r\rrrrtZ patch_outsZstagerrYZx_outr=r>Z patch_outZdet_tgtZ features_0Z features_1Z features_2Z features_3rrrrs\"     &    "( zSwinTransformer.forwardcstt|||dS)z?Convert the model into training mode while keep layers freezed.N)rrtrainr)rmoderrrrs zSwinTransformer.traincCsnd}||j7}t|jD] \}}||7}q||j|jd|jdd|j7}||j|j7}|S)Nrr"r!)rflopsrrrrrZ num_classes)rrrrrrrrs zSwinTransformer.flops)T)rrrrr rvr rr*Zjitignorerrrrrr rrrrrs@`  9Or)mathosnumpyrr*Ztorch.nnr Ztorch.nn.functionalZ functionalrmZtorch.utils.checkpointutilsrZtimm.models.layersrrrModulerpir5rArCrDr_rwrrrrrrrs.  -L>02