乐于分享
好东西不私藏

Transformer源码详解

Transformer源码详解

Transformer

Transformer 是 2017 年论文 《Attention Is All You Need》 提出的深度学习模型,完全基于自注意力机制(Self‑Attention),抛弃 RNN/LSTM 循环结构,是现在大语言模型、BERT、GPT、ViT 的基础底座。整体模型结构如图所示

位置编码

位置编码是一种向Transformer模型输入序列中注入词序信息的技术。由于Transformer的自注意力机制本身不具备处理顺序的能力,位置编码被用来显示地表示每个词在序列中的位置。

它通过以下方式工作:

基本形式:位置编码是将词嵌入(代表词义)与位置向量(代表词位置)进行相加,得到一个同时包含语义和位置信息的新向量。在Transformer中位置编码通过一组正弦和余弦函数预先计算得出。对于位置pos和维度i其计算公式为:

  • 偶数维度:
  • 奇数维度:

这个方法的优势在于它的值域在[-1, 1]之间,数值稳定;每个位置都有其独特的编码向量;并且其正弦函数的周期性使得模型更容易捕捉到词与词之间的相对位置关系

具体流程如图所示代码如下

classPositionalEncoding(nn.Module):'''为输入序列注入位置信息,使用正弦和余弦函数生成固定的位置编码'''def__init__(self, d_model, dropout=0.1, max_len=5000):'''        参数:        :param d_model: 词嵌入的维度        :param dropout: dropout概率        :param max_len: 支持的最大序列长度        '''        super(PositionalEncoding, self).__init__()        self.dropout = nn.Dropout(p=dropout) # Dropout层,防止过拟合# 创建一个形状为(max_len, d_model)的零矩阵,用于存储位置编码        pe = torch.zeros(max_len, d_model)# 生成位置索引,形状(max_len, 1)        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)# 计算分母中的频率项,div_term形状为 (d_model/2)        div_term = torch.exp(torch.arange(0, d_model, 2).float() *                             (-math.log(10000.0)/ d_model))# 偶数维度使用正弦函数(0,2,4,...)        pe[:, 0::2] = torch.sin(position / div_term)# 奇数维度使用余弦函数(1,3,5...)        pe[:, 1::2] = torch.cos(position / div_term)# 增加batch维度,变为(1, max_len, d_model),方便广播        pe = pe.unsqueeze(0)# 将 pe 注册为缓冲区,不参与梯度更新,但会保存到模型中        self.register_buffer('pe', pe)

多头注意力机制

多头注意力机制是一种让模型同时从多个角度关注输入序列不同位置的注意力方法。它通过将查询、键、值投影到多个子空间,并执行多次注意力计算,最后将结果拼接起来。

多头注意力 = 多个并行的"单头注意力" + 结果拼接 + 输出投影

因此我们需要了解单头注意力,其计算公式为:

直观理解:

单头注意力 = 一个人看一段文字,只从一个角度理解输入:“我 爱 自然语言处理” -> 一个人看这段文字 -> 只关注一种关系(比如语法关系)->输出:一种理解

多头注意力,其计算公式为:

其中每个头:

多头注意力前向传播的流程为:

线性变换并拆分多头 -> 计算注意力分数 -> 应用掩码(可选)-> 计算注意力权重 -> 加权求和 -> 拼接多头并投影

具体流程如图所示矩阵变换如图所示代码如下

