o

    h!                     @   s.  d dl Z d dlZd dlZd dlZd dlmZ d dlm  mZ	 ddl
mZ d dlm
Z
 ddlmZ ddlmZ ddlmZ d d	lmZ eejed
krZd dlmZ eejdd
ZnejZG dd deZG dd deZG dd deZG dd deZG dd deZG dd deZ G dd deZ!dS )    N   )
build_dropout)
trunc_normal_   )
BaseModule)
LayerScale)	to_2tuple)
digit_versionz1.10.0)partialij)indexingc                       sN   e Zd ZdZ					d
 fdd	Z fddZdd	d
Zedd Z  Z	S )	WindowMSAa  Window based multi-head self-attention (W-MSA) module with relative
    position bias.

    Args:
        embed_dims (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
            Defaults to True.
        qk_scale (float, optional): Override default qk scale of
            ``head_dim ** -0.5`` if set. Defaults to None.
        attn_drop (float, optional): Dropout ratio of attention weight.
            Defaults to 0.
        proj_drop (float, optional): Dropout ratio of output. Defaults to 0.
        init_cfg (dict, optional): The extra config for initialization.
            Defaults to None.
    TN        c	                    s   t  | || _|| _|| _|| }	|p|	d | _tt	d|d  d d|d  d  || _
| j\}
}| d| d |
d|}||j }
|

d }
| d|
 tj||d |d| _t|| _t||| _t|| _tjdd	| _d S )
N      r   r   r   relative_position_index   biasdim)super__init__
embed_dimswindow_size	num_headsscalenn	Parametertorchzerosrelative_position_bias_tabledouble_step_seqTflip
contiguousregister_bufferLinearqkvDropout	attn_dropproj	proj_dropSoftmaxsoftmax)selfr   r   r   qkv_biasqk_scaler*   r,   init_cfghead_embed_dimsWhWwZrel_index_coordsZrel_position_index	__class__ 3/root/Awesome-Backbones/configs/common/attention.pyr   +   s*   
"

zWindowMSA.__init__c                    s    t t|   t| jdd d S Ng{Gz?)std)r   r
   init_weightsr   r!   r/   r6   r8   r9   r<   O   s   zWindowMSA.init_weightsc           
      C   sj  |j \}}}| |||d| j|| j ddddd}|d |d |d }}}	|| j }||dd }
| j| j	d 	| j
d | j
d  | j
d | j
d  d}|ddd }|
|d }
|dur|j d }|
	|| || j|||dd }
|
	d| j||}
| 
|
}
n| 
|
}
| |
}
|
|	 dd|||}| |}| |}|S )	
        Args:

            x (tensor): input features with shape of (num_windows*B, N, C)
            mask (tensor, Optional): mask with shape of (num_windows, Wh*Ww,
                Wh*Ww), value should be between (-inf, 0].
        r   r   r   r      r   N)shaper(   reshaper   permuter   	transposer!   r   viewr   r%   	unsqueezer.   r*   r+   r,   )
r/   xmaskB_NCr(   qkvattnrelative_position_biasnWr8   r8   r9   forwardT   sR   






zWindowMSA.forwardc                 C   sL   t d| | | }t d|| |}|d d d f |d d d f  ddS )Nr   r   r   )r   arangerB   )step1len1step2len2seq1seq2r8   r8   r9   r"   ~   s   (zWindowMSA.double_step_seq)TNr   r   NN)
__name__
__module____qualname____doc__r   r<   rR   staticmethodr"   
__classcell__r8   r8   r6   r9   r
      s    $
*r
   c                       sr   e Zd ZdZdddddeddddddee df fd	d
	Zdd Zed
