动手学深度学习（Pytorch版）代码实践 -注意力机制-Transformer

原创

已于 2024-07-11 09:07:04 修改 · 1.9k 阅读

17 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #pytorch #transformer

于 2024-07-10 19:09:07 首次发布

68`Transformer`

在这里插入图片描述

1. PositionWiseFFN

基于位置的前馈网络

原理：这是一个应用于每个位置的前馈神经网络。它使用相同的多层感知机（MLP）对序列中的每个位置独立进行变换。
作用：对输入序列的每个位置独立地进行非线性变换，增强模型的表达能力。
组成：
- 两个线性层
- 一个ReLU激活函数

class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    # 基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP）
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,
                 **kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)  # 第一个全连接层
        self.relu = nn.ReLU()  # ReLU激活函数
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)  # 第二个全连接层

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))  # 前向传播

2. AddNorm

残差连接和层规范化

原理：将残差连接（输入直接加到输出）和层规范化结合在一起。层规范化在归一化每个输入样本后，通过一个可训练的缩放和平移参数来调整输出。
作用：通过残差连接解决深度神经网络训练中的梯度消失问题，并通过层规范化稳定网络训练。
组成：
- 一个Dropout层
- 一个LayerNorm层

class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)  # Dropout层
        self.ln = nn.LayerNorm(normalized_shape)  # LayerNorm层

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)  # 残差连接后进行层规范化

3. PositionalEncoding

位置编码

原理：位置编码为每个输入位置添加一个固定的向量，使模型能够感知序列中元素的位置。这些位置编码向量使用正弦和余弦函数生成。
作用：在不改变模型结构的情况下，为Transformer提供序列的位置信息，使其能够处理顺序数据。
组成：
- 一个Dropout层
- 一个存储位置编码的张量

class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)  # Dropout层
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))  # 初始化位置编码张量
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)  # 偶数位置使用sin函数
        self.P[:, :, 1::2] = torch.cos(X)  # 奇数位置使用cos函数

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)  # 添加位置编码
        return self.dropout(X)  # 应用Dropout

4. MultiHeadAttention

多头注意力

原理：多头注意力机制将输入拆分成多个头，每个头独立地计算注意力，然后将这些头的输出拼接并进行一次线性变换。每个头有自己的查询、键和值的线性变换。
作用：通过并行计算多个注意力机制，增强模型的表示能力和捕捉不同特征的能力。
组成：
- 多个线性层
- 一个点积注意力层

class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads  # 注意力头数
        self.attention = d2l.DotProductAttention(dropout)  # 点积注意力
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)  # 线性变换层，用于生成查询向量
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)  # 线性变换层，用于生成键向量
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)  # 线性变换层，用于生成值向量
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)  # 线性变换层，用于生成输出向量

    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状:
        # (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状:
        # (batch_size，)或(batch_size，查询的个数)
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数，
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)  # 变换查询向量
        keys = transpose_qkv(self.W_k(keys), self.num_heads)  # 变换键向量
        values = transpose_qkv(self.W_v(values), self.num_heads)  # 变换值向量

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])

def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)