classMultiHeadAttention(nn.Module):"""    多头注意力机制,将输入映射到多个子空间,并进行缩放点积注意力,    最后拼接并线性投影。    """def__init__(self, d_model, n_head, dropout=0.1):"""        参数:        :param d_model: 模型总维度        :param n_head: 注意力头的数量        :param dropout: dropout概率        """        super(MultiHeadAttention, self).__init__()assert d_model % n_head == 0# d_model 必须能被 n_head整除        self.d_model = d_model       # 模型维度        self.n_head = n_head         # 头数        self.d_k = d_model // n_head # 每个头的维度# 定义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) # 值变换矩阵# 输出投影层,将多头拼接后的结果映射回 d_model 维度        self.fc_out = nn.Linear(d_model, d_model)# Dropout层        self.dropout = nn.Dropout(dropout)# 缩放因子,用于点积注意力的缩放        self.scale = math.sqrt(self.d_k)defforward(self, query, key, value, mask=None):"""        前向传播        参数:        :param query: 查询张量,形状 (batch_size, q_len, d_model)        :param key: 键张量,形状 (batch_size, k_len, d_model)        :param value: 值张量,形状 (batch_size, v_len, d_model)        :param mask: 掩码张量,可选,形状 (batch_size, 1, 1, seq_len) 或可广播        :return: 注意力输出,形状 (batch_size, q_len, d_model)        """        batch_size = query.size(0# 获取 batch 大小# 线性变换并拆分为多个头# 每个头: (batch_size, seq_len, d_model) -> (batch_size, seq_len, n_head, d_k)# 然后转置为 (batch_size, n_head, seq_len, d_k)        Q = self.w_q(query).view(batch_size, -1, self.n_head, self.d_k).transpose(12)        K = self.w_k(key).view(batch_size, -1, self.n_head, self.d_k).transpose(12)        V = self.w_v(value).view(batch_size, -1, self.n_head, self.d_k).transpose(12)# 计算注意力分数(缩放点积)# Q 与 K 的转置相乘,形状:(batch_size, n_head, q_len, k_len)        attn_scores = torch.matmul(Q, K.transpose(-2-1)) / self.scale# 应用掩码if mask isnotNone:# 掩码中值为0的位置会被替换为一个极大的负数,使得 softmax 后接近 0            attn_scores = attn_scores.masked_fill(mask == 0-1e9)# 计算注意力权重(softmax 归一化)        attn_weights = F.softmax(attn_scores, dim=-1)# 对注意力权重应用 dropout        attn_weights = self.dropout(attn_weights)# 加权求和: 注意力权重乘以 V# 输出形状: (batch_size, n_head, q_len, d_k)        output = torch.matmul(attn_weights, V)# 将多个头的输出拼接回 d_model 维度# 先转置: (batch_size, q_len, n_head, d_k)# 再拼合: (batch_size, q_len, d_model)        output = output.transpose(12).contiguous().view(batch_size, -1, self.d_model)# 线性投影        output = self.fc_out(output)return output

前馈网络

前馈网络是Transformer中每个编码器和解码器层都包含的、独立应用于每个位置的全连接神经网络

前馈网络 = 两层全连接层 + ReLU激活函数,对每个位置独立处理

公式如下:

为社么需要前馈网络?

多头注意力可以捕捉序列中词与词之间的关系(依赖关系),但是它只是线性变换的组合(虽然有非线性softmax,但主要是特征混合),需要更强的非线性变换来学习复杂特征。前馈网络的作用:对每个位置的表示进行非线性变换,可以增强模型的表达能力(非线性),特征变换和增强,引入更多的可学习参数。

具体流程如图代码如下

classPositionwiseFeedForward(nn.Module):"""两层全连接网络,中间使用 ReLU 激活,对每个位置独立作用"""def__init__(self, d_model, d_ff, dropout=0.1):"""        参数:        :param d_model: 输入输出维度        :param d_ff: 前馈网络隐藏层维度(通常 2048)        :param dropout: dropout概率        """        super(PositionwiseFeedForward, self).__init__()        self.w_1 = nn.Linear(d_model, d_ff)     # 第一线性层,扩展维度        self.w_2 = nn.Linear(d_ff, d_model)     # 第二线性层,还原维度        self.dropout = nn.Dropout(dropout)      # Dropout层defforward(self, x):"""        前向传播        参数        :param x: 输入张量,形状 (batch_size, seq_len, d_model)        :return: 输出张量,形状同输入        """# 第一层线性变换后接 ReLU 激活,然后 dropout,再通过第二线性层return self.w_2(self.dropout(F.relu(self.w_1(x))))

编码器层

编码器层是Transformer编码器的基本单元,由两个核心子层组成:多头自注意力机制和前馈网络,每个子层都配有残差连接和层归一化。

编码器层 = 两个处理步骤(注意力 + 前馈),每一步都配有"残差连接 + 归一化"的增强装置

编码器层

子层1:多头自注意力:让每个词关注序列中的所有其它词 输出:融合了上下文信息的表示|连接1:Add & Norm:残差连接:输入+子层输出 LayerNorm:对结果进行归一化|子层2:前馈网络:对每个位置独立进行非线性变换 输出:增强的特征表示|连接2:Add & Norm:残差连接:输入+子层输出 LayerNorm:对结果进行归一化

编码器层结构如下

层归一化

层归一化是对单个样本的所有特征维度进行归一化,使其均值为0、方差为1的标准化方法。