d Z	edd Z
edddZ  ZS )ShiftWindowMSAa  Shift Window Multihead Self-Attention Module.

    Args:
        embed_dims (int): Number of input channels.
        num_heads (int): Number of attention heads.
        window_size (int): The height and width of the window.
        shift_size (int, optional): The shift step of each window towards
            right-bottom. If zero, act as regular window-msa. Defaults to 0.
        qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
            Defaults to True
        qk_scale (float | None, optional): Override default qk scale of
            head_dim ** -0.5 if set. Defaults to None.
        attn_drop (float, optional): Dropout ratio of attention weight.
            Defaults to 0.0.
        proj_drop (float, optional): Dropout ratio of output. Defaults to 0.
        dropout_layer (dict, optional): The dropout_layer used before output.
            Defaults to dict(type='DropPath', drop_prob=0.).
        pad_small_map (bool): If True, pad the small feature map to the window
            size, which is common used in detection and segmentation. If False,
            avoid shifting window and shrink the window size to the size of
            feature map, which is common used in classification.
            Defaults to False.
        version (str, optional): Version of implementation of Swin
            Transformers. Defaults to `v1`.
        init_cfg (dict, optional): The extra config for initialization.
            Defaults to None.
    r   TNDropPathr   type	drop_probFc              
      s   t  | |d us|d urtdt || _|| _d| j  kr'| jk s*J  J t|
ts>J dt	t dt	|
 d|
d|t
| j|||||d|| _t|	| _
|
| _d S )NzThe ShiftWindowMSA in new version has supported auto padding and dynamic input shape in all condition. And the argument `auto_pad` and `input_resolution` have been deprecated.r   z?Expect Window based multi-head self-attention Module is type ofz
, but got .)r   r   r   r0   r1   r*   r,   r8   )r   r   warningswarnDeprecationWarning
shift_sizer   
issubclassr   rd   r   w_msar   drop
pad_small_map)r/   r   r   r   rj   r0   r1   r*   r,   
dropout_layerrn   Zinput_resolutionZauto_padZ
window_msaZmsa_cfgr2   r6   r8   r9   r      s:    

zShiftWindowMSA.__init__c              	   C   s  |j \}}}|\}}||| ksJ d| d| d| d|||||}| j}| j}	t|||kr;| jr8|	nd}	nt|||k rS| jsSJ d| d| d| d|||  | }
|||  | }t|ddd|
d|f}|j d	 |j d
 }}
|	dkrtj	||	 |	 fdd}| j
||
f||	|jd
}| ||}|d|d
 |}| j
||d}|d|||}| |||
|}| jdkrtj	||	|	fdd}n|}||ks||
kr|d d d |d |d d f  }|||| |}| |}|S )NzThe query length z  doesn't match the input shape (z, z).r   zThe input shape (z#) is smaller than the window size (zB). Please set `pad_small_map=True`, or decrease the `window_size`.r   r   )r   r   )shiftsdims)r   rj   devicer   )rH   )rA   rE   r   rj   minrn   Fpadr   roll
get_attn_maskrr   window_partitionrl   window_reverser%   rm   )r/   queryhw_shapeBLrK   HWr   rj   pad_rpad_bZH_padZW_pad	attn_maskZ
query_windowsZattn_windowsZ	shifted_xrG   r8   r8   r9   rR      s`   




$
zShiftWindowMSA.forwardc                 C   sb   t | jd || | |  }| ||| || ||d}|dddddd |||d}|S )Nr   r   r   r   r   r?      )intrA   rE   rC   r%   )windowsr~   r   r   r|   rG   r8   r8   r9   ry     s   $zShiftWindowMSA.window_reversec                 C   sV   | j \}}}}| ||| ||| ||} | dddddd }|d|||}|S )Nr   r   r   r   r?   r   r   )rA   rE   rC   r%   )rG   r   r|   r~   r   rK   r   r8   r8   r9   rx     s   zShiftWindowMSA.window_partitionc              	   C   s   |dkr|t jdg| dR d|i}td| t| | t| d f}td| t| | t| d f}d}|D ]}|D ]}	||d d ||	d d f< |d7 }q>q:t||}
|
d|| }
|
d|
d }||dkd}||dkd}|S d }|S )Nr   r   rr   r   r   g      Yr   )r   r    slicera   rx   rE   rF   masked_fill)r{   r   rj   rr   Zimg_maskh_slicesw_slicescnthwZmask_windowsr   r8   r8   r9   rw   (  s8   


zShiftWindowMSA.get_attn_maskrZ   )
r[   r\   r]   r^   dictr
   r   rR   r_   ry   rx   rw   r`   r8   r8   r6   r9   ra      s,     
