o ) iW@sVddlmZddlmZddlZddlmmZddlmZdd d Z dddZ dS)) annotations)TupleN)nnq_sizeintk_sizerel_pos torch.TensorreturncCst}tdt||d}|jd|kr6tj|d|jddddd|dd}|d|dd}n|}t |dddft||d}t |dddft||d}|||dt||d}|| S) aY Get relative positional embeddings according to the relative positions of query and key sizes. Args: q_size (int): size of query q. k_size (int): size of key k. rel_pos (Tensor): relative position embeddings (L, C). Returns: Extracted positional embeddings according to relative positions. rlinear)sizemodeNg?) torchTensorrmaxshapeF interpolatereshapepermutearangelong)rrrZrel_pos_resizedZ max_rel_distZq_coordsZk_coordsrelative_coordsrg/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/blocks/attention_utils.py get_rel_poss  $$ rattnq rel_pos_lstnn.ParameterListrc Cst|d|d|d}t|d|d|d}|j\}}} t|dkr|dd\} } |dd\} } ||| | | }td||}td||}||| | | | |dddddddddf|dddddddddf|| | | | }|St|dkr |dd\} } }|dd\} } }t|||d}||| | || }td||}td ||}td ||}||| | || | ||ddddddddddf|ddddddddddf|ddddddddddf|| | || | |}|S) a Calculate decomposed Relative Positional Embeddings from mvitv2 implementation: https://github.com/facebookresearch/mvit/blob/19786631e330df9f3622e5402b4a419a263a2c80/mvit/models/attention.py Only 2D and 3D are supported. Encoding the relative position of tokens in the attention matrix: tokens spaced a distance `d` apart will have the same embedding value (unlike absolute positional embedding). .. math:: Attn_{logits}(Q, K) = (QK^{T} + E_{rel})*scale where .. math:: E_{ij}^{(rel)} = Q_{i}.R_{p(i), p(j)} with :math:`R_{p(i), p(j)} \in R^{dim}` and :math:`p(i), p(j)`, respectively spatial positions of element :math:`i` and :math:`j` When using "decomposed" relative positional embedding, positional embedding is defined ("decomposed") as follow: .. math:: R_{p(i), p(j)} = R^{d1}_{d1(i), d1(j)} + ... + R^{dn}_{dn(i), dn(j)} with :math:`n = 1...dim` Decomposed relative positional embedding reduces the complexity from :math:`\mathcal{O}(d1*...*dn)` to :math:`\mathcal{O}(d1+...+dn)` compared with classical relative positional embedding. Args: attn (Tensor): attention map. q (Tensor): query q in the attention layer with shape (B, s_dim_1 * ... * s_dim_n, C). rel_pos_lst (ParameterList): relative position embeddings for each axis: rel_pos_lst[n] for nth axis. q_size (Tuple): spatial sequence size of query q with (q_dim_1, ..., q_dim_n). k_size (Tuple): spatial sequence size of key k with (k_dim_1, ..., k_dim_n). Returns: attn (Tensor): attention logits with added relative positional embeddings. rr r Nzbhwc,hkc->bhwkzbhwc,wkc->bhwkzbhwdc,hkc->bhwdkzbhwdc,wkc->bhwdk)rrlenrreinsumview)rr r!rrrhrwbatch_dimq_hq_wk_hk_wr_qZrel_hZrel_wZq_dk_drdZrel_drrradd_decomposed_rel_pos4s>+  V"""r3)rrrrrr r r ) rr r r r!r"rrrrr r ) __future__rtypingrrtorch.nn.functionalr functionalrrr3rrrrs   !