U Ph@sddlmZddlZddlmZddlZddlmZddlm Z ddl m Z ddl m Z ddlmZmZmZdgZGd ddejZdS) ) annotationsN)Sequence)PatchEmbeddingBlockTransformerBlock)Conv)deprecated_argensure_tuple_repis_sqrt ViTAutoEnccs`eZdZdZeddddddddddddddddddddddddfdd ZddZZS)r a Vision Transformer (ViT), based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " Modified to also give same dimension outputs as the input size of the image pos_embedz1.2z1.4 proj_typezplease use `proj_type` instead.)namesinceremovednew_name msg_suffix convFintzSequence[int] | intstrfloatboolNone) in_channelsimg_size patch_size out_channels deconv_chns hidden_sizemlp_dim num_layers num_headsr r dropout_rate spatial_dimsqkv_bias save_attnreturnc stt|s"td|dt|| |_t|| |_| |_t|j|jD]*\}}||dkrNtd|d|dqNt |||| |jd|_ t fddt |D|_t |_ttj|jf}d d|jD}||||d |_|||||d |_d S) 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 out_channels: number of output channels. Defaults to 1. deconv_chns: number of channels for the deconvolution layers. Defaults to 16. hidden_size: dimension of hidden layer. Defaults to 768. mlp_dim: dimension of feedforward layer. Defaults to 3072. num_layers: number of transformer blocks. Defaults to 12. num_heads: number of attention heads. Defaults to 12. proj_type: position embedding layer type. Defaults to "conv". 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. Defaults to False. .. deprecated:: 1.4 ``pos_embed`` is deprecated in favor of ``proj_type``. Examples:: # for single channel input with image size of (96,96,96), conv position embedding and segmentation backbone # It will provide an output of same size as that of the input >>> net = ViTAutoEnc(in_channels=1, patch_size=(16,16,16), img_size=(96,96,96), proj_type='conv') # for 3-channel with image size of (128,128,128), output will be same size as of input >>> net = ViTAutoEnc(in_channels=3, patch_size=(16,16,16), img_size=(128,128,128), proj_type='conv') z(patch_size should be square number, got .rz patch_size=z! should be divisible by img_size=)r r!r"r%r(r r)r*c sg|]}tqSr.0ir)r%r&r(r+r,r/S/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/networks/nets/vitautoenc.py msz'ViTAutoEnc.__init__..cSsg|]}tt|qSr/)rmathsqrtr0r/r/r4r5vs) kernel_sizestride)r r#r8r9N)super__init__r ValueErrorr r"r!r*ziprpatch_embeddingnn ModuleListrangeblocks LayerNormnormr CONVTRANSconv3d_transposeconv3d_transpose_1)selfr r!r"r#r$r%r&r'r(r r r)r*r+r,mpZ conv_transup_kernel_size __class__r3r4r;$sD4      zViTAutoEnc.__init__cCs|jdd}||}g}|jD]}||}||q"||}|dd}ddt||jD}t ||jd|jdf|}| |}| |}||fS)z Args: x: input tensor must have isotropic spatial dimensions, such as ``[batch_size, channels, sp_size, sp_size[, sp_size]]``. NrcSsg|]\}}||qSr/r/)r1srJr/r/r4r5sz&ViTAutoEnc.forward..r) shaper>rBappendrD transposer=r"torchreshaperFrG)rHx spatial_sizehidden_states_outblkdr/r/r4forward|s        zViTAutoEnc.forward) rrrrrrrrrrFF)__name__ __module__ __qualname____doc__rr;rZ __classcell__r/r/rLr4r s,4U) __future__rr6collections.abcrrStorch.nnr?Z$monai.networks.blocks.patchembeddingrZ&monai.networks.blocks.transformerblockrmonai.networks.layersr monai.utilsrr r __all__Moduler r/r/r/r4 s