o ( i%@sddlmZddlmZmZmZddlZddlmZddl mm Z ddl m Z ddlmZeddd\ZZGd d d ejZdS) ) annotations)OptionalTupleUnionN)get_rel_pos_embedding_layer)optional_importzeinops.layers.torch Rearrange)namecsJeZdZdZ             d$d%fdd Zd&d'd"d#ZZS)(SABlockz A self-attention block, based on: "Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale " FNT hidden_sizeint num_heads dropout_ratefloatqkv_biasbool save_attndim_head int | Nonehidden_input_sizecausalsequence_lengthrel_pos_embedding str | None input_size Tuple | Noneattention_dtypetorch.dtype | None include_fcuse_combined_linearuse_flash_attentionreturnNonec sxtd|krdkstdtd||dkr!td|r,|||_||_n||dkr6td||_|||_|rH| durHtd|rP|rPtd|rZ| durZtd||_|ra|n||_|| rrt|j|j|_ nt |_ |||||rtj|j|jd |d |_ t |_ |_ |_td d |d |_n-tj|j|j|d |_ tj|j|j|d |_ tj|j|j|d |_t |_ td |d|_td|_t||_t||_||_|jd|_||_t|_| |_||_| |_| |_||_||_ |r!| dur!|!dt"t#| | $dd| | |nt|_%| dur4t&| | |j|jnd|_'| |_(dS)a Args: hidden_size (int): dimension of hidden layer. num_heads (int): number of attention heads. dropout_rate (float, optional): fraction of the input units to drop. Defaults to 0.0. qkv_bias (bool, optional): bias term for the qkv linear layer. Defaults to False. save_attn (bool, optional): to make accessible the attention matrix. Defaults to False. dim_head (int, optional): dimension of each head. Defaults to hidden_size // num_heads. hidden_input_size (int, optional): dimension of the input tensor. Defaults to hidden_size. causal: whether to use causal attention (see https://arxiv.org/abs/1706.03762). sequence_length: if causal is True, it is necessary to specify the sequence length. rel_pos_embedding (str, optional): Add relative positional embeddings to the attention map. For now only "decomposed" is supported (see https://arxiv.org/abs/2112.01526). 2D and 3D are supported. input_size (tuple(spatial_dim), optional): Input resolution for calculating the relative positional parameter size. attention_dtype: cast attention operations to this dtype. include_fc: whether to include the final linear layer. Default to True. use_combined_linear: whether to use a single linear layer for qkv projection, default to True. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). rz'dropout_rate should be between 0 and 1.z-hidden size should be divisible by num_heads.Nz2sequence_length is necessary for causal attention.zsave_attn has been set to True, but use_flash_attention is also setto True. save_attn can only be used if use_flash_attention is False.z@rel_pos_embedding must be None if you are using flash_attention.)biaszb h (qkv l d) -> qkv b l h d)qkvlzb h (l d) -> b l h d)r(zb l h d -> b h (l d)g causal_mask))super__init__ ValueErrorZ inner_dimrrrnnLinearout_projIdentityr'to_qto_kto_vrinput_rearrange out_rearrangeDropout drop_output drop_weightsrscalertorchTensoratt_matrrrrr r!register_buffertrilonesviewr)rrel_positional_embeddingr)selfr rrrrrrrrrrrrr r! __class__e/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/selfattention.pyr+ s )                zSABlock.__init__ attn_maskOptional[torch.Tensor]c Cs|jr|||}|d|d|d}}}n|||}|||}|||}|jdurC||j}||j}|jrVt j |||||j |j |j d}nytd|||j }|jdurl||||}|j r|durwtd||jddddd|jdd|jdfdktd }|dur|dd}|d |jd d }||dktd }|jd d }|jr||_||}td ||}||}|jr| |}|!|}|S) aF Args: x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C attn_mask (torch.Tensor, optional): mask to apply to the attention matrix. B x (s_dim_1 * ... * s_dim_n). Defaults to None. Return: torch.Tensor: B x (s_dim_1 * ... * s_dim_n) x C rr$N)querykeyvaluerGr9 dropout_p is_causalzblxd,blyd->blxyz2Causal attention does not support attention masks.z-inf)dimzbhxy,bhyd->bhxd)"r r4r'r1r2r3rtor!Fscaled_dot_product_attentionr9rrr:einsumrAr, masked_fillr)shaper unsqueezeexpandrsoftmaxrdetachr<r8r5rr/r7)rBxrGoutputqkvr<rErErFforwardsP     >      zSABlock.forward) r FFNNFNNNNTTF) r r rr rrrrrrrrrrrrrrrrrrrrrrr rr!rr"r#)N)rGrH)__name__ __module__ __qualname____doc__r+ra __classcell__rErErCrFr s" |r ) __future__rtypingrrrr:torch.nnr-torch.nn.functional functionalrSmonai.networks.layers.utilsr monai.utilsrr_Moduler rErErErFs