attention

multi-head attention

pytorch代码实现

import math
import torch
from torch import nn
import torch.nn.functional as Fclass MultiHeadAttention(nn.Module):def __init__(self, heads=8, d_model=128, droput=0.1):super().__init__()self.d_model = d_model  # 128self.d_k = d_model // heads  # 128//8=16self.h = heads  # 8self.q_linear = nn.Linear(d_model, d_model)  # (50,128)*(128,128)=(50,128)，其中(128*128)属于权重，在网络训练中学习。self.k_linear = nn.Linear(d_model, d_model)self.v_linear = nn.Linear(d_model, d_model)self.dropout = nn.Dropout(droput)self.out = nn.Linear(d_model, d_model)def attention(self, q, k, v, d_k, mask=None, dropout=None):scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # 矩阵乘法 (32,8,50,16)*(32,8,16,50)->(32,8,50,50)if mask is not None:mask = mask.unsqueeze(1)scores = scores.masked_fill(mask == 0, -1e9)scores = F.softmax(scores, dim=-1)if dropout is not None:scores = dropout(scores)output = torch.matmul(scores, v)  # (32,8,50,50)*(32,8,50,16)->(32,8,50,16)return outputdef forward(self, q, k, v, mask=None):bs = q.size(0)  # batch_size 大小  这里的例子是32k = self.k_linear(k).view(bs, -1, self.h, self.d_k)q = self.k_linear(q).view(bs, -1, self.h, self.d_k)v = self.k_linear(v).view(bs, -1, self.h, self.d_k)# (32,50,128)->(32,50,128)->(32,50,8,16)  8*16=128 每个embedding拆成的8份，也就是8个头k = k.transpose(1, 2)  # (32,50,8,16)->(32,8,50,16)q = q.transpose(1, 2)v = v.transpose(1, 2)scores = self.attention(q, k, v, self.d_k, mask, self.dropout)  # (32,8,50,16)concat = scores.transpose(1, 2).contiguous().view(bs, -1, self.d_model)  # (32,50,128)output = self.out(concat)  # (32,50,128)return outputif __name__ == '__main__':multi_head_attention = MultiHeadAttention(8, 128)normal_tensor = torch.randn(32, 50, 128)  # 随机生成均值为0，方差为1的正态分布。batch_size=32,序列长度=50，embedding维度=128。x = torch.sigmoid(normal_tensor)  # 把每个数缩放到(0,1)output = multi_head_attention(x, x, x)print('done')