公式如下:

其中:

  • μ:均值

  • σ2:方差

  • γ:可学习的缩放参数(scale)

  • β:可学习的偏移参数(shift)

  • ϵ:防止除零的小数(如1e-5)

  • ⊙:逐元素相乘

使用层归一化的作用

1、内部协变量偏移:在训练过程中由于网络参数不断变化,导致每层输入的分布发生变化,上层需要不断适应下层的输出分布,训练变得不稳定,收敛慢。层归一化可以让每层输入保持相同的分布(均值0,方差1),让每层都在“熟悉”的输入分布上学习。

2、梯度消失/爆炸:深层网络中梯度在反向传播时不断相乘,没有归一化可能导致梯度指数级衰减(消失)或增长(爆炸)。层归一化可以让梯度保持在合理范围内,训练深层网络更稳定。

3、序列长度变化:每个样本独立归一化,不受长度影响,完全适配变长序列

前向传播

反向传播(梯度流)

损失 L 对 x 的梯度

具体情况如图LayerNorm 和 BatchNorm的区别如图所示

残差连接

残差连接是一种“跳跃连接”技术,将子层的输入直接加到输出上,让网络学习“残差”(即输入与输出之间的差异)

公式如下:

残差连接的作用

解决梯度消失/爆炸:随着网络层数的增加,反向传播时梯度从输出传到输入,需要连乘N次,如果每层梯度都是0.9: 0.9¹⁰⁰ ≈ 0.000026 (消失),如果每层都是1.1:1.1¹⁰⁰ ≈ 13780 (爆炸)。梯度消失

-> 浅层无法更新 -> 训练停滞

梯度计算(残差连接)

即使很小,梯度也能通过直接路径传播,浅层始终保持很充足的梯度信号。

具体流程如图所示编码器层代码如下

classEncoderLayer(nn.Module):"""    一个编码器层,包含:    1. 多头自注意力子层    2. 前馈网络子层    每个子层前后都有层归一化(LayerNorm)和残差连接。    """def__init__(self, d_model, n_head, d_ff, dropout=0.1):"""        参数:        :param d_model: 模型维度        :param n_head: 注意力头数        :param d_ff: 前馈网络隐藏层维度        :param dropout: dropout概率        """        super(EncoderLayer, self).__init__()        self.self_attn = MultiHeadAttention(d_model, n_head, dropout) # 自注意力层        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout) # 前馈网络        self.norm1 = nn.LayerNorm(d_model) # 第一个子层的层归一化        self.norm2 = nn.LayerNorm(d_model) # 第二个子层的层归一化        self.dropout1 = nn.Dropout(dropout) # 自注意力后的 dropout        self.dropout2 = nn.Dropout(dropout) # 前馈网络后的 dropoutdefforward(self, x, mask=None):"""        前向传播        :param x: 输入张量,形状 (batch_size, seq_len,, d_model)        :param mask: 自注意力掩码(例如填充掩码)        :return: 输出张量,形状同输入        """# 自注意力层: 残差连接 + 层归一化# 首先计算自注意力的输出,这里 query, key, value 都是x        attn_output = self.self_attn(x, x, x, mask)# 应用 dropout,然后残差连接:x + Dropout(Attn(x))        x = self.norm1(x + self.dropout1(attn_output))# 前馈子层: 残差连接 + 层归一化        ff_output = self.feed_forward(x)        x = self.norm2(x + self.dropout2(ff_output))return x

解码器层

解码器层是Transformer编码器的基本单元,由三个核心子层组成:掩码自注意力和交叉注意力和前馈网络,每个子层都配有残差连接和层归一化。

解码器层结构如下:

目标掩码(Target Mask)

目标掩码(又称因果掩码、未来掩码)用于解码器的自注意力,确保每个位置只能关注到它之前(包括自己)的位置,不能看到未来的词。

为什么需要目标掩码?

解码器是自回归生成翻译任务: "I love AI" → "我爱人工智能"生成过程:Step 1: 输入 <SOS> → 预测 "我"Step 2: 输入 <SOS> 我 → 预测 "爱"  Step 3: 输入 <SOS> 我爱 → 预测 "人工智能"关键: Step 2 预测"爱"时,不能看到"人工智能"       否则就是作弊!(提前知道了答案)目标掩码保证了:"爱" 的位置只能看到 [<SOS>, 我]不能看到 [<SOS>, 我, 爱, 人工智能]

