o - iS @sddlmZddlZddlZddlZddlmZddlmZddl Z ddl m Z ddl m Z ddl mZddlmZmZmZmZdd lmZdd lmZgd Zd d ddddddddd Zidddddddddd d!d"d#d$d%d&d'd(d)d*d+d,d-d.d/d0d1d2d3d4d5d6d7d8ZGd9d:d:e jZGd;d<dd>eZGd?d@d@eZGdAdBdBeeZ dudGdHZ!dvdOdPZ"dwdWdXZ#dxd_d`Z$dydcddZ%dzdidjZ&d{dmdnZ'd|dqdrZ(GdsdtdteZ)dS)}) annotationsN)reduce) NamedTuple)nn) model_zoo) BaseEncoder)ActConvPadPool)get_norm_layer)look_up_option) EfficientNetEfficientNetBNget_efficientnet_image_size drop_connectEfficientNetBNFeatures BlockArgsEfficientNetEncoder)?r皙?r)r皙?rr)r333333?i333333?r)rffffff?i,rr)r?i|皙?r)g?皙@irr)rg@i?r)@g@iXr r)rg @ir r)g333333@g333333@i r r efficientnet-b0efficientnet-b1efficientnet-b2efficientnet-b3efficientnet-b4efficientnet-b5efficientnet-b6efficientnet-b7zefficientnet-b8zefficientnet-l2r#zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pthr$zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b1-f1951068.pthr%zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b2-8bb594d6.pthr&zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b3-5fb5a3c3.pthr'zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b4-6ed6700e.pthr(zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b5-b6417697.pthr)zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b6-c76e70fd.pthr*zdhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b7-dcc49843.pthzb0-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b0-b64d5a18.pthzb1-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b1-0f3ce85a.pthzb2-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b2-6e9d97e5.pthzb3-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b3-cdd7c0f4.pthzb4-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b4-44fb3a87.pthzb5-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b5-86493f6b.pthzb6-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b6-ac80338e.pthzb7-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b7-4652b6dd.pthzb8-apzhhttps://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/adv-efficientnet-b8-22a8fe65.pthcsFeZdZdddddfdfd#fdd Zd$ddZd%d&d!d"ZZS)' MBConvBlockTbatchMbP?{Gz?epsmomentumr spatial_dimsint in_channels out_channels kernel_sizestride image_size list[int] expand_ratiose_ratio float | Noneid_skip bool | Nonenorm str | tupledrop_connect_ratereturnNonec sttd|f} td|f} ||_||_| |_||_||_| |_ |dur:d|kr1dkr:nnd|_ ||_ nd|_ |}||}|jdkra| ||ddd |_ t |j ||_t| ||d |_nt|_ t|_t|_| |||||jdd |_t |j||_t| ||d |_t||j}|j r| d|_tdt||j }| ||dd |_t |jdg||_| ||dd |_t |jdg||_|}| ||ddd |_t |j||_t| ||d |_ t!d dd|_"dS)a Mobile Inverted Residual Bottleneck Block. Args: spatial_dims: number of spatial dimensions. in_channels: number of input channels. out_channels: number of output channels. kernel_size: size of the kernel for conv ops. stride: stride to use for conv ops. image_size: input image resolution. expand_ratio: expansion ratio for inverted bottleneck. se_ratio: squeeze-excitation ratio for se layers. id_skip: whether to use skip connection. norm: feature normalization type and arguments. Defaults to batch norm. drop_connect_rate: dropconnect rate for drop connection (individual weights) layers. References: [1] https://arxiv.org/abs/1704.04861 (MobileNet v1) [2] https://arxiv.org/abs/1801.04381 (MobileNet v2) [3] https://arxiv.org/abs/1905.02244 (MobileNet v3) conv adaptiveavgNrTF)r4r5r6biasnamer2channels)r4r5groupsr6r7rH)r4r5r6memswishinplace)#super__init__r r r4r5r=r7r:rAhas_ser; _expand_conv_make_same_padder_expand_conv_paddingr _bn0rIdentity_depthwise_conv_depthwise_conv_padding_bn1_calculate_output_image_size _se_adaptpoolmaxr3 _se_reduce_se_reduce_padding _se_expand_se_expand_padding _project_conv_project_conv_padding_bn2r_swish)selfr2r4r5r6r7r8r:r;r=r?rA conv_typeadaptivepool_typeinpoupZnum_squeezed_channelsZ final_oup __class__b/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/efficientnet.pyrQMsZ #         zMBConvBlock.__init__inputs torch.TensorcCs|}|jdkr|||}||}||}|||}||}||}|jrO| |}| | |}||}| | |}t||}|||}||}|jrz|jdkrz|j|jkrz|jrvt||j|jd}||}|S)zMBConvBlock"s forward function. Args: inputs: Input tensor. Returns: Output of this block after processing. rG)ptraining)r:rSrUrVrerXrYrZrRr\r^r_r`ratorchsigmoidrbrcrdr=r7r4r5rArrr)rfroxZ x_squeezedrmrmrnforwards*        zMBConvBlock.forwardmemory_efficientboolcCs,|r tddd|_dStddd|_dS)zSets swish function as memory efficient (for training) or standard (for export). Args: memory_efficient (bool): Whether to use memory-efficient version of swish. rMTrNswishralphaN)rre)rfrwrmrmrn set_swishs,zMBConvBlock.set_swish)r2r3r4r3r5r3r6r3r7r3r8r9r:r3r;r<r=r>r?r@rAr<rBrCrorpTrwrxrBrC)__name__ __module__ __qualname__rQrvr| __classcell__rmrmrkrnr+Ks  h(r+c s^eZdZdddddddddd d fdd f d*fdd Zd+d,d"d#Zd-d&d'Zd.d(d)ZZS)/rrrrr,r-r.r/blocks_args_str list[str]r2r3r4 num_classeswidth_coefficientfloatdepth_coefficient dropout_rater8r?r@rA depth_divisorrBrCc st|dvr tdtd|f} td|f} dd|D}t|ts)td|gkr1td||_||_||_ | |_ |g|}d }t d || }| |j |d |d d |_ t |j ||_t| ||d|_t||}d}t|_g|_t|jD]/\}}|jt |j|| t |j|| t|j|d}||j|<||j7}|jdkr|j|qx|jt|jd}t|jD]\}}|j }|r|t||9}t}| t!|t"||j|j|j#|j||j$|j%|j&| |d |d7}t||j}|jdkr|j|jdd}t'|jdD]1}|j }|r|t||9}| t!|t"||j|j|j#|j||j$|j%|j&| |d |d7}q|j t!||q||krMtd|j}t d|| }| ||dd d|_(t |j(||_)t| ||d|_*| d|_+t,||_-t.||j|_/t0d|_1|2dS)a EfficientNet based on `Rethinking Model Scaling for Convolutional Neural Networks `_. Adapted from `EfficientNet-PyTorch `_. Args: blocks_args_str: block definitions. spatial_dims: number of spatial dimensions. in_channels: number of input channels. num_classes: number of output classes. width_coefficient: width multiplier coefficient (w in paper). depth_coefficient: depth multiplier coefficient (d in paper). dropout_rate: dropout rate for dropout layers. image_size: input image resolution. norm: feature normalization type and arguments. drop_connect_rate: dropconnect rate for drop connection (individual weights) layers. depth_divisor: depth divisor for channel rounding. )rGrrz#spatial_dims can only be 1, 2 or 3.rDrEcSsg|]}t|qSrm)r from_string.0srmrmrn sz)EfficientNet.__init__..zblocks_args must be a listzblock_args must be non-emptyr rF)r6r7rHrIr) input_filtersoutput_filters num_repeatrG) r2r4r5r6r7r8r:r;r=r?rA)rr7z,total number of blocks created != num_blocksi)r6rHrMN)3rPrQ ValueErrorr r isinstancelistZ _blocks_argsrr4rA_round_filters _conv_stemrT_conv_stem_paddingr rVr[r Sequential_blocksextract_stacks enumerate_replacerr_round_repeatsrr7appendlenr add_modulestrr+r6r:r;r=range _conv_head_conv_head_paddingrZ _avg_poolingDropout_dropoutLinear_fcrre_initialize_weights)rfrr2r4rrrrr8r?rArrgrhZ blocks_argsZcurrent_image_sizer7r5 num_blocksidxZ block_argsZ stack_idxZblk_drop_connect_rate sub_stack_Zhead_in_channelsrkrmrnrQs                       zEfficientNet.__init__TrwrxcCsB|rtdntddd|_|jD] }|D]}||qqdS)z Sets swish function as memory efficient (for training) or standard (for JIT export). Args: memory_efficient: whether to use memory-efficient version of swish. rMryrrzN)rrerr|)rfrwrblockrmrmrnr|s   zEfficientNet.set_swishrorpcCsx|||}|||}||}|||}|||}||}|j dd}| |}| |}|S)a! Args: inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. Returns: a torch Tensor of classification prediction in shape ``(Batch, num_classes)``. rG) start_dim) rrrerVrrrrZrflattenrr)rfrorurmrmrnrvs      zEfficientNet.forwardcCs|D]r\}}t|tjtjtjfr7ttj|j d|j }|j j dtd||jdur6|jj qt|tjtjtjfrP|j j d|jj qt|tjrv|j d}d}dt||}|j j | ||jj qdS)a Args: None, initializes weights for conv/linear/batchnorm layers following weight init methods from `official Tensorflow EfficientNet implementation `_. Adapted from `EfficientNet-PyTorch's init method `_. rGrr!Nr) named_modulesrrConv1dConv2dConv3droperatormulr6r5weightdatanormal_mathsqrtrHzero_ BatchNorm1d BatchNorm2d BatchNorm3dfill_rsizeuniform_)rfrmfan_outfan_in init_rangermrmrnrs$      z EfficientNet._initialize_weights)rrr2r3r4r3rr3rrrrrrr8r3r?r@rArrr3rBrCr~rr})rBrC)rrrrQr|rvrrrmrmrkrnrs   3 rc s8eZdZdddddddddfd fdfdd ZZS)rTrrrr,r-r.r/F model_namer pretrainedrxprogressr2r3r4rr?r@adv_proprBrCc  gd} |tvrdt} td|d| dt|\} } } }}tj| |||| | || ||d |rB|dkrDt||||dSdSdS) a Generic wrapper around EfficientNet, used to initialize EfficientNet-B0 to EfficientNet-B7 models model_name is mandatory argument as there is no EfficientNetBN itself, it needs the N in [0, 1, 2, 3, 4, 5, 6, 7, 8] to be a model Args: model_name: name of model to initialize, can be from [efficientnet-b0, ..., efficientnet-b8, efficientnet-l2]. pretrained: whether to initialize pretrained ImageNet weights, only available for spatial_dims=2 and batch norm is used. progress: whether to show download progress for pretrained weights download. spatial_dims: number of spatial dimensions. in_channels: number of input channels. num_classes: number of output classes. norm: feature normalization type and arguments. adv_prop: whether to use weights trained with adversarial examples. This argument only works when `pretrained` is `True`. Examples:: # for pretrained spatial 2D ImageNet >>> image_size = get_efficientnet_image_size("efficientnet-b0") >>> inputs = torch.rand(1, 3, image_size, image_size) >>> model = EfficientNetBN("efficientnet-b0", pretrained=True) >>> model.eval() >>> outputs = model(inputs) # create spatial 2D >>> model = EfficientNetBN("efficientnet-b0", spatial_dims=2) # create spatial 3D >>> model = EfficientNetBN("efficientnet-b0", spatial_dims=3) # create EfficientNetB7 for spatial 2D >>> model = EfficientNetBN("efficientnet-b7", spatial_dims=2) zr1_k3_s11_e1_i32_o16_se0.25zr2_k3_s22_e6_i16_o24_se0.25zr2_k5_s22_e6_i24_o40_se0.25zr3_k3_s22_e6_i40_o80_se0.25zr3_k5_s11_e6_i80_o112_se0.25zr4_k5_s22_e6_i112_o192_se0.25zr1_k3_s11_e6_i192_o320_se0.25, invalid model_name  found, must be one of  rr2r4rrrrr8rAr?rNefficientnet_paramsjoinkeysrrPrQ_load_state_dictrfrrrr2r4rr?rrmodel_name_stringZ weight_coeffZ depth_coeffr8rZdropconnect_raterkrmrnrQs(0  zEfficientNetBN.__init__rrrrxrrxr2r3r4r3rr3r?r@rrxrBrC)rrrrQrrmrmrkrnrs rc sBeZdZdddddddddfd fdfdd ZdddZZS) rTrrrr,r-r.r/Frrrrxrr2r3r4rr?r@rrBrCc  r) a{ Initialize EfficientNet-B0 to EfficientNet-B7 models as a backbone, the backbone can be used as an encoder for segmentation and objection models. Compared with the class `EfficientNetBN`, the only different place is the forward function. This class refers to `PyTorch image models `_. rrrrrrrNrrrkrmrnrQ7s(  zEfficientNetBNFeatures.__init__rorpcCsp|||}|||}g}d|jvr||t|jD]\}}||}|d|jvr5||q!|S)z Args: inputs: input should have spatially N dimensions ``(Batch, in_channels, dim_0[, dim_1, ..., dim_N])``, N is defined by `dimensions`. Returns: a list of torch Tensors. rrG)rrrerVrrrr)rfrorufeaturesirrmrmrnrvns    zEfficientNetBNFeatures.forwardrr})rrrrQrvrrmrmrkrnr5s 7rc@sPeZdZdZgdZedddZeddd Zedd d ZedddZ dS)rzI Wrap the original efficientnet to an encoder for flexible-unet. r"rB list[dict]c Cs>g}|jD]}||dddddddddfd |vd q|S) zN Get the initialization parameter for efficientnet backbones. Trrrr,r-r.r/ap)rrrr2r4rr?r)backbone_namesr)clsparameter_list backbone_namermrmrnget_encoder_parameterss   z*EfficientNetEncoder.get_encoder_parameterslist[tuple[int, ...]]cCsgdS)zS Get number of efficientnet backbone output feature maps' channel. ) (pi@r)rr0xi`)rrri)rr8i)rr@i)rrHi@)rrPri)rrXi)rhrii`rmrrmrmrnnum_channels_per_outputsz+EfficientNetEncoder.num_channels_per_outputr9cCs dgdS)z Get number of efficientnet backbone output feature maps. Since every backbone contains the same 5 output feature maps, the number list should be `[5] * 10`.  rmrrmrmrn num_outputss zEfficientNetEncoder.num_outputsrcCs|jS)z2 Get names of efficient backbone. )rrrmrmrnget_encoder_namessz%EfficientNetEncoder.get_encoder_namesN)rBr)rBr)rBr9)rBr) rrr__doc__r classmethodrrrrrmrmrmrnrs   rrrrBr3cCsB|tvrdt}td|d|dt|\}}}}}|S)z Get the input image size for a given efficientnet model. Args: model_name: name of model to initialize, can be from [efficientnet-b0, ..., efficientnet-b7]. Returns: Image size for single spatial dimension as integer. rrrr)rrrr)rrrresrmrmrnrs  rrorprqrrrrxc Cs|dks|dkrtd||s|S|jd}d|}t|jd}|dgdg|}tj||j|jd}||7}t|}|||} | S)ah Drop connect layer that drops individual connections. Differs from dropout as dropconnect drops connections instead of whole neurons as in dropout. Based on `Deep Networks with Stochastic Depth `_. Adapted from `Official Tensorflow EfficientNet utils `_. This function is generalized for MONAI's N-Dimensional spatial activations e.g. 1D activations [B, C, H], 2D activations [B, C, H, W] and 3D activations [B, C, H, W, D] Args: inputs: input tensor with [B, C, dim_1, dim_2, ..., dim_N] where N=spatial_dims. p: probability to use for dropping connections. training: whether in training or evaluation mode. Returns: output: output tensor after applying drop connection. rFrz$p must be in range of [0, 1], found rrGr)dtypedevice)rshaperrsrandrrfloor) rorqrr batch_size keep_probnum_dimsZrandom_tensor_shape random_tensorZ binary_tensoroutputrmrmrnrs   rmodel nn.ModulearchrrrCc Cs|r |ddd}t|td}|durtd|ddSt|}tj||d}|}td}| D]\}} t |d|} | |vrT| j || j krT|| ||<q7| |dS) Nz efficientnet-z-apzpretrained weights of z is not provided)rz(.+)\.\d+(\.\d+\..+)z\1\2) splitr url_mapprintrload_url state_dictrecompileitemssubr load_state_dict) rrrr model_urlZpretrain_state_dictmodel_state_dictpatternkeyvalueZ pretrain_keyrmrmrnrs   rr8r9r6tuple[int, ...]dilationr7cCslt|}t|dkr||}t|dkr||}ddt||||D}dd|D}ddt|D}|S)a/ Helper for getting padding (nn.ConstantPadNd) to be used to get SAME padding conv operations similar to Tensorflow's SAME padding. This function is generalized for MONAI's N-Dimensional spatial operations (e.g. Conv1D, Conv2D, Conv3D) Args: image_size: input image/feature spatial size. kernel_size: conv kernel's spatial size. dilation: conv dilation rate for Atrous conv. stride: stride for conv operation. Returns: paddings for ConstantPadNd padder to be used on input tensor to conv op. rGcSsDg|]\}}}}tt||d||d|d|dqS)rGr)r]rceil)rZ_i_sZ_k_s_d_srmrmrnrEs .z-_get_same_padding_conv_nd..cSs g|] }|d||dfqS)rrm)r_prmrmrnrJs cSsg|] }|D]}|qqSrmrm)rinnerouterrmrmrnrNs)rzipreversed)r8r6r%r7r _pad_size _paddingsZ _paddings_retrmrmrn_get_same_padding_conv_nd(s   r0conv_op!nn.Conv1d | nn.Conv2d | nn.Conv3dcCsHt||j|j|j}tdt|df}t|dkr ||ddStS)a  Helper for initializing ConstantPadNd with SAME padding similar to Tensorflow. Uses output of _get_same_padding_conv_nd() to get the padding size. This function is generalized for MONAI's N-Dimensional spatial operations (e.g. Conv1D, Conv2D, Conv3D) Args: conv_op: nn.ConvNd operation to extract parameters for op from image_size: input image/feature spatial size Returns: If padding required then nn.ConstandNd() padder initialized to paddings otherwise nn.Identity() Z constantpadrrrF)paddingr#) r0r6r%r7r rsumrrW)r1r8r3padderrmrmrnrTRs   rTfiltersrr<rcCsR|s|S|}|}||}t|t||d||}|d|kr%||7}t|S)aN Calculate and round number of filters based on width coefficient multiplier and depth divisor. Args: filters: number of input filters. width_coefficient: width coefficient for model. depth_divisor: depth divisor to use. Returns: new_filters: new number of filters after calculation. rg?)r]r3)r6rr multiplierdivisorZ filters_floatZ new_filtersrmrmrnrjs  rrepeatsrcCs|s|Stt||S)a Re-calculate module's repeat number of a block based on depth coefficient multiplier. Args: repeats: number of original repeats. depth_coefficient: depth coefficient for model. Returns: new repeat: new number of repeat after calculating. r3rr&)r9rrmrmrnrs rinput_image_sizeint | tuple[int]csLttrtfddD}|stddfdd|DS)a5 Calculates the output image size when using _make_same_padder with a stride. Required for static padding. Args: input_image_size: input image/feature spatial size. stride: Conv2d operation"s stride. Returns: output_image_size: output image/feature spatial size. c3s|] }d|kVqdS)rNrmrr7rmrn sz/_calculate_output_image_size..z&unequal strides are not possible, got rcsg|] }tt|qSrmr:)rZim_szr=rmrnrsz0_calculate_output_image_size..)rtupleallr)r;r7Zall_strides_equalrmr=rnr[s r[c@sleZdZUdZded<ded<ded<ded<ded<ded<d ed <d Zd ed <edddZddZd S)rzq BlockArgs object to assist in decoding string notation of arguments for MBConvBlock definition. r3rr6r7r:rrrxr=Nr<r; block_stringrc Cs<|d}i}|D]}td|}t|dkr#|dd\}}|||<q d|vr0t|ddkpct|ddkoC|dd|ddkpct|ddkoc|dd|ddkoc|dd|ddk}|sjtd tt|d t|d t|ddt|d t|d t|dd|vd|vrt|ddSddS)a> Get a BlockArgs object from a string notation of arguments. Args: block_string (str): A string notation of arguments. Examples: "r1_k3_s11_e1_i32_o16_se0.25". Returns: BlockArgs: namedtuple defined at the top of this function. rz(\d.*)rNrrGrrzinvalid stride option receivedrkeroZnoskipse)rr6r7r:rrr=r;)rrrrrr3r)rAopsoptionsopsplitsr"r#Z stride_checkrmrmrnrs6   &>     zBlockArgs.from_stringcCsTd|jd|jd|j|jd|jd|jd|jd|j}|js(|d7}|S) z Return a block string notation for current BlockArgs object Returns: A string notation of BlockArgs object arguments. Example: "r1_k3_s11_e1_i32_o16_se0.25_noskip". rB_kr(_e_i_o_seZ_noskip)rr6r7r:rrr;r=)rfstringrmrmrn to_strings zBlockArgs.to_string)rAr) rrrr__annotations__r; staticmethodrrQrmrmrmrnrs    'r)rrrBr3)rorprqrrrrxrBrp) rrrrrrxrrxrBrC) r8r9r6r$r%r$r7r$rBr9)r1r2r8r9)r6r3rr<rrrBr3)r9r3rr<rBr3)r;r9r7r<)* __future__rrrr functoolsrtypingrrsr torch.utilsrmonai.networks.blocksrmonai.networks.layers.factoriesrr r r monai.networks.layers.utilsr monai.utils.moduler __all__rrModuler+rrrrrrrr0rTrrr[rrmrmrmrns             yWP J  /  *