.G

ra   c                
       sF   e Zd ZdZdddedddddddddf
 fdd		Zd
d Z  ZS )MultiheadAttentiona  Multi-head Attention Module.

    This module implements multi-head attention that supports different input
    dims and embed dims. And it also supports a shortcut from ``value``, which
    is useful if input dims is not the same with embed dims.

    Args:
        embed_dims (int): The embedding dimension.
        num_heads (int): Parallel attention heads.
        input_dims (int, optional): The input dimension, and if None,
            use ``embed_dims``. Defaults to None.
        attn_drop (float): Dropout rate of the dropout layer after the
            attention calculation of query and key. Defaults to 0.
        proj_drop (float): Dropout rate of the dropout layer after the
            output projection. Defaults to 0.
        dropout_layer (dict): The dropout config before adding the shortcut.
            Defaults to ``dict(type='Dropout', drop_prob=0.)``.
        qkv_bias (bool): If True, add a learnable bias to q, k, v.
            Defaults to True.
        qk_scale (float, optional): Override default qk scale of
            ``head_dim ** -0.5`` if set. Defaults to None.
        proj_bias (bool) If True, add a learnable bias to output projection.
            Defaults to True.
        v_shortcut (bool): Add a shortcut from value to output. It's usually
            used if ``input_dims`` is different from ``embed_dims``.
            Defaults to False.
        init_cfg (dict, optional): The Config for initialization.
            Defaults to None.
    Nr   r)   rc   TFc
           
         s   t t| j|d |p|| _|| _|| _|
| _|| | _|p"| jd | _t	j
| j|d |d| _t	|| _
t	j
|||	d| _t	|| _t|| _|rSt|| _d S t	 | _d S )Nr2   r   r   r   )r   r   r   
input_dimsr   r   
v_shortcut	head_dimsr   r   r'   r(   r)   r*   r+   r,   r   out_dropr   gamma1Identity)
r/   r   r   r   r*   r,   ro   r0   r1   	proj_biasr   Zuse_layer_scaler2   r6   r8   r9   r   c  s   


