# 注意力计算函数
def attention(Q, K, V, mask):
# b句话,每句话50个词,每个词编码成32维向量,4个头,每个头分到8维向量
# Q,K,V = [b, 4, 50, 8]
# [b, 4, 50, 8] * [b, 4, 8, 50] -> [b, 4, 50, 50]
# Q,K矩阵相乘,求每个词相对其他所有词的注意力
score = torch.matmul(Q, K.permute(0, 1, 3, 2))
# 除以每个头维数的平方根,做数值缩放
score /= 8 ** 0.5
# mask遮盖,mask是true的地方都被替换成-inf,这样在计算softmax的时候,-inf会被压缩到0
# mask = [b, 1, 50, 50]
score = score.masked_fill_(mask, -float('inf'))
score = torch.softmax(score, dim=-1)
# 以注意力分数乘以V,得到最终的注意力结果
# [b, 4, 50, 50] * [b, 4, 50, 8] -> [b, 4, 50, 8]
score = torch.matmul(score, V)
# 每个头计算的结果合一
# [b, 4, 50, 8] -> [b, 50, 32]
score = score.permute(0, 2, 1, 3).reshape(-1, 50, 32)
return score
LN和BN
位置编码矩阵:
# 位置编码层
class PositionEmbedding(torch.nn.Module):
def __init__(self):
super().__init__()
# pos是第几个词,i是第几个维度,d_model是维度总数
def get_pe(pos, i, d_model):
fenmu = 1e4 ** (i / d_model)
pe = pos / fenmu
if i % 2 == 0:
return math.sin(pe)
return math.cos(pe)
# 初始化位置编码矩阵
pe = torch.empty(50, 32)
for i in range(50):
for j in range(32):
pe[i, j] = get_pe(i, j, 32)
pe = pe.unsqueeze(0)
# 定义为不更新的常量
self.register_buffer('pe', pe)
# 词编码层
self.embed = torch.nn.Embedding(39, 32)
# 初始化参数
self.embed.weight.data.normal_(0, 0.1)
def forward(self, x):
# [8, 50] -> [8, 50, 32]
embed = self.embed(x)
# 词编码和位置编码相加
# [8, 50, 32] + [1, 50, 32] -> [8, 50, 32]
embed = embed + self.pe
return embed