理解Transformer架构

结合前文(链接已置顶)我们知道一个transfomer架构大致由这几部分组成:

  • 使用多头注意力机制输出混合了上下文的向量。
  • 将混合了上下文的向量,做残差连接(和原本的向量做加法)和LayerNorm做归一化处理。
  • FNN前馈网络做深度学习。
  • 对词表中的词做概率计算,输出预测。

尝试用代码实现多头注意力模块

需要先了解,torch包链接:https://pytorch-cn.com/tutorials:PyTorch 是一个开源的深度学习框架,主要由 Meta(原 Facebook)的 AI 研究团队开发和维护,现在已经成为机器学习/深度学习领域最主流的框架之一。

下面用到的主要方法:

  • torch.nn.Linear:是 PyTorch 中用来实现全连接层(线性层/仿射变换)的一个类.
  • torch.matmul:两个张量的乘积。
  • view方法将张量的数据重新排列,不会增加也不会减少数据。比如 12 个元素,可以变成 (3,4)、(4,3)、(2,6)、(12,) 等。
  • transpose:张量的转置方法,需要注意的是转置操作只是改变了张量的步幅(stride),并没有真正移动内存中的数据,所以转置后的张量通常在内存中是不连续的。要将张量转为连续,可以调用contiguous。
import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制模块
    """
    def __init__(self, d_model, num_heads):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除 后面拼接heads要用"

        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # 定义 Q, K, V 和输出的线性变换层
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def scaled_dot_product_attention(self, Q, K, V, mask=None):
        # 1. 计算注意力得分 (QK^T)
        attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)

        # 2. 应用掩码 (如果提供,主要是前文提到的Activation激活函数)
        if mask is not None:
            # 将掩码中为 0 的位置设置为一个非常小的负数,这样 softmax 后会接近 0
            attn_scores = attn_scores.masked_fill(mask == 0, -1e9)

        # 3. 计算注意力权重 (Softmax)
        attn_probs = torch.softmax(attn_scores, dim=-1)

        # 4. 加权求和 (权重 * V)
        output = torch.matmul(attn_probs, V)
        return output

    # 拆分向量为多头多批次处理,提高处理效率
    def split_heads(self, x):
        # 将输入 x 的形状从 (batch_size, seq_length, d_model)
        # 变换为 (batch_size, num_heads, seq_length, d_k)
        batch_size, seq_length, d_model = x.size()
        return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)

    def combine_heads(self, x):
        # 将输入 x 的形状从 (batch_size, num_heads, seq_length, d_k)
        # 变回 (batch_size, seq_length, d_model)
        batch_size, num_heads, seq_length, d_k = x.size()
        return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)

    def forward(self, Q, K, V, mask=None):
        # 1. 对 Q, K, V 进行线性变换
        Q = self.split_heads(self.W_q(Q))
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))

        # 2. 计算缩放点积注意力
        attn_output = self.scaled_dot_product_attention(Q, K, V, mask)

        # 3. 合并多头输出并进行最终的线性变换
        output = self.W_o(self.combine_heads(attn_output))
        return output

测试多头函数:

# 初始化一个att的多头实例
att = MultiHeadAttention(d_model=4, num_heads=2);

# 模拟输入的Q、K、V张量
my_q = torch.rand(2, 1, 4);
my_k = torch.rand(2, 1, 4);
my_v = torch.rand(2, 1, 4);

# 调用forward函数,输出进行了最终线性变换的
theOutput = att.forward(my_q, my_k, my_v);
print('output:', theOutput, theOutput.shape);

控制台看到结果output(一个Size(2,1,4)的张量):

继续用代码实现FNN

class PositionWiseFeedForward(nn.Module):
    """
    位置前馈网络模块
    """
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionWiseFeedForward, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        # x 形状: (batch_size, seq_len, d_model)
         x = self.linear1(x)
        print('larger:', x, x.shape);
        x = self.relu(x)
        x = self.dropout(x)
        print('after activation:', x);
        x = self.linear2(x)
        print('recur:',x, x.shape);
        # 最终输出形状: (batch_size, seq_len, d_model)
        return x

测试FNN:

# d_ff通常为d_model的4倍,前文有提到
theFFN = PositionWiseFeedForward(d_model=4, d_ff=16);
# 将多头注意后的结果给FNN
theFinal = theFFN.forward(theOutput);

print("Final:", theFinal);

最终运行结果:

由结果可见,整个FNN将张量的最后一个维度从4变为16后,经过relu激活函数将负数置为0后,再将维度还原回4。

至此,按照Encoder-Decoder结构,整体代码如下:

import torch
import torch.nn as nn
import math

# --- 占位符模块,将在后续小节中实现 ---

class PositionalEncoding(nn.Module):
    """
    位置编码模块,先不管pass
    """
    def forward(self, x):
        pass

class MultiHeadAttention(nn.Module):
    """
    多头注意力机制模块,已实现
    """
    def forward(self, query, key, value, mask):
        pass

class PositionWiseFeedForward(nn.Module):
    """
    位置前馈网络模块,已实现
    """
    def forward(self, x):
        pass

# --- 编码器核心层 ---

class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention() # 待实现
        self.feed_forward = PositionWiseFeedForward() # 待实现
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        # 残差连接与层归一化将在 3.1.2.4 节中详细解释
        # 1. 多头自注意力
        attn_output = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 2. 前馈网络
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))

        return x

# --- 解码器核心层 ---

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention() # 待实现
        self.cross_attn = MultiHeadAttention() # 待实现
        self.feed_forward = PositionWiseFeedForward() # 待实现
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, encoder_output, src_mask, tgt_mask):
        # 1. 掩码多头自注意力 (对自己)
        attn_output = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))

        # 2. 交叉注意力 (对编码器输出)
        # encoder_output作为编码器的输出传给交叉注意力头
        cross_attn_output = self.cross_attn(x, encoder_output, encoder_output, src_mask)
        x = self.norm2(x + self.dropout(cross_attn_output))

        # 3. 前馈网络
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))

        return x

类似文章

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注