zMultiheadAttention.__init__c           
      C   s   |j \}}}| |||d| j| jddddd}|d |d |d }}}||dd | j }	|	jdd}	| 	|	}	|	| dd||| j
}| |}| | 
| |}| jrg|d| }|S )	Nr   r   r   r   r?   r@   r   r   )rA   r(   rB   r   r   rC   rD   r   r.   r*   r   r+   r   r   r,   r   squeeze
r/   rG   r|   rJ   _r(   rL   rM   rN   rO   r8   r8   r9   rR     s    

zMultiheadAttention.forwardr[   r\   r]   r^   r   r   rR   r`   r8   r8   r6   r9   r   D  s    !
#r   c                       s8   e Zd ZdZ						d fdd	Zdd	d
Z  ZS )
WindowMSAV2a  Window based multi-head self-attention (W-MSA) module with relative
    position bias.

    Based on implementation on Swin Transformer V2 original repo. Refers to
    https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer_v2.py
    for more details.

    Args:
        embed_dims (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional): If True, add a learnable bias to q, k, v.
            Defaults to True.
        attn_drop (float, optional): Dropout ratio of attention weight.
            Defaults to 0.
        proj_drop (float, optional): Dropout ratio of output. Defaults to 0.
        pretrained_window_size (tuple(int)): The height and width of the window
            in pre-training.
        init_cfg (dict, optional): The extra config for initialization.
            Defaults to None.
    Tr      r   r   Nc
              	      s~  t  |	 || _|| _|| _ttjd|ddtjddtj||dd| _	tj
tdt
|ddf dd| _tj| jd	 d  | jd	 tjd
}tj| jd d  | jd tjd
}tt||gddd	 d	}
|d	 d	kr|
d d d d d d d	f  |d	 d   < |
d d d d d d df  |d d   < n.|
d d d d d d d	f  | jd	 d   < |
d d d d d d df  | jd d   < |
d9 }
t|
tt|
d  td }
| d
|
 t| jd	 }t| jd }tjt||gd	d}tj|dd}|d d d d d f |d d d d d f  }|ddd	 }|d d d d d	f  | jd	 d 7  < |d d d d df  | jd d 7  < |d d d d d	f  d| jd  d 9  < |d}| d| tj||d dd| _|rt
t|| _t
t|| _ nd | _d | _ t!|| _"t||| _#t!|| _$tj%dd| _&d S )Nr   T)in_featuresout_featuresr   )inplaceF
   r   
requires_gradr   )dtype   g      ?relative_coords_tabler   )	start_dimr   r   r   r   )'r   r   r   r   r   r   
Sequentialr'   ReLUcpb_mlpr   r   logoneslogit_scalerS   float32stackmeshgridrC   r%   rF   signlog2absnpr&   flattensumr(   r    q_biasv_biasr)   r*   r+   r,   r-   r.   )r/   r   r   r   r0   r*   r,   Zcpb_mlp_hidden_dimsZpretrained_window_sizer2   kwargsrelative_coords_hrelative_coords_wr   Z	indexes_hZ	indexes_wcoordinatesZrelative_coordinatesr   r6   r8   r9   r     s   




..(,
zWindowMSAV2.__init__c                 C   s  |j \}}}d}| jdurt| jtj| jdd| jf}tj|| jj	|d}|
||d| j|| j ddddd	}|d |d |d }}	}
tj
|d
dtj
|	d
ddd
 }tj| jtd
d }|| }| | jd
| j}
|
| jd
 | jd | jd  | jd | jd  d
}|ddd }dt| }||d }|dur|j d }||| || j|||dd }|d
| j||}| |}n| |}| |}||
 dd
|||}| |}| |}|S )r>   NFr   inputweightr   r   r   r   r   r?   r   r   r@   g      Y@)max   ) rA   r   r   cat
zeros_liker   rt   linearr(   r   rB   r   rC   	normalizerD   clampr   r   r   expr   r   rE   r   r   r%   sigmoidrF   r.   r*   r+   r,   )r/   rG   rH   rI   rJ   rK   r0   r(   rL   rM   rN   rO   r   r!   rP   rQ   r8   r8   r9   rR   	  s   
"






zWindowMSAV2.forward)Tr   r   r   r   NrZ   )r[   r\   r]   r^   r   rR   r`   r8   r8   r6   r9   r     s    Zr   c                       sR   e Zd ZdZ					d fdd	Zdd Zd	d
 Z fddZdd
dZ  Z	S )
BEiTAttentiona  Window based multi-head self-attention (W-MSA) module with relative
    position bias.

    The initial implementation is in MMSegmentation.

    Args:
        embed_dims (int): Number of input channels.
        num_heads (int): Number of attention heads.
        window_size (tuple[int]): The height and width of the window.
        use_rel_pos_bias (bool): Whether to use unique relative position bias,
            if False, use shared relative position bias defined in backbone.
        bias (str): The option to add leanable bias for q, k, v. If bias is
            True, it will add leanable bias. If bias is 'qv_bias', it will only
            add leanable bias for q, v. If bias is False, it will not add bias
            for q, k, v. Default to 'qv_bias'.
        qk_scale (float | None, optional): Override default qk scale of
            head_dim ** -0.5 if set. Default: None.
        attn_drop_rate (float): Dropout ratio of attention weight.
            Default: 0.0
        proj_drop_rate (float): Dropout ratio of output. Default: 0.
        init_cfg (dict | None, optional): The Config for initialization.
            Default: None.
    qv_biasNr   c
           
         s   t  j|	d || _|| _|| }|| _|p|d | _|}|dkr'|   d}|| _|| _| 	  t
j||d |d| _t

|| _t
||| _t

