o , i.@s`ddlmZddlmZddlZddlmZddlmZddl m Z dgZ Gdddej Z dS)) annotations)SequenceN)PatchEmbeddingBlockTransformerBlockViTcsFeZdZdZ          d&d'fd"d# Zd$d%ZZS)(rz Vision Transformer (ViT), based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " ViT supports Torchscript but only works for Pytorch after 1.8.  conv learnableFTanh in_channelsintimg_sizeSequence[int] | int patch_size hidden_sizemlp_dim num_layers num_heads proj_typestrpos_embed_typeclassificationbool num_classes dropout_ratefloat spatial_dimsqkv_bias save_attnreturnNonec stdkrdkstdtddkr!td| |_t||||| | d |_tfddt|D|_ t |_ |jrvt t dd|_|dkrmtt| t|_d St| |_d Sd S) aO Args: in_channels (int): dimension of input channels. img_size (Union[Sequence[int], int]): dimension of input image. patch_size (Union[Sequence[int], int]): dimension of patch size. hidden_size (int, optional): dimension of hidden layer. Defaults to 768. mlp_dim (int, optional): dimension of feedforward layer. Defaults to 3072. num_layers (int, optional): number of transformer blocks. Defaults to 12. num_heads (int, optional): number of attention heads. Defaults to 12. proj_type (str, optional): patch embedding layer type. Defaults to "conv". pos_embed_type (str, optional): position embedding type. Defaults to "learnable". classification (bool, optional): bool argument to determine if classification is used. Defaults to False. num_classes (int, optional): number of classes if classification is used. Defaults to 2. dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0. spatial_dims (int, optional): number of spatial dimensions. Defaults to 3. post_activation (str, optional): add a final acivation function to the classification head when `classification` is True. Default to "Tanh" for `nn.Tanh()`. Set to other values to remove this function. qkv_bias (bool, optional): apply bias to the qkv linear layer in self attention block. Defaults to False. save_attn (bool, optional): to make accessible the attention in self attention block. Defaults to False. Examples:: # for single channel input with image size of (96,96,96), conv position embedding and segmentation backbone >>> net = ViT(in_channels=1, img_size=(96,96,96), proj_type='conv', pos_embed_type='sincos') # for 3-channel with image size of (128,128,128), 24 layers and classification backbone >>> net = ViT(in_channels=3, img_size=(128,128,128), proj_type='conv', pos_embed_type='sincos', classification=True) # for 3-channel with image size of (224,224), 12 layers and classification backbone >>> net = ViT(in_channels=3, img_size=(224,224), proj_type='conv', pos_embed_type='sincos', classification=True, >>> spatial_dims=2) rz'dropout_rate should be between 0 and 1.z-hidden_size should be divisible by num_heads.) rrrrrrrr r"c sg|] }tqSr).0ir rrrr#r$r(Y/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/vit.py lsz ViT.__init__..rN)super__init__ ValueErrorrrpatch_embeddingnn ModuleListrangeblocks LayerNormnorm Parametertorchzeros cls_token SequentialLinearrclassification_head)selfrrrrrrrrrrrr r"Zpost_activationr#r$ __class__r+r,r/!s> 6   z ViT.__init__cCs||}t|dr|j|jddd}tj||fdd}g}|jD] }||}||q#| |}t|drD| |dddf}||fS)Nr;rr')dimr>) r1hasattrr;expandshaper9catr5appendr7r>)r?xr;hidden_states_outblkr(r(r,forwardys      z ViT.forward) rr r r r r Fr rrrFF) rrrrrrrrrrrrrrrrrrrrrrr r!r"rr#rr$rr%r&)__name__ __module__ __qualname____doc__r/rL __classcell__r(r(r@r,rs" X) __future__rcollections.abcrr9torch.nnr2Z$monai.networks.blocks.patchembeddingrZ&monai.networks.blocks.transformerblockr__all__Modulerr(r(r(r,s