o ߥiH@sddlZddlZddlmZddlmmZddlmZm Z dde de fddZ Gd d d ej ZGd d d ej ZGd ddej ZGdddej ZGdddej ZGdddej ZGdddej ZdS)N)_pair_tripleF drop_probtrainingcCsd|dks|s|Sd|}|jdfd|jd}|tj||j|jd}||||}|S)a From https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py. Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the argument. rr)r)dtypedevice)shapendimtorchrandrr floor_div)xrr keep_probr random_tensoroutputri/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/cv/action_recognition/tada_convnext.py drop_path s   rcs*eZdZdZdfdd ZddZZS)DropPathz From https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py. Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). Ncstt|||_dSN)superr__init__r)selfr __class__rrr)s zDropPath.__init__cCst||j|jSr)rrrrrrrrforward-szDropPath.forwardr__name__ __module__ __qualname____doc__rr __classcell__rrrrr#srcs8eZdZdZfddZddZddZdd ZZS) TadaConvNeXta ConvNeXt A PyTorch impl of : `A ConvNet for the 2020s` - https://arxiv.org/pdf/2201.03545.pdf Args: in_chans (int): Number of input image channels. Default: 3 num_classes (int): Number of classes for classification head. Default: 1000 depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] drop_path_rate (float): Stochastic depth rate. Default: 0. layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. c stjjj}jjjjjj}jjj}jjjt jjj dr,jjj j nd}t jjj dr<jjj j nd}t |_t t j|d|ddf|ddf|ddddfdtddd d }|j|td D]"t tdd d t jdd d d }|j|qqt |_ddtd|t|DdtdD]#t jfddt|D} |j| |7qt jddd|_dS)N T_KERNEL_SIZET_STRIDErr) kernel_sizestridepaddingư>channels_firsteps data_format)rr(r()r+r,cSsg|]}|qSr)item).0rrrr isz)TadaConvNeXt.__init__..cs(g|]}t|dqS))dimrlayer_scale_init_value)TAdaConvNeXtBlock)r5jcfgcurdimsdp_ratesir8rrr6ns r1)rrVIDEOBACKBONENUM_INPUT_CHANNELS NUM_FILTERS DROP_PATHDEPTHLARGE_SCALE_INIT_VALUEhasattrSTEMr'r)nn ModuleListdownsample_layers SequentialConv3d LayerNormappendrangestagesr linspacesumnorm) rr<in_chansdrop_path_ratedepthsstem_t_kernel_sizet_stridestemdownsample_layerstagerr;rr@sf                zTadaConvNeXt.__init__cCs>tdD]}|j||}|j||}q||gdS)Nr*)rA)rSrNrTrWmean)rrr@rrrforward_features{s zTadaConvNeXt.forward_featurescCs t|tr |d}||}|S)Nvideo) isinstancedictrcrrrrrs  zTadaConvNeXt.forwardcCsdS)N) rrrrrrget_num_layersszTadaConvNeXt.get_num_layers) r!r"r#r$rrcrrir%rrrrr&1s  ;r&c*eZdZdZdfdd ZddZZS) ConvNeXtBlocka ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back We use (2) as we find it slightly faster in PyTorch Args: dim (int): Number of input channels. drop_path (float): Stochastic depth rate. Default: 0.0 layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. rr.csttj||dd|d|_t|dd|_t|d||_t |_ td|||_ |dkr>tj |t |dd nd|_|d krLt||_dSt|_dS) Nrrmrr3r3)r+r-groupsr.rBr*rT requires_gradr)rrrLrPdwconvrQrWLinearpwconv1GELUactpwconv2 Parameterr onesgammarIdentityr)rr<r7rr8rrrrs2    zConvNeXtBlock.__init__cCs|}||}|ddddd}||}||}||}||}|jdur.|j|}|ddddd}|||}|SNrr(r3r*r)rrpermuterWrtrvrwrzrrrinputrrrrs       zConvNeXtBlock.forwardrr.r rrrrrks rkcs.eZdZdZ  dfdd ZddZZS) rQaF LayerNorm that supports two data formats: channels_last (default) or channels_first. The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). r. channels_lastcsTttt||_tt||_||_ ||_ |j dvr$t |f|_ dS)N)rr/) rrrLrxr ryweightzerosbiasr1r2NotImplementedErrornormalized_shape)rrr1r2rrrrs   zLayerNorm.__init__cCs|jdkrt||j|j|j|jS|jdkrP|jddd}||djddd}||t ||j}|jdddddf||jdddddf}|SdS)Nrr/rT)keepdimr() r2F layer_normrrrr1rbpowr sqrt)rrusrrrrs  &zLayerNorm.forward)r.rr rrrrrQs  rQcrj) r9a ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_fi rst) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back We use (2) as we find it slightly faster in PyTorch Args: dim (int): Number of input channels. drop_path (float): Stochastic depth rate. Default: 0.0 layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. rr.cs0tt|}t||dd|dd|_|jjjj}|dkr4t ||jjjj |jjjj |jj dud|_ n!|dkrNt||jjjj |jjjj |jj dud|_ ntd|t|d d |_t|d ||_t|_td |||_|d krtj|t|d dnd|_|dkrt||_dSt|_dS)Nrlrncout)r+r-rocal_dimnormal)c_inratiokernels with_bias_cal normal_lngeluzUnknown route_func_type: {}r.rBr*rTrpr) rrfloat TAdaConv2drrrCrDBRANCHROUTE_FUNC_TYPE RouteFuncMLP ROUTE_FUNC_R ROUTE_FUNC_Kr dwconv_rfRouteFuncMLPLnGelu ValueErrorformatrQrWrLrsrtrurvrwrxr ryrzrr{r)rr<r7rr8route_func_typerrrrs`             zTAdaConvNeXtBlock.__init__cCs|}||||}|ddddd}||}||}||}||}|jdur2|j|}|ddddd}|||}|Sr|) rrrr}rWrtrvrwrzrr~rrrr s      zTAdaConvNeXtBlock.forwardrr rrrrr9s &r9cs0eZdZdZ   d fdd ZddZZS) rzF The routing function for generating the calibration weights. Fh㈵>皙?cs@tt|||_||_td|_td|_tj ||ddd|_ tj |t |||dddg|ddddgd|_ t t ||ddd|_t|_tj t ||||dddg|ddddgd d |_d |j_|jjj|rtj t ||||dddg|ddddgd d |_d |j_|jjjd Sd S) z Args: c_in (int): number of input channels. ratio (int): reduction ratio for the routing function. kernels (list): temporal kernel size of the stacked 1D convolutions )Nrrrr) in_channels out_channelsr+r-r(r.r/r0F)rrr+r-rTN)rrrrrrLAdaptiveAvgPool3davgpool globalpoolrPgintarQlnrugelub skip_initrdatazero_b_bias)rrrrrbn_epsbn_mmtrrrr!sR          zRouteFuncMLPLnGelu.__init__cCsl||}||}||||}||}||}|jr/||d||dgS||dS)Nr) rrrrrrrrr)rrrrrrrXs    zRouteFuncMLPLnGelu.forward)Frrr rrrrrs7rcs>eZdZdZ      d fdd Zdd Zd d ZZS) rz Performs temporally adaptive 2D convolution. Currently, only application on 5D tensors is supported, which makes TAdaConv2d essentially a 3D convolution with temporal kernel size of 1. rrTcinc sZtt| t|}t|}t|}t|}|ddks J|ddks(J|ddks0J|ddks8J| dvs>J||_||_||_||_||_||_ ||_ | |_ t tdd||||d|d|_|rxt tdd||_n|ddt jj|jtdd|jdurt j|j\} } dt| } t j|j| | dSdS)Nrr)rrr(r)r)rrrrrrr+r,r-dilationrorrLrxr Tensorrrregister_parameterinitkaiming_uniform_mathr_calculate_fan_in_and_fan_outuniform_) rrrr+r,r-rrorrfan_in_boundrrrrmsB     zTAdaConv2d.__init__c Cst|tr|d|d}}n|}d}|j\}}}}}} |\} }} } } |ddddddd| | }|jdkrT|dddddd|jd||j|| }n|jd krq|dddddd|jd||j|| }d}|j dur|dur|ddddd |j d}n |j | | dd}t j ||||jdd|jdd|jdd|j| | d }|| | ||d |dddddd}|S) z Args: x (tensor): feature to perform convolution on. alpha (tensor): calibration weight for the base weights. W_t = alpha_t * W_b rrNr(r3r*rArr)rrr,r-rrora)relistrsizer}reshaper unsqueezerorsqueezerepeatrconv2dr,r-rview)rralphaw_alphab_alpharc_outrkhkwrthwrrrrrrrsV         zTAdaConv2d.forwardc CsHd|jd|jd|jdd|jd|jd|jdud|jd S) Nz TAdaConv2d(z, z, kernel_size=zstride=z , padding=z, bias=z , cal_dim="z"))rrr+r,r-rrrhrrr__repr__s(zTAdaConv2d.__repr__)rrrrTr)r!r"r#r$rrrr%rrrrrfs <1r)rF)rr torch.nnrLtorch.nn.functional functionalrtorch.nn.modules.utilsrrrboolrModulerr&rkrQr9rrrrrrs [-!BJ