o ߥi3@sddlZddlmZddedefddZGdd d ejZdd d Zdd dZ GdddejZ GdddejZ GdddejZ dS)NF drop_probtrainingcCsd|dks|s|Sd|}|jdfd|jd}|tj||j|jd}||||}|S)a/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.sh different form of dropout in a.sh 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.sh layer name and use 'survival rate' as the argument. rrr)dtypedevice)shapendimtorchrandrrfloor_div)xrr keep_probr random_tensoroutputr\/home/ubuntu/.local/lib/python3.10/site-packages/modelscope/models/multi_modal/ofa/resnet.py drop_paths   rcs*eZdZdZdfdd ZddZZS)DropPathz^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'srrc Cstj||d|||d|dS)z3x3 convolution with paddingF) kernel_sizestridepaddinggroupsbiasdilationnnConv2d) in_planes out_planesr&r(r*rrrconv3x33sr0cCstj||d|ddS)z1x1 convolutionrF)r%r&r)r+)r.r/r&rrrconv1x1@s r1cs6eZdZdZ      dfdd ZddZZS) BasicBlockrN@c stt||durtj}|dks|dkrtd|dkr"tdt||||_|||_ tj dd|_ t|||_ |||_ ||_||_dS)Nrr3z3BasicBlock only supports groups=1 and base_width=64z(Dilation > 1 not supported in BasicBlockTinplace)rr2rr, BatchNorm2d ValueErrorNotImplementedErrorr0conv1bn1ReLUreluconv2bn2 downsampler&) rinplanesplanesr&r?r( base_widthr* norm_layerrrrrIs$     zBasicBlock.__init__cCsJr)r9r:r<r=r>r?rridentityoutrrrrdszBasicBlock.forward)rNrr3rNrr r! expansionrrr#rrrrr2Fsr2cs8eZdZdZ       d fdd Zdd ZZS) BottleneckrNr3rc stt||durtj}t||d|} t|| |_|| |_t | | ||||_ || |_ t| ||j |_ |||j |_tjdd|_||_||_| dkrZt| |_dSt|_dS)NgP@Tr4r)rrIrr,r6intr1r9r:r0r=r>rHconv3bn3r;r<r?r&rIdentityr) rr@rAr&r?r(rBr*rCdrop_path_ratewidthrrrrs&     zBottleneck.__init__cCs|}||}||}||}||}||}||}||}||}|jdur4||}|||}||}|Sr) r9r:r<r=r>rLrMr?rrDrrrrs           zBottleneck.forward)rNrr3rNrrGrrrrrIxsrIcsNeZdZdZ      dfdd Z   dd d Zd d Zd dZZS)ResNeta Deep residual network, copy from https://github.com/pytorch/vision/blob/main/torchvision/models/resnet.py. You can see more details from https://arxiv.org/abs/1512.03385 step 1. Get image embedding with `7` as the patch image size, `2` as stride. step 2. Do layer normalization, relu activation and max pooling. step 3. Go through three times residual branch. Frr3Nrc stt||durtj}||_d|_d|_|durgd}t|dkr,t d |||_ ||_ tj d|jdddd d |_||j|_tjd d |_tjdddd |_|jtd|d|d|_|jtd|dd|d|d|_|jtd|dd|d|d|_|D].}t|tj rtjj|jdddqt|tjtjtjfrtj |jdtj |j!dq|r|D]!}t|trtj |j"jdqt|t#rtj |j$jdqdSdS)a Args: layers (`Tuple[int]`): There are three layers in resnet, so the length of layers should greater then three. And each element in `layers` is the number of `Bottleneck` in relative residual branch. zero_init_residual (`bool`, **optional**, default to `False`): Whether or not to zero-initialize the last BN in each residual branch. groups (`int`, **optional**, default to `1`): The number of groups. So far, only the value of `1` is supported. width_per_group (`int`, **optional**, default to `64`): The width in each group. So far, only the value of `64` is supported. replace_stride_with_dilation (`Tuple[bool]`, **optional**, default to `None`): Whether or not to replace stride with dilation in each residual branch. norm_layer (`torch.nn.Module`, **optional**, default to `None`): The normalization module. If `None`, will use `torch.nn.BatchNorm2d`. drop_path_rate (`float`, **optional**, default to 0.0): Drop path rate. See more details about drop path from https://arxiv.org/pdf/1605.07648v4.pdf. Nr3r)FFFr$zHreplace_stride_with_dilation should be None or a 3-element tuple, got {}F)r%r&r'r)Tr4)r%r&r'r)rO)r&dilaterOfan_outr<)mode nonlinearity)%rrQrr,r6 _norm_layerr@r*lenr7formatr(rBr-r9r:r;r< MaxPool2dmaxpool _make_layerrIlayer1layer2layer3modules isinstanceinitkaiming_normal_weight SyncBatchNorm GroupNorm constant_r)rMr2r>) rlayerszero_init_residualr(width_per_groupreplace_stride_with_dilationrCrOmrrrrsz        zResNet.__init__c Cs|j}d}|j} |r|j|9_d}|dks|j||jkr2tt|j||j||||j}g} | ||j||||j|j | |||j|_ddt d||D} t d|D]} | ||j||j|j |j|| | dq\tj| S)a5 Making a single residual branch. step 1. If dilate==`True`, switch the value of dilate and stride. step 2. If the input dimension doesn't equal to th output output dimension in `block`, initialize a down sample module. step 3. Build a sequential of `blocks` number of `block`. Args: block (`torch.nn.Module`): The basic block in residual branch. planes (`int`): The output dimension of each basic block. blocks (`int`): The number of `block` in residual branch. stride (`int`, **optional**, default to `1`): The stride using in conv. dilate (`bool`, **optional**, default to `False`): Whether or not to replace dilate with stride. drop_path_rate (`float`, **optional**, default to 0.0): Drop path rate. See more details about drop path from https://arxiv.org/pdf/1605.07648v4.pdf. Returns: A sequential of basic layer with type `torch.nn.Sequential[block]` NrcSsg|]}|qSr)item).0rrrr Csz&ResNet._make_layer..r)r(rBr*rCrO) rZr*r@rHr, Sequentialr1appendr(rBr linspacerange) rblockrAblocksr&rUrOrCr?previous_dilationrkdprirrrr_s@   zResNet._make_layercCsJ||}||}||}||}||}||}||}|Sr)r9r:r<r^r`rarbrrrr _forward_implQs       zResNet._forward_implcCs ||Sr)r|rrrrr[s zResNet.forward)Frr3NNr)rFr) rr r!r"rr_r|rr#rrrrrQs Y > rQ)rF)rrrr) r torch.nnr,floatboolrModulerr0r1r2rIrQrrrrs  2;