|| _d S )Nr   r   r   Fr   r   )r   r   r   r   r   r   
_init_qv_biasr   use_rel_pos_bias_init_rel_pos_embeddingr   r'   r(   r)   r*   r+   r,   )
r/   r   r   r   r   r   r1   Zattn_drop_rateZproj_drop_rater2   r   r3   r0   r6   r8   r9   r   Z  s"   zBEiTAttention.__init__c                 C   s,   t t| j| _t t| j| _d S rZ   )r   r   r   r    r   r   r   r=   r8   r8   r9   r   z  s   zBEiTAttention._init_qv_biasc           	      C   s  | j r| j\}}d| d d| d  d | _tt| j| j| _t	|}t	|}t
t||g}t|d}|d d d d d f |d d d d d f  }|
ddd }|d d d d df  |d 7  < |d d d d df  |d 7  < |d d d d df  d| d 9  < tj|| d fd |jd}|d|dd dd f< | jd |ddd f< | jd |dd df< | jd |d< | d| d S d | _d | _d | _d S )	Nr   r   r   r   )sizer   r   r   r   )r   r   Znum_relative_distancer   r   r   r    r   r!   rS   r   torch_meshgridr   rC   r%   r   r   r&   r   )	r/   r4   r5   coords_hcoords_wcoordscoords_flattenrelative_coordsr   r8   r8   r9   r   ~  s:   


*""&
z%BEiTAttention._init_rel_pos_embeddingc                    s&   t    | jrt| jdd d S d S r:   )r   r<   r   r   r!   r=   r6   r8   r9   r<     s   
zBEiTAttention.init_weightsc                 C   st  |j \}}}| jdkr(tj| jdd}t| j|| jf}tj|| j	j
|d}n| 	|}|||d| jd
dddd	d
}|d |d	 |d }	}
}|	| j }	|	|
dd }| jdur| jd }
| jd	 }| j| jd |
| d	 |
| d	 d}|
ddd	 }||d }|dur|| }|jdd
}| |}|| d	d|||}| |}| |}|S )z
        Args:
            x (tensor): input features with shape of (num_windows*B, N, C).
            rel_pos_bias (tensor): input relative position bias with shape of
                (num_heads, N, N).
        r   Fr   r   r   r   r   r   r   r?   r@   Nr   )rA   r   r   r   r   r   r   rt   r   r(   r   rB   r   rC   r   rD   r!   r   r   rE   r%   rF   r.   r*   r+   r,   )r/   rG   Zrel_pos_biasr|   rJ   rK   Zk_biasr0   r(   rL   rM   rN   rO   r4   r5   rP   r8   r8   r9   rR     sD   

"







zBEiTAttention.forward)r   Nr   r   NrZ   )
r[   r\   r]   r^   r   r   r   r<   rR   r`   r8   r8   r6   r9   r   A  s     (r   c                       sH   e Zd ZdZ				d
 fdd	Ze d fd	d
	Zdd Z  Z	S )LeAttentionuC  LeViT Attention. Multi-head attention with attention bias,  which is
    proposed in `LeViT: a Vision Transformer in ConvNet’s Clothing for Faster
    Inference<https://arxiv.org/abs/2104.01136>`_

    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads. Default: 8.
        key_dim (int): Dimension of key. Default: None.
        attn_ratio (int): Ratio of attention heads. Default: 8.
        resolution (tuple[int]): Input resolution. Default: (16, 16).
        init_cfg (dict, optional): The Config for initialization.
    r   r?      r   Nc                    st  t  j|d t|trt|dksJ || _|d | _|| _||  | _}t	|| | _
