cnn中實(shí)現(xiàn)attention主要是有Sparial Domain和Channel Domain
soft-attention是可微的，可以通過(guò)梯度來(lái)實(shí)現(xiàn)

import torch
import torch.nn as nn
import numpy as np
import math

class SelfAttention(nn.Module):
    
    def __init__(self, hidden_size, num_attention_heads, dropout_prob):   
        """
        假設(shè) hidden_size = 128, num_attention_heads = 8, dropout_prob = 0.2
        即隱層維度為128，注意力頭設(shè)置為8個(gè)
        """
        super(SelfAttention, self).__init__()
        if hidden_size % num_attention_heads != 0:   # 整除
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (hidden_size, num_attention_heads))
        # 參數(shù)定義
        self.num_attention_heads = num_attention_heads    # 8
        self.attention_head_size = int(hidden_size / num_attention_heads)  # 16  每個(gè)注意力頭的維度
        self.all_head_size = int(self.num_attention_heads * self.attention_head_size)   
        # all_head_size = 128 即等于hidden_size, 一般自注意力輸入輸出前后維度不變
        
        # query, key, value 的線(xiàn)性變換（上述公式2）
        self.query = nn.Linear(hidden_size, self.all_head_size)    # 128, 128
        self.key = nn.Linear(hidden_size, self.all_head_size)
        self.value = nn.Linear(hidden_size, self.all_head_size)
        
        # dropout
        self.dropout = nn.Dropout(dropout_prob)

    def transpose_for_scores(self, x):
        # INPUT:  x'shape = [bs, seqlen, hid_size]  假設(shè)hid_size=128
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) # [bs, seqlen, 8, 16]
        x = x.view(*new_x_shape)   # 
        return x.permute(0, 2, 1, 3)   # [bs, 8, seqlen, 16]

    def forward(self, hidden_states, attention_mask):
        # eg: attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])  shape=[bs, seqlen]
        attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)   # [bs, 1, 1, seqlen] 增加維度
        attention_mask = (1.0 - attention_mask) * -10000.0   # padding的token置為-10000，exp(-1w)=0
        
        # 線(xiàn)性變換
        mixed_query_layer = self.query(hidden_states)   # [bs, seqlen, hid_size]
        mixed_key_layer = self.key(hidden_states)       # [bs, seqlen, hid_size]
        mixed_value_layer = self.value(hidden_states)   # [bs, seqlen, hid_size]

        query_layer = self.transpose_for_scores(mixed_query_layer)    # [bs, 8, seqlen, 16]
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)   # [bs, 8, seqlen, 16]

        # Take the dot product between "query" and "key" to get the raw attention scores.
        # 計(jì)算query與title之間的點(diǎn)積注意力分?jǐn)?shù)，還不是權(quán)重（個(gè)人認(rèn)為權(quán)重應(yīng)該是和為1的概率分布）
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        # [bs, 8, seqlen, 16]*[bs, 8, 16, seqlen]  ==> [bs, 8, seqlen, seqlen]
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)   # [bs, 8, seqlen, seqlen]
        # 除以根號(hào)注意力頭的數(shù)量，可看原論文公式，防止分?jǐn)?shù)過(guò)大，過(guò)大會(huì)導(dǎo)致softmax之后非0即1
        attention_scores = attention_scores + attention_mask
        # 加上mask，將padding所在的表示直接-10000

        # 將注意力轉(zhuǎn)化為概率分布，即注意力權(quán)重
        attention_probs = nn.Softmax(dim=-1)(attention_scores)    # [bs, 8, seqlen, seqlen]

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)
        
        # 矩陣相乘，[bs, 8, seqlen, seqlen]*[bs, 8, seqlen, 16] = [bs, 8, seqlen, 16]
        context_layer = torch.matmul(attention_probs, value_layer)   # [bs, 8, seqlen, 16]
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()   # [bs, seqlen, 8, 16]
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)   # [bs, seqlen, 128]
        context_layer = context_layer.view(*new_context_layer_shape)
        return context_layer    # [bs, seqlen, 128] 得到輸出

attention=SelfAttention(4,2,0.2)
x_in=torch.randn(3,5,4)
x_mask=torch.Tensor([[1,1,1,0,0],
                    [1,1,0,0,0],
                    [1,1,1,1,1],])
print(x_mask.shape)
x_out=attention(x_in,x_mask)