首先是注意力公式：
其計算圖：
代碼：
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，即當前batch中有多少一個序列； # 8=head.num，注意力頭的個數； # 10=目標序列中詞的個數，64是每個詞對應的向量表示； # 11=源語言序列傳過來的memory中，當前序列的詞的個數， # 64是每個詞對應的向量表示。 # 類似於，這里假定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)相乘， #（注意是最后兩個維度相乘）得到(30,8,10,11)， #代表10個目標語言序列中每個詞和11個源語言序列的分別的“親密度”。 #然后除以sqrt(d_k)=8，防止過大的親密度。 #這里的scores的shape是(30, 8, 10, 11) if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) #使用mask，對已經計算好的scores，按照mask矩陣，填-1e9， #然后在下一步計算softmax的時候，被設置成-1e9的數對應的值~0,被忽視 p_attn = F.softmax(scores, dim = -1) #對scores的最后一個維度執行softmax，得到的還是一個tensor, #(30, 8, 10, 11) if dropout is not None: p_attn = dropout(p_attn) #執行一次dropout return torch.matmul(p_attn, value), p_attn #返回的第一項，是(30,8,10, 11)乘以（最后兩個維度相乘） #value=(30,8,11,64)，得到的tensor是(30,8,10,64)， #和query的最初的形狀一樣。另外，返回p_attn，形狀為(30,8,10,11). #注意，這里返回p_attn主要是用來可視化顯示多頭注意力機制。
多頭注意力：
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) #定義四個Linear networks, 每個的大小是(512, 512)的， #每個Linear network里面有兩類可訓練參數，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下回細細分解。 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))] # 這里是前三個Linear Networks的具體應用， #例如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) #調用上面定義好的attention函數，輸出的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) #執行第四個Linear network，把(30, 10, 512)經過一次linear network， #得到(30, 10, 512).