o & i"@sddlmZddlmZddlZddlZddlmZddl mm Z ddlm Z ddl mZddlmZmZddlmZmZddlmZed d d \ZZd d hZhdZGdddejZGdddejZdS)) annotations)SequenceN) LayerNorm)build_sincos_position_embedding)Conv trunc_normal_)ensure_tuple_repoptional_import)look_up_optionzeinops.layers.torch Rearrange)nameconv perceptron>sincos learnablenonecs<eZdZdZ    ddfdd ZddZddZZS)PatchEmbeddingBlocka A patch embedding block, based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " Example:: >>> from monai.networks.blocks import PatchEmbeddingBlock >>> PatchEmbeddingBlock(in_channels=4, img_size=32, patch_size=8, hidden_size=32, num_heads=4, >>> proj_type="conv", pos_embed_type="sincos") r r in_channelsintimg_sizeSequence[int] | int patch_size hidden_size num_heads proj_typestrpos_embed_type dropout_ratefloat spatial_dimsreturnNonec sttd|krdksntd|d||dkr)td|d|dt|t|_t|t|_t|| }t|| }t ||D]\} } | | krPtd|jd kr_| | dkr_td qDt d d t ||D|_ t |t ||_||jd krttj| f||||d|_nP|jd krdd| } dddd| D} dddd | Ddddd | Dd}ddt|D}tt| d|fi|t|j||_ttd|j ||_t||_|jdkrn<|jdkrt|jdd d!d"d#n+|jd$kr)g}t ||D] \}}|||qt ||| |_n td%|jd&|!|j"dS)'a Args: in_channels: dimension of input channels. img_size: dimension of input image. patch_size: dimension of patch size. hidden_size: dimension of hidden layer. num_heads: number of attention heads. proj_type: patch embedding layer type. pos_embed_type: position embedding layer type. dropout_rate: fraction of the input units to drop. spatial_dims: number of spatial dimensions. rz dropout_rate z should be between 0 and 1.z hidden size z" should be divisible by num_heads .z+patch_size should be smaller than img_size.rz:patch_size should be divisible by img_size for perceptron.cSsg|]\}}||qSr&).0Zim_dp_dr&r&f/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/patchembedding.py Xsz0PatchEmbeddingBlock.__init__..r r out_channels kernel_sizestride))hp1)wp2)dp3Nzb c  css&|]\}}d|d|dVqdS)(r5)Nr&)r'kvr&r&r) cs$z/PatchEmbeddingBlock.__init__..zb (cSg|]}|dqS)rr&r'cr&r&r)r*dz) (cSr;)r$r&r<r&r&r)r*dr>z c)cSs i|] \}}d|d|qS)pr$r&)r'ir?r&r&r) es z0PatchEmbeddingBlock.__init__..z -> rrr{Gz?@meanstdabrzpos_embed_type z not supported.)#super__init__ ValueErrorr SUPPORTED_PATCH_EMBEDDING_TYPESrSUPPORTED_POS_EMBEDDING_TYPESrrzipnpprodZ n_patchesrZ patch_dimrCONVpatch_embeddingsjoin enumeratenn Sequentialr Linear Parametertorchzerosposition_embeddingsDropoutdropoutrappendrapply _init_weights)selfrrrrrrrrr!mr?charsZ from_charsZto_charsZaxes_len grid_sizein_sizeZpa_size __class__r&r)rK-sV          2$    zPatchEmbeddingBlock.__init__cCst|tjr)t|jdddddt|tjr%|jdur'tj|jddSdSdSt|tjrAtj|jdtj|jddSdS)NrrBrCrDrErg?) isinstancerVrXrweightbiasinit constant_r)rbrcr&r&r)ra{s  z!PatchEmbeddingBlock._init_weightscCs>||}|jdkr|ddd}||j}||}|S)Nr )rSrflatten transposer\r^)rbx embeddingsr&r&r)forwards    zPatchEmbeddingBlock.forward)r rrr)rrrrrrrrrrrrrrrr r!rr"r#)__name__ __module__ __qualname____doc__rKraru __classcell__r&r&rgr)r sN rcs8eZdZdZdddejdfdfdd ZddZZS) PatchEmbeda0 Patch embedding block based on: "Liu et al., Swin Transformer: Hierarchical Vision Transformer using Shifted Windows " https://github.com/microsoft/Swin-Transformer Unlike ViT patch embedding block: (1) input is padded to satisfy window size requirements (2) normalized if specified (3) position embedding is not used. Example:: >>> from monai.networks.blocks import PatchEmbed >>> PatchEmbed(patch_size=2, in_chans=1, embed_dim=48, norm_layer=nn.LayerNorm, spatial_dims=3) rnr$0rrrin_chansr embed_dim norm_layertype[LayerNorm]r!r"r#cslt|dvr tdt||}||_||_ttj|f||||d|_|dur1|||_ dSd|_ dS)a Args: patch_size: dimension of patch size. in_chans: dimension of input channels. embed_dim: number of linear projection output channels. norm_layer: normalization layer. spatial_dims: spatial dimension. )rnrz#spatial dimension should be 2 or 3.r+N) rJrKrLrrr~rrRprojnorm)rbrr}r~rr!rgr&r)rKs    zPatchEmbed.__init__c Cs|}t|dkri|\}}}}}||jddkr,t|d|jd||jdf}||jddkrIt|ddd|jd||jdf}||jddkrht|ddddd|jd||jdf}nDt|dkr|\}}}}||jddkrt|d|jd||jdf}||jddkrt|ddd|jd||jdf}||}|jdur |}|ddd}||}t|dkr|d|d|d}}}|dd d|j |||}|St|dkr |d|d}}|dd d|j ||}|S)Nrnrr$rro) sizelenrFpadrrrqrrviewr~) rbrsx_shape_r3r/r1whwwr&r&r)rus: $(,  $(    zPatchEmbed.forward) rrr}rr~rrrr!rr"r#) rvrwrxryrVrrKrurzr&r&rgr)r{s!r{) __future__rcollections.abcrnumpyrPrZtorch.nnrVtorch.nn.functional functionalrrZ%monai.networks.blocks.pos_embed_utilsrmonai.networks.layersrr monai.utilsrr monai.utils.moduler r rrMrNModulerr{r&r&r&r)s     m