目标掩码对注意力的影响

没有掩码的注意力 (Q=K=V=[<SOS>, 我, 爱]):         <SOS>   我      爱<SOS>   [0.33,  0.33,  0.34]我      [0.32,  0.34,  0.34]  ← "我" 能看 "爱" (未来!)爱      [0.33,  0.33,  0.34]有掩码后:         <SOS>   我      爱<SOS>   [1.00,  0.00,  0.00]  ← 只能看自己我      [0.50,  0.50,  0.00]  ← 只能看 <SOS> 和我爱      [0.33,  0.33,  0.34]  ← 能看所有"我" 无法看到 "爱" 

具体情况如图

源掩码(Source Mask)

源掩码用于忽略输入序列中的填充(Padding)部分,确保注意力机制只关注真实词,不关注无意义的填充符。

为什么需要源掩码?

问题: 序列长度不固定批处理中,所有序列必须等长:样本1: "I love AI" (长度3)样本2: "Deep learning is great" (长度4)样本3: "Transformer" (长度1)处理: 统一填充到最大长度4样本1: "I love AI <PAD>"样本2: "Deep learning is great"样本3: "Transformer <PAD> <PAD> <PAD>"问题: 如果关注 <PAD>,会引入噪声      "I" 不应该关注 "love" 和 "AI" 之间的 <PAD>

源掩码在交叉注意力中的应用

编码器输出: (batch=1, src_len=4, d_model=512)位置0: "I" 的向量位置1: "love" 的向量位置2: "AI" 的向量位置3: <PAD> 的向量 (无意义)解码器查询: "我" 的向量没有源掩码:注意力权重: [0.2, 0.3, 0.2, 0.3]  ← 关注了 <PAD>!预测会受影响有源掩码:注意力分数中,位置3设为 -∞Softmax后: [0.28, 0.43, 0.29, 0.00]  ← <PAD> 被忽略模型只关注真实词 

具体情况如图目标掩码与源掩码的区别解码器层代码如下:

classDecoderLayer(nn.Module):"""    一个解码器层,包含:    1. 掩码多头自注意力子层    2. 编码器-解码器交叉注意力子层    3. 前馈网络子层    每个子层前后都有层归一化和残差连接    """def__init__(self, d_model, n_head, d_ff, dropout=0.1):"""        参数:        :param d_model: 模型维度        :param n_head: 注意力头数        :param d_ff: 前馈网络隐藏层维度        :param dropout: dropout概率        """        super(DecoderLayer, self).__init__()        self.self_attn = MultiHeadAttention(d_model, n_head, dropout) # 解码器自注意力        self.cross_attn = MultiHeadAttention(d_model, n_head, dropout) # 编码器-解码器注意力        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout) # 前馈网络        self.norm1 = nn.LayerNorm(d_model) # 自注意力后的层归一化        self.norm2 = nn.LayerNorm(d_model) # 交叉注意力后的层归一化        self.norm3 = nn.LayerNorm(d_model) # 前馈网络后的层归一化        self.dropout1 = nn.Dropout(dropout)        self.dropout2 = nn.Dropout(dropout)        self.dropout3 = nn.Dropout(dropout)defforward(self, x, enc_output, src_mask=None, tgt_mask=None):"""        前向传播        参数        :param x: 解码器输入,形状 (batch_size, tgt_len, d_model)        :param enc_output: 编码器输出,形状 (batch_size, src_len, d_model)        :param src_mask: 源序列掩码(用于交叉注意力,可选)        :param tgt_mask: 目标序列掩码(用于自注意力,防止看到未来信息)        :return: 输出张量,形状同 x        """# 掩码自注意力子层# query,key, value 均为 x,但使用 tgt_mask 进行遮掩        attn_output = self.self_attn(x, x, x, tgt_mask)        x = self.norm1(x + self.dropout1(attn_output))# 编码器-解码器交叉注意力子层# query 来自解码器 x,key 和 value 来自编码器输出 enc_output        cross_output = self.cross_attn(x, enc_output, enc_output, src_mask)        x = self.norm2(x + self.dropout2(cross_output))# 前馈子层        ff_output = self.feed_forward(x)        x = self.norm3(x + self.dropout3(ff_output))return x

编码器

编码器是Transformer架构中负责将输入序列(如源语言句子)转换为连续表示向量的组件,这些表示向量捕捉了序列中词与词之间的依赖关系和上下文信息。

