o - i%@sddlmZddlmZddlZddlZddlmZddl m Z ddl m Z ddl mZddlmZddlmZdd lmZhd Zd gZGd d d ejZdS) ) annotations)SequenceN)PatchEmbeddingBlock)build_sincos_position_embeddingTransformerBlock) trunc_normal_)ensure_tuple_rep)look_up_option>sincos learnablenoneMaskedAutoEncoderViTcsdeZdZdZ           d1d2fd'd( Zd)d*Zd3d4d-d.Zd3d4d/d0ZZS)5ra3 Masked Autoencoder (ViT), based on: "Kaiming et al., Masked Autoencoders Are Scalable Vision Learners " Only a subset of the patches passes through the encoder. The decoder tries to reconstruct the masked patches, resulting in improved training speed.  ?convr F in_channelsintimg_sizeSequence[int] | int patch_size hidden_sizemlp_dim num_layers num_heads masking_ratiofloatdecoder_hidden_sizedecoder_mlp_dimdecoder_num_layersdecoder_num_heads proj_typestrpos_embed_typedecoder_pos_embed_type dropout_rate spatial_dimsqkv_biasbool save_attnreturnNonec stdkrdksntdddkr"tddkr,tdt|||_t|||_||_t|j|jD]\}}||dkrWtd|d|dqB|_|dksc|dkrktd |d||_ t t dd|_t|||| ||jd |_fd d t|D}t jg|t R|_t |_t t dd|_t|t|_t t d|jj|_fd d t| D}t jg|t R|_t tt |j||_!|"dS)a Args: in_channels: dimension of input channels or the number of channels for input. img_size: dimension of input image. patch_size: dimension of patch size hidden_size: dimension of hidden layer. Defaults to 768. mlp_dim: dimension of feedforward layer. Defaults to 512. num_layers: number of transformer blocks. Defaults to 12. num_heads: number of attention heads. Defaults to 12. masking_ratio: ratio of patches to be masked. Defaults to 0.75. decoder_hidden_size: dimension of hidden layer for decoder. Defaults to 384. decoder_mlp_dim: dimension of feedforward layer for decoder. Defaults to 512. decoder_num_layers: number of transformer blocks for decoder. Defaults to 4. decoder_num_heads: number of attention heads for decoder. Defaults to 12. proj_type: position embedding layer type. Defaults to "conv". pos_embed_type: position embedding layer type. Defaults to "sincos". decoder_pos_embed_type: position embedding layer type for decoder. Defaults to "sincos". dropout_rate: fraction of the input units to drop. Defaults to 0.0. spatial_dims: number of spatial dimensions. Defaults to 3. qkv_bias: apply bias to the qkv linear layer in self attention block. Defaults to False. save_attn: to make accessible the attention in self attention block. Defaults to False. Examples:: # for single channel input with image size of (96,96,96), and sin-cos positional encoding >>> net = MaskedAutoEncoderViT(in_channels=1, img_size=(96,96,96), patch_size=(16,16,16), pos_embed_type='sincos') # for 3-channel with image size of (128,128,128) and a learnable positional encoding >>> net = MaskedAutoEncoderViT(in_channels=3, img_size=128, patch_size=16, pos_embed_type='learnable') # for 3-channel with image size of (224,224) and a masking ratio of 0.25 >>> net = MaskedAutoEncoderViT(in_channels=3, img_size=(224,224), patch_size=(16,16), masking_ratio=0.25, spatial_dims=2) rz,dropout_rate should be between 0 and 1, got .z-hidden_size should be divisible by num_heads.z=decoder_hidden_size should be divisible by decoder_num_heads.z patch_size=z! should be divisible by img_size=z1masking_ratio should be in the range (0, 1), got ) rrrrr r'r)r+r,c sg|] }tqSr.0_)r+rrr r-r/r4l/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/masked_autoencoder_vit.py z1MaskedAutoEncoderViT.__init__..c sg|] }tqSr4rr5)r#r$r&r+r-r/r4r8r9r:N)#super__init__ ValueErrorr rrr,zipr#r!nn Parametertorchzeros cls_tokenrpatch_embeddingrange Sequential LayerNormblocksLinear decoder_embed mask_tokensr SUPPORTED_POS_EMBEDDING_TYPESr* n_patchesdecoder_pos_embeddingdecoder_blocksrnpprod decoder_pred _init_weights)selfrrrrrrr r!r#r$r%r&r'r)r*r+r,r-r/mprHrO __class__) r#r$r&r+rrr r-r/r8r<(sX 6        zMaskedAutoEncoderViT.__init__cCs|jdkrn=|jdkrt|jdddddn-|jdkr:g}t|j|jD] \}}|||q$t||j|j |_n t d |jd t|j dddddt|j dddddd S) z similar to monai/networks/blocks/patchembedding.py for the decoder positional encoding and for mask and classification tokens r r rg{Gz?gg@)meanstdabr zdecoder_pos_embed_type z not supported.N) r*rrNr>rrappendrr#r,r=rKrC)rT grid_sizein_sizepa_sizer4r4r8rSs    z"MaskedAutoEncoderViT._init_weightsN float | Nonec Cs|j\}}}|durd|nd|j}tjt||t||dd}|t|d|f}tj||tjd|j } d| t|d|f<||| fS)Nr2F) replacement)dtyper) shaper!rA multinomialonesrarange unsqueezetodevice) rTxr! batch_size num_tokensr7Zpercentage_to_keepselected_indicesx_maskedmaskr4r4r8_maskings  zMaskedAutoEncoderViT._maskingcCs||}|j||d\}}}|j|jddd}tj||fdd}||}||}|j |jd|jdd}|ddddddf|t |jd d|f<||j }tj|ddddddf|gdd}||}||}|ddddddf}||fS)N)r!rrdr2)dim)rDrrrCexpandrerAcatrHrJrKrepeatrhrirNrOrR)rTrlr!rorqZ cls_tokensx_r4r4r8forwards   4 (  zMaskedAutoEncoderViT.forward)rrrrrrrrrrr r rrFF)(rrrrrrrrrrrrr rr!r"r#rr$rr%rr&rr'r(r)r(r*r(r+r"r,rr-r.r/r.r0r1)N)r!ra) __name__ __module__ __qualname____doc__r<rSrrrx __classcell__r4r4rWr8r s, r  ) __future__rcollections.abcrnumpyrPrAtorch.nnr?Z$monai.networks.blocks.patchembeddingr%monai.networks.blocks.pos_embed_utilsrZ&monai.networks.blocks.transformerblockrmonai.networks.layersr monai.utilsr monai.utils.moduler rL__all__Modulerr4r4r4r8s