t	|| | | _|| _| j|d  }t
|| _t
||| _t
| j|| _ttt|d t|d }	t|	}
i }g }|	D ],}
|	D ]'}t|
d |d  t|
d |d  f}||vrt|||< |||  qsqotj
t|t|| _| jdt||
|
dd d S )	Nr   r   r   r   r   attention_bias_idxsF)
persistent) r   r   
isinstancetuplelenr   r   key_dimnh_kdr   ddh
attn_ratior   	LayerNormnormr'   r(   r+   list	itertoolsproductranger   appendr   r   r    attention_biasesr&   
LongTensorrE   )r/   r   r   r   r   
resolutionr2   r   r   pointsrJ   Zattention_offsetsidxsp1p2offsetr6   r8   r9   r     sD   
(
zLeAttention.__init__Tc                    s<   t  | |rt| dr| `d S | jd d | jf | _d S )Nab)r   trainhasattrr   r   r   )r/   moder6   r8   r9   r     s   zLeAttention.trainc           
      C   s   |j \}}}| |}| |}|||| jdj| j| j| jgdd\}}}|dddd}|dddd}|dddd}||	dd | j
 | jrV| jd d | j
f n| j }	|	jdd}	|	| 	dd||| j}| |}|S )Nr   r   r   r   r   r   r@   )rA   r   r(   rE   r   splitr   r   rC   rD   r   trainingr   r   r   r.   rB   r   r+   r   r8   r8   r9   rR     s*   


zLeAttention.forward)r   r?   r   N)T)
r[   r\   r]   r^   r   r   no_gradr   rR   r`   r8   r8   r6   r9   r     s    )r   c                       sH   e Zd ZdZddddedddddd	dddf fd
d	Zdd
 Z  ZS )ChannelMultiheadAttentiona  Channel Multihead Self-attention Module.

    This module implements channel multi-head attention that supports different
    input dims and embed dims.
    Args:
        embed_dims (int): The embedding dimension.
        num_heads (int): Parallel attention heads.
        input_dims (int, optional): The input dimension, and if None,
            use ``embed_dims``. Defaults to None.
        attn_drop (float): Dropout rate of the dropout layer after the
            attention calculation of query and key. Defaults to 0.
        proj_drop (float): Dropout rate of the dropout layer after the
            output projection. Defaults to 0.
        dropout_layer (dict): The dropout config before adding the shoutcut.
            Defaults to ``dict(type='Dropout', drop_prob=0.)``.
        qkv_bias (bool): If True, add a learnable bias to q, k, v.
            Defaults to False.
        proj_bias (bool) If True, add a learnable bias to output projection.
            Defaults to True.
        qk_scale_type (str): The scale type of qk scale.
            Defaults to 'learnable'. It can be 'learnable', 'fixed' or 'none'.
        qk_scale (float, optional): If set qk_scale_type to 'none', this
            should be specified with valid float number. Defaults to None.
        v_shortcut (bool): Add a shortcut from value to output. It's usually
            used if ``input_dims`` is different from ``embed_dims``.
            Defaults to False.
        init_cfg (dict, optional): The Config for initialization.
            Defaults to None.
    r   Nr   r)   rc   FT	learnablec
           
         s   t  | |p	|| _|| _|| _|| _|| | _|	dkr)tt	
|dd| _n|	dkr4| jd | _n
|	dkrA|
d us>J |
| _tj| j|d |d| _
t|| _tj|||d| _t|| _t|| _d S )Nr   r   fixedr   noner   r   )r   r   r   r   r   r   r   r   r   r   r   r   r'   r(   r)   r*   r+   r,   r   r   )
r/   r   r   r   r*   r,   ro   r0   r   Z
qk_scale_typer1   r   r2   r6   r8   r9   r   J  s$   

z"ChannelMultiheadAttention.__init__c           
      C   s   |j \}}}| |||d| j| jddddd}dd |d |d |d fD \}}}tj|dd	tj|dd	}}||d
d | j	 }	|	j
dd	}	|	| dddd||| j}| |}| 
| |}| jrw|d d| }|S )Nr   r   r   r   r?   c                 S   s   g | ]}| d dqS )r@   r   )rD   ).0itemr8   r8   r9   
<listcomp>s  s    z5ChannelMultiheadAttention.forward.<locals>.<listcomp>r   r   r@   )rA   r(   rB   r   r   rC   rt   r   rD   r   r.   r   r+   r   r,   r   r   r   r8   r8   r9   rR   n  s    & 
z!ChannelMultiheadAttention.forwardr   r8   r8   r6   r9   r   +  s     
$r   )"rg   r   numpyr   r   torch.nnr   Ztorch.nn.functional
functionalrt   Zbasic.build_layerr   Zcore.initialize.weight_initr   base_moduler   layer_scaler   helpersr   Zutils.version_utilsr	   __version__	functoolsr
   r   r   r
   ra   r   r   r   r   r   r8   r8   r8   r9   <module>   s4   m @T * W