编码器 = "阅读理解器",将原始文本转化为计算机能理解的深层表示

编码器的五大组件

词嵌入层:将词索引转为稠密向量|缩放操作:匹配位置编码的尺度|位置编码:注入位置信息|N个编码器层:多头自注意力、前馈网络|上下文表示:每个词都融合了序列信息

编码器的核心特性

双向上下文理解

每个词能关注到序列中的所有其他词句子: "The cat sat on the mat""sat" 关注:  - 前面: "The", "cat" (主语是谁)  - 后面: "on", "the", "mat" (在哪里)每个词都获得了完整的上下文信息这就是"双向"的含义

位置无关的并行处理

编码器可以同时处理序列的所有位置输入: [The, cat, sat, on, the, mat]         ↓    ↓    ↓    ↓    ↓    ↓所有位置同时计算自注意力         ↓    ↓    ↓    ↓    ↓    ↓输出: [v1,  v2,  v3,  v4,  v5,  v6]不需要像RNN那样按顺序处理训练速度快,可并行

不变输出维度

输入: (batch, src_len, d_model)输出: (batch, src_len, d_model)序列长度不变!每个位置的词向量都升级为"上下文感知"的表示

具体结构如图

词嵌入层

词嵌入层是将离散的词ID(索引)映射为稠密、连续的向量表示的查找表,为后续的神经网络处理提供数值化的输入。

词嵌入层 = "字典查询器",把词转换成计算机能理解的数字向量

公式如下:

其中:

  • :第个词的索引

  • :嵌入矩阵

  • 输出: 维向量

可视化

词汇表 (vocab_size=5):索引: 0   1     2     3     4词:  [PAD] [我]  [爱]  [AI]  [EOS]      ↓     ↓     ↓     ↓     ↓嵌入: [0,0,0] [0.8,0.2] [-0.2,0.9] [0.5,-0.4] [0,0,0]      (d_model=2)输入: [1, 2, 3] ("我爱AI")        ↓ 查表输出: [[0.8, 0.2], [-0.2, 0.9], [0.5, -0.4]]      (形状: 3 × 2)

词嵌入的三个关键特性

分布式表示

独热编码 (One-Hot):"猫": [0, 0, 1, 0, 0, ...]  (10000维,只有一个1)内存大、稀疏、无语义关系分布式表示 (Distributed Representation):"猫": [0.2, -0.5, 0.8, 0.1, ...]  (512维,密集)内存小、稠密、能捕捉语义关系相似词有相似的向量:"猫" → [0.2, -0.5, 0.8, 0.1]"狗" → [0.3, -0.4, 0.7, 0.2]  (相近!)"电视"→ [0.9, 0.3, -0.1, 0.5]  (不同!)

语义空间

词嵌入将词映射到一个向量空间这个空间中的几何关系反映语义关系:国王 - 男 + 女 ≈ 女王(king - man + woman ≈ queen)巴黎 - 法国 + 意大利 ≈ 罗马(paris - france + italy ≈ rome)Transformer的词嵌入在训练中会自动学习这些语义关系

可学习参数

# 嵌入矩阵是模型参数,通过训练学习embedding = nn.Embedding(vocab_size, d_model)# 训练过程中,梯度会更新嵌入矩阵# 类似词的向量会靠拢,不同词的向量会分开# 训练前: 随机初始化# 训练后: 学习到语义关系

具体情况如图

缩放操作

缩放操作是在词嵌入后,将嵌入向量乘以,使词嵌入的尺度与位置编码相匹配的关键步骤。

公式如下:

其中:

  • :词嵌入向量

  • :缩放因子

  • :位置编码

为什么需要缩放操作?根本原因是尺度不匹配

问题:词嵌入和位置编码的"尺度"不匹配词嵌入: 可学习的参数,数值范围较小 (标准差 ~1)位置编码: 正弦/余弦函数,值域固定 [-1, 1]直接相加: 词嵌入信息被位置编码"淹没"位置编码: [-1, 1] 范围词嵌入:   [-0.5, 0.5] 范围 (假设)词嵌入 + 位置编码 ≈ 位置编码 (词嵌入贡献小)解决方案: 放大词嵌入,使其与位置编码"旗鼓相当"

数学分析

词嵌入的初始化

