自然语言处理：transfomer架构

# 定义位置编码层
class PositionEmbedding(torch.nn.Module) :
    def __init__(self):
        super().__init__()
        # pos是第几个词，i是第几个词向量维度，d_model是编码维度总数
        def get_pe(pos, i, d_model):
            d = 1e4**(i / d_model)
            pe = pos / d
            if i % 2 == 0:
                return math.sin(pe) # 偶数维度用sin
            return math.cos(pe) # 奇数维度用cos
        # 初始化位置编码矩阵
        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) # 增加一个维度，shape变为[1,50,32]
        # 定义为不更新的常量
        self.register_buffer('pe', pe)
        # 词编码层
        self.embed = torch.nn.Embedding(39, 32) # 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

# 定义mask_pad函数
def mask_pad(data):
    # b句话，每句话50个词，这里是还没embed的
    # data = [b, 50]
    # 判断每个词是不是<PAD>
    mask = data == vocab_x['<PAD>']
    # [b, 50] -> [b, 1, 1, 50]
    mask = mask.reshape(-1, 1, 1, 50)
    # 在计算注意力时，计算50个词和50个词相互之间的注意力，所以是个50*50的矩阵
    # PAD的列为True，意味着任何词对PAD的注意力都是0，但是PAD本身对其它词的注意力并不是0，所以是PAD的行不为True
    # 复制n次
    # [b, 1, 1, 50] -> [b, 1, 50, 50]
    mask = mask.expand(-1, 1, 50, 50) # 根据指定的维度扩展
    return mask

# 定义mask_tril函数
def mask_tril(data):
    # b句话，每句话50个词，这里是还没embed的
    # data = [b, 50]
    # 50*50的矩阵表示每个词对其它词是否可见
    # 上三角矩阵，不包括对角线，意味着对每个词而言它只能看到它自己和它之前的词，而看不到之后的词
    # [1, 50, 50]
    """
    [[0, 1, 1, 1, 1],
    [0, 0, 1, 1, 1],
    [0, 0, 0, 1, 1],
    [0, 0, 0, 0, 1],
    [0, 0, 0, 0, 0]]
    """
    tril = 1 - torch.tril(torch.ones(1, 50, 50, dtype=torch.long)) # torch.tril返回下三角矩阵，则1-tril返回上三角矩阵
    # 判断y当中每个词是不是PAD, 如果是PAD, 则不可见
    # [b, 50]
    mask = data == vocab_y['<PAD>'] # mask的shape为[b, 50]
    # 变形+转型，为了之后的计算
    # [b, 1, 50]
    mask = mask.unsqueeze(1).long() # 在指定位置插入维度，mask的shape为[b, 1, 50]
    # mask和tril求并集
    # [b, 1, 50] + [1, 50, 50] -> [b, 50, 50]
    mask = mask + tril
    # 转布尔型
    mask = mask > 0 # mask的shape为[b, 50, 50]
    # 转布尔型，增加一个维度，便于后续的计算
    mask = (mask == 1).unsqueeze(dim=1) # mask的shape为[b, 1, 50, 50]
    return mask

# 定义注意力计算函数
def attention(Q, K, V, mask):
    """
    Q：torch.randn(8, 4, 50, 8)
    K：torch.randn(8, 4, 50, 8)
    V：torch.randn(8, 4, 50, 8)
    mask：torch.zeros(8, 1, 50, 50)
    """
    # 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)) # K.permute(0, 1, 3, 2)表示将K的第3维和第4维交换
    # 除以每个头维数的平方根，做数值缩放
    score /= 8**0.5
    # mask遮盖，mask是True的地方都被替换成-inf，这样在计算softmax时-inf会被压缩到0
    # mask = [b, 1, 50, 50]
    score = score.masked_fill_(mask, -float('inf')) # masked_fill_()函数的作用是将mask中为1的位置用value填充
    score = torch.softmax(score, dim=-1) # 在最后一个维度上做softmax
    # 以注意力分数乘以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