首先是注意力公式：

其計(jì)算圖：

代碼：

def attention(query, key, value, mask=None, dropout=None): 
# query, key, value的形狀類似于(30, 8, 10, 64), (30, 8, 11, 64), 
#(30, 8, 11, 64)，例如30是batch.size，即當(dāng)前batch中有多少一個(gè)序列；
# 8=head.num，注意力頭的個(gè)數(shù)；
# 10=目標(biāo)序列中詞的個(gè)數(shù)，64是每個(gè)詞對(duì)應(yīng)的向量表示；
# 11=源語言序列傳過來的memory中，當(dāng)前序列的詞的個(gè)數(shù)，
# 64是每個(gè)詞對(duì)應(yīng)的向量表示。
# 類似于，這里假定query來自target language sequence；
# key和value都來自source language sequence.
  "Compute 'Scaled Dot Product Attention'" 
  d_k = query.size(-1) # 64=d_k
  scores = torch.matmul(query, key.transpose(-2, -1)) / 
    math.sqrt(d_k) # 先是(30,8,10,64)和(30, 8, 64, 11)相乘，
    #（注意是最后兩個(gè)維度相乘）得到(30,8,10,11)，
    #代表10個(gè)目標(biāo)語言序列中每個(gè)詞和11個(gè)源語言序列的分別的“親密度”。
    #然后除以sqrt(d_k)=8，防止過大的親密度。
    #這里的scores的shape是(30, 8, 10, 11)
  if mask is not None: 
    scores = scores.masked_fill(mask == 0, -1e9) 
    #使用mask，對(duì)已經(jīng)計(jì)算好的scores，按照mask矩陣，填-1e9，
    #然后在下一步計(jì)算softmax的時(shí)候，被設(shè)置成-1e9的數(shù)對(duì)應(yīng)的值~0,被忽視
  p_attn = F.softmax(scores, dim = -1) 
    #對(duì)scores的最后一個(gè)維度執(zhí)行softmax，得到的還是一個(gè)tensor, 
    #(30, 8, 10, 11)
  if dropout is not None: 
    p_attn = dropout(p_attn) #執(zhí)行一次dropout
  return torch.matmul(p_attn, value), p_attn
#返回的第一項(xiàng)，是(30,8,10, 11)乘以（最后兩個(gè)維度相乘）
#value=(30,8,11,64)，得到的tensor是(30,8,10,64)，
#和query的最初的形狀一樣。另外，返回p_attn，形狀為(30,8,10,11). 
#注意，這里返回p_attn主要是用來可視化顯示多頭注意力機(jī)制。

多頭注意力：

class MultiHeadedAttention(nn.Module): 
  def __init__(self, h, d_model, dropout=0.1): 
    # h=8, d_model=512
    "Take in model size and number of heads." 
    super(MultiHeadedAttention, self).__init__() 
    assert d_model % h == 0 # We assume d_v always equals d_k 512%8=0
    self.d_k = d_model // h # d_k=512//8=64
    self.h = h #8
    self.linears = clones(nn.Linear(d_model, d_model), 4) 
    #定義四個(gè)Linear networks, 每個(gè)的大小是(512, 512)的，
    #每個(gè)Linear network里面有兩類可訓(xùn)練參數(shù)，Weights，
    #其大小為512*512，以及biases，其大小為512=d_model。

    self.attn = None 
    self.dropout = nn.Dropout(p=dropout)
  def forward(self, query, key, value, mask=None): 
   # 注意，輸入query的形狀類似于(30, 10, 512)，
   # key.size() ~ (30, 11, 512), 
   #以及value.size() ~ (30, 11, 512)
    
    if mask is not None: # Same mask applied to all h heads. 
      mask = mask.unsqueeze(1) # mask下回細(xì)細(xì)分解。
    nbatches = query.size(0) #e.g., nbatches=30
    # 1) Do all the linear projections in batch from 
    #d_model => h x d_k 
    query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k)
      .transpose(1, 2) for l, x in 
      zip(self.linears, (query, key, value))] 
      # 這里是前三個(gè)Linear Networks的具體應(yīng)用，
      #例如query=(30,10, 512) -> Linear network -> (30, 10, 512) 
      #-> view -> (30,10, 8, 64) -> transpose(1,2) -> (30, 8, 10, 64)
      #，其他的key和value也是類似地，
      #從(30, 11, 512) -> (30, 8, 11, 64)。
    # 2) Apply attention on all the projected vectors in batch. 
    x, self.attn = attention(query, key, value, mask=mask, 
      dropout=self.dropout) 
      #調(diào)用上面定義好的attention函數(shù)，輸出的x形狀為(30, 8, 10, 64)；
      #attn的形狀為(30, 8, 10=target.seq.len, 11=src.seq.len)
    # 3) "Concat" using a view and apply a final linear. 
    x = x.transpose(1, 2).contiguous().
      view(nbatches, -1, self.h * self.d_k) 
      # x ~ (30, 8, 10, 64) -> transpose(1,2) -> 
      #(30, 10, 8, 64) -> contiguous() and view -> 
      #(30, 10, 8*64) = (30, 10, 512)
return self.linears[-1](x) 
#執(zhí)行第四個(gè)Linear network，把(30, 10, 512)經(jīng)過一次linear network，
#得到(30, 10, 512).

摘自：https://zhuanlan.zhihu.com/p/107889011

總結(jié)

以上是生活随笔為你收集整理的transformer中自注意力和多头注意力的pytorch实现的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。