# PyTorch Embedding 默认初始化nn.Embedding(vocab_size, d_model)# 权重 ~ Uniform(-√(1/d_model), √(1/d_model))# 标准差 ≈ 1/√d_model# 例如 d_model=512:# 标准差 ≈ 1/√512 ≈ 0.044# 词嵌入的值 ≈ [-0.13, 0.13] (很小!)

词嵌入 x 后:

x = embedding * √512# 标准差 ≈ √512 * 1/√512 = 1# 值 ≈ [-1.5, 1.5] (正常范围)# 现在词嵌入的尺度与位置编码 [-1, 1] 相当

可视化对比

d_model=512, 原始词嵌入值: [-0.13, -0.08, 0.02, 0.11, -0.05]                                    ↓ 缩放 √512 ≈ 22.6缩放后:                      [-2.94, -1.81, 0.45, 2.49, -1.13]位置编码:                    [0.84, 0.54, -0.42, 0.91, -0.14]相加:不缩放:  [-0.13, -0.08, 0.02, 0.11, -0.05] + [0.84, 0.54, -0.42, 0.91, -0.14]        = [0.71, 0.46, -0.40, 1.02, -0.19]  ← 几乎全是位置编码缩放后:  [-2.94, -1.81, 0.45, 2.49, -1.13] + [0.84, 0.54, -0.42, 0.91, -0.14]        = [-2.10, -1.27, 0.03, 3.40, -1.27]  ← 词嵌入和位置编码均衡

缩放操作的深层原因

  • 保持训练稳定性
  • 与模型初始化一致
  • 保持词嵌入的相对重要性

具体情况如图代码如下:

classEncoder(nn.Module):"""    完整编码器,由词嵌入 + 位置编码 + N 个编码器堆叠而成    """def__init__(self, vocab_size, d_model, n_head, d_ff, n_layers, dropout=0.1, max_len=5000):"""        参数:        :param vocab_size: 源语言词汇表大小        :param d_model: 模型维度        :param n_head: 注意力头数        :param d_ff: 前馈网络隐藏层维度        :param n_layers: 编码器层数        :param dropout: dropout概率        :param max_len: 位置编码支持的最大长度        """        super(Encoder, self).__init__()        self.token_embedding = nn.Embedding(vocab_size, d_model) # 词嵌入层        self.pos_encoding = PositionalEncoding(d_model, dropout, max_len) # 位置编码# 构建 n_layers 个编码器层        self.layers = nn.ModuleList([            EncoderLayer(d_model, n_head, d_ff, dropout) for _ in range(n_layers)        ])        self.d_model = d_model # 保存模型维度,用于后续缩放defforward(self, src, src_mask=None):"""        前向传播        :param src: 源序列索引,形状 (batch_size, src_len)        :param src_mask: 原序列掩码(填充掩码),形状可广播        :return: 编码器输出,形状 (batch_size, src_len, d_model)        """# 词嵌入并乘以 sqrt(d_model)(论文中的缩放操作)        x = self.token_embedding(src) * math.sqrt(self.d_model)# 添加位置编码        x = self.pos_encoding(x)# 逐层通过编码器层for layer in self.layers:            x = layer(x, src_mask)return x

解码器

解码器是Transformer架构中负责根据编码器输出的源语言表示,自回归地生成目标语言序列的组件。

解码器 = "翻译生成器",根据源语言理解,一步步生成目标语言

解码器的四大组件

词嵌入层:(vocab_size, d_model) 查找表|位置编码 :注入位置信息|N个解码器层:每个层:掩码自注意力 + 交叉注意力 + 前馈网络|输出投影层:(d_model -> vocab_size) 映射到词汇表

具体结构如图代码如下:

classDecoder(nn.Module):"""    完整解码器,由词嵌入 + 位置编码 + N 个解码器层堆叠 + 线性输出层(映射到词汇表)组成    """def__init__(self, vocab_size, d_model, n_head, d_ff, n_layers, dropout=0.1, max_len=5000):"""        参数:        :param vocab_size: 目标语言词汇表大小        :param d_model: 模型维度        :param n_head: 注意力头数        :param d_ff: 前馈网络隐藏层维度        :param n_layers: 解码器层数        :param dropout: dropout概率        :param max_lne: 位置编码支持的最大长度        """        super(Decoder, self).__init__()        self.token_embedding = nn.Embedding(vocab_size, d_model) # 词嵌入层        self.pos_encoding = PositionalEncoding(d_model, dropout, max_len) # 位置编码# 构建 n_layers 个解码器层        self.layers = nn.ModuleList([            DecoderLayer(d_model, n_head, d_ff, dropout) for _ in range(n_layers)        ])# 最终线性层,将 d_model 映射到词汇表大小,用于预测下一个词        self.fc_out = nn.Linear(d_model, vocab_size)        self.d_model = d_modeldefforward(self, tgt, enc_output, src_mask=None, tgt_mask=None):"""        前向传播        参数:        :param tgt: 目标序列索引,形状 (batch_size, tgt_len)        :param enc_output: 编码器输出,形状 (batch_size, src_len, d_model)        :param src_mask: 源序列掩码        :param tgt_mask: 目标序列掩码 (遮掩未来信息)        :return: 输出 logits,形状 (batch_size, tgt_len, vocab_size)        """# 词嵌入并缩放        x = self.token_embedding(tgt) * math.sqrt(self.d_model)# 添加位置编码        x = self.pos_encoding(x)# 逐层通过解码器层for layer in self.layers:            x = layer(x, enc_output, src_mask, tgt_mask)# 最后通过线性层映射到词汇表大小        output = self.fc_out(x)return output

完整Transformer模型

具体结构如图代码如下:

classTransformer(nn.Module):"""    完整的 Transformer 模型,包含编码器、解码器以及掩码生成逻辑。    """def__init__(self, src_vocab_size, tgt_vocab_size, d_model=512, n_head=8, d_ff=2048, n_layers=6, dropout=0.1, max_len=5000):"""        参数:        :param src_vocab_size: 源语言词汇表大小        :param tgt_vocab_size: 目标语言词汇表大小        :param d_model: 模型维度(默认512)        :param n_head: 注意力头数(默认8)        :param d_ff: 前馈网络隐藏层维度(默认2048)        :param n_layers: 编码器和解码器的层数(默认6)        :param dropout: dropout概率(默认0.1)        :param max_len: 最大序列长度(默认5000)        """        super(Transformer, self).__init__()        self.encoder = Encoder(src_vocab_size, d_model, n_head, d_ff, n_layers, dropout, max_len)        self.decoder = Decoder(tgt_vocab_size, d_model, n_head, d_ff, n_layers, dropout, max_len)        self.d_model = d_modeldefgenerate_padding_mask(self, seq, pad_idx=0):"""        生成填充掩码,屏蔽 <PAD> 位置        参数:        :param seq: 序列索引,形状 (batch_size, seq_len)        :param pad_idx: 填充标记的索引,默认为0        :return: 掩码,形状 (batch_size, 1, 1, sel_len),填充位置为0,其余为1        """# 将非填充位置设为 1, 填充位置设为 0        mask = (seq != pad_idx).unsqueeze(1).unsqueeze(2# (batch_size, 1, 1, seq_len)return maskdefgenerate_subsequent_mask(self, seq):"""        生成序列掩码(下三角矩阵),防止解码器看到未来信息        参数:        :param seq: 序列索引,形状 (batch_size, seq_len)        :return: 掩码,形状 (1, 1, seq_len, seq_len),上三角区域为 0        """        seq_len = seq.size(1)# 创建一个上三角矩阵,主对角线及其下方为1,上方为0        subsequent_mask = torch.triu(torch.ones((11, seq_len, seq_len), device=seq.device), diagonal=1)# 反转:下三角为 1,上三角为 0        subsequent_mask = (subsequent_mask == 0).to(seq.device)return subsequent_maskdefforward(self, src, tgt):"""        前向传播        参数:        :param src: 源序列索引,形状 (batch_size, src_len)        :param tgt: 目标序列索引,形状 (batch_size, tgt_len)        :return: 解码器输出 logits,形状 (batch_size, tgt_len, tgt_vocab_size)        """# 生成源序列的填充掩码        src_mask = self.generate_padding_mask(src)# 生成慕目标序列的填充掩码 + 下三角掩码(防止看到未来信息)        tgt_pad_mask = self.generate_padding_mask(tgt) # 填充掩码        tgt_sub_mask = self.generate_subsequent_mask(tgt) # 下三角掩码# 组合掩码:两者取逻辑与,即填充位置或未来位置为 0        tgt_mask = tgt_pad_mask & tgt_sub_mask# 编码器前向传播        enc_output = self.encoder(src, src_mask)# 解码器前向传播        dec_output = self.decoder(tgt, enc_output, src_mask, tgt_mask)return dec_output