Transformer实例代码全解析

下面将逐行/逐块解析这个Transformer实现的代码，详细说明每个组件的功能和设计逻辑：

1. 导入依赖库

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

• 导入PyTorch核心库（torch）、神经网络模块（nn）、功能函数（F，包含激活函数等）和数学库（math用于位置编码等计算）。

2. 位置编码类（PositionalEncoding）

位置编码用于给输入序列添加位置信息（Transformer本身是并行处理序列，没有内置的顺序感知能力）。

class PositionalEncoding(nn.Module):
    """位置编码，为输入序列添加位置信息"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)  # dropout层，防止过拟合
        
        # 计算位置编码矩阵 (max_len, d_model)
        pe = torch.zeros(max_len, d_model)  # 初始化位置编码矩阵
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # 位置索引 (max_len, 1)
        # 计算衰减因子：10000^(-2i/d_model)，用于控制正弦/余弦函数的周期
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        # 偶数索引用正弦函数，奇数索引用余弦函数（原始Transformer设计）
        pe[:, 0::2] = torch.sin(position * div_term)  # 0,2,4...列
        pe[:, 1::2] = torch.cos(position * div_term)  # 1,3,5...列
        
        # 调整形状为 (1, max_len, d_model) 便于广播，再转置为 (max_len, 1, d_model)
        pe = pe.unsqueeze(0).transpose(0, 1)
        # 将位置编码注册为非可学习参数（不参与梯度更新）
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        # x: (seq_len, batch_size, d_model) 输入序列嵌入
        # 给输入添加位置编码（广播机制适配不同batch和seq_len）
        x = x + self.pe[:x.size(0), :]  # 只取与输入序列长度匹配的位置编码
        return self.dropout(x)  # 应用dropout

• 核心逻辑：通过正弦/余弦函数生成位置编码，不同位置的编码具有不同周期，模型可通过学习捕捉位置关系。

• register_buffer：将位置编码存储为缓冲区（非参数），节省内存且不参与训练。

• 代码解释

  • super(MultiHeadAttention, self).__init__() 是 Python 中调用父类构造方法的语法：
    1. super(MultiHeadAttention, self) 用于获取 MultiHeadAttention 类的父类（这里是 nn.Module）
    2. .__init__() 调用父类的构造函数，确保父类中定义的初始化逻辑（如参数管理、设备配置等）被执行
    3. 在 PyTorch 中，所有自定义网络模块都需继承 nn.Module，通过这行代码可以正确初始化父类的核心功能，为当前模块提供参数管理、前向传播等基础能力。
  • nn.Dropout 是 PyTorch 中用于防止神经网络过拟合的正则化层。
    1. 其核心原理是：在训练过程中，以指定概率（如 dropout=0.1 表示 10% 概率）随机将输入张量中的部分元素置为 0，同时将剩余元素按比例放大（乘以 1/(1-概率)）以保持总体输入能量不变。
    2. 作用是：通过随机"丢弃"部分神经元的输出，强制网络学习更加鲁棒的特征（不依赖特定神经元），避免过度依赖训练数据中的噪声，从而提升模型泛化能力。
    3. 使用时通常在全连接层或卷积层后添加，测试/推理阶段会关闭（不进行丢弃操作）。
  • torch.zeros 是 PyTorch 中用于创建全为 0 的张量的函数。
    1. 其基本用法是根据指定的形状（如 (2, 3) 表示 2 行 3 列）生成元素均为 0 的张量，还可指定数据类型（如 dtype=torch.float32）和设备（如 CPU 或 GPU）等参数。
    2. 在上述位置编码实现中，pe = torch.zeros(max_len, d_model) 就是创建了一个形状为 (max_len, d_model)、初始值全为 0 的张量，用于后续填充位置编码值。
    3. 该函数在初始化参数矩阵、创建占位张量等场景中经常使用，是构建神经网络时初始化数据的基础工具之一。
  • torch.arange 是 PyTorch 中用于生成一维张量的函数，其功能是创建一个包含从起始值到结束值（不包含结束值）的均匀间隔数值的张量。
    1. 它的基本用法为 torch.arange(start=0, end, step=1, dtype=None, device=None, requires_grad=False)，其中：
       • start 为起始值，默认是 0；
       • end 为结束值（必选参数），生成的数值不包含该值；
       • step 是步长，默认是 1，决定相邻两个数值的间隔。
    2. 例如，torch.arange(0, 5, 1) 会生成张量 tensor([0, 1, 2, 3, 4])，在神经网络中常被用于生成索引、位置信息等，如在上述位置编码的实现中用于获取序列的位置索引（ position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 就是生成序列位置的索引张量）。
  • .unsqueeze(dim) 是用于给张量（Tensor）增加一个维度的方法，参数 dim 指定新增维度的位置。
    1. 例如，若有一个形状为 (3, 4) 的二维张量 x：
       • x.unsqueeze(0) 会在第 0 维增加维度，结果形状变为 (1, 3, 4)
       • x.unsqueeze(1) 会在第 1 维增加维度，结果形状变为 (3, 1, 4)
       • x.unsqueeze(-1) 会在最后一维增加维度，结果形状变为 (3, 4, 1)
    2. 该方法常用于调整张量维度以满足运算需求（如广播机制、矩阵乘法的维度匹配等），返回一个新的张量，不改变原张量。
  • torch.exp 是 PyTorch 中的一个张量运算函数，用于计算输入张量中每个元素的自然指数（以欧拉数 $e$ 为底的指数），即对张量 x 中的每个元素 $x_i$，计算结果为 $e^{x_i}$。
    1. 它的作用是对张量进行逐元素的指数变换，广泛应用于神经网络的激活函数（如 softmax 计算中需先对 logits 做指数运算）、概率分布转换等场景。例如在多头注意力机制中，计算注意力分数时可能会用到指数函数处理相关数值。
    2. 该函数支持广播机制，能处理不同形状但可广播的张量输入，返回与输入形状相同的张量。
  • math.log 是 Python 标准库 math 模块中的函数，用于计算对数。
    1. 默认情况下，math.log(x) 计算的是 自然对数（以常数 e 为底，e≈2.71828），即求满足 e^y = x 的 y 值，结果返回一个浮点数。
    2. 若要计算以其他数值为底的对数，可传入第二个参数，如 math.log(x, base) 表示计算以 base 为底 x 的对数（需保证 x>0 且 base>0、base≠1）。
    3. 例如：
       • math.log(math.e) 结果为 1.0（自然对数的底数的自然对数为1）
       • math.log(100, 10) 结果为 2.0（以10为底100的对数是2）
  • torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))用于计算Transformer位置编码中的衰减因子，核心作用是生成随维度增加而指数衰减的系数，具体解析如下：
    1. 计算逻辑：
       • torch.arange(0, d_model, 2) 生成从0到d_model（模型维度）的偶数序列（步长为2）
       • 乘以 -math.log(10000.0) / d_model 得到衰减基数，控制衰减速率
       • torch.exp() 将结果转换为指数形式，得到 10000^(-2i/d_model)（i为维度索引）
       • $$\frac{1}{10000^{2k/d_{\text{model}}}}=10000^{-2k/d_{\text{model}}}=e^{\log (10000)\cdot (-2k/d_{\text{model}})}=e^{2k\cdot (-\log (10000)/d_{\text{model}})}$$
    2. 本质作用：
       生成一组随维度增加而逐渐减小的系数，使得位置编码中不同维度的正弦/余弦函数具有不同周期（从2π到10000×2π），让模型能够学习到不同尺度的位置依赖关系。
    3. 与位置编码的关系：
       后续会将该系数与位置索引相乘，再分别作为正弦（偶数维度）和余弦（奇数维度）函数的参数，形成具有位置区分性的编码向量。
  • .transpose(dim0, dim1) 是用于交换张量维度的方法。
    1. 它接收两个参数 dim0 和 dim1，表示要交换的两个维度的索引，交换后张量的形状会相应改变，但数据本身不变。
    2. 例如，若有一个形状为 (2, 3) 的张量 x，执行 x.transpose(0, 1) 后，会得到一个形状为 (3, 2) 的张量，实现了矩阵的转置。
    3. 该方法常用于调整张量维度以满足运算需求（如注意力机制中对序列长度和批次维度的调整），返回的是原张量的视图（不复制数据），修改结果会影响原张量。
  • register_buffer 是 nn.Module 提供的方法，用于注册非可学习参数（缓冲区）。
    1. 这些参数会被视为模型状态的一部分，随模型一起保存（如通过 state_dict()）和加载，但不参与梯度计算和参数更新。
    2. 常见用途是存储模型中需要固定的常量或辅助数据（例如 Transformer 中的位置编码矩阵），既保证其随模型持久化，又避免被优化器误更新。
  • def forward(self, x) 是继承自 nn.Module 的自定义网络模块中必须实现的核心方法，用于定义前向传播逻辑。
    1. 其作用是：规定输入数据 x 在网络中的流动路径——即数据如何经过各层（如卷积层、注意力层等）的处理，最终得到输出结果。
    2. 当调用模型实例（如 model(input)）时，PyTorch 会自动触发 forward 方法，无需显式调用。方法中的 x 通常是输入张量，经过层运算后返回的张量即为网络输出。
    3. 这一方法是构建神经网络的核心，决定了数据的处理流程和网络功能。
  • x = x + self.pe[:x.size(0), :]的作用是将位置编码（self.pe）添加到输入序列的嵌入向量（x）中，为序列注入位置信息。具体解析：
    1. x 是输入序列的嵌入向量，形状通常为 (seq_len, batch_size, d_model)（序列长度、批次大小、模型维度）
    2. self.pe 是预计算的位置编码矩阵，形状为 (max_len, 1, d_model)（最大序列长度、1、模型维度）
    3. self.pe[:x.size(0), :] 截取与输入序列长度（x.size(0)）匹配的位置编码部分
    4. 通过加法 x = x + ... 将位置信息与嵌入向量融合，使模型能够感知序列中元素的位置关系。这一操作是 Transformer 中位置编码的核心应用，解决了自注意力机制本身不具备顺序感知能力的问题。
    5. 广播机制（Broadcasting）：由于 x 和截取后的 self.pe 形状不完全相同（x 有 batch_size 维度，self.pe 对应维度为 1），PyTorch 会通过 广播机制 自动扩展两者的维度，使形状匹配后再进行加法。广播规则：维度为 1 的维度会被扩展为另一个 tensor 对应维度的大小。因此 self.pe 的第 1 维（大小为 1）会被扩展为 batch_size，扩展后的形状为 (seq_len, batch_size, d_model)，与 x 的形状完全一致。
    6. 元素级加法（Element-wise Addition）：形状匹配后，加法按 元素一一对应相加：
       • 对于 x 中位置 (i, j, k) 的元素（第 i 个序列位置、第 j 个样本、第 k 个特征维度），与 self.pe 中位置 (i, j, k) 的元素（广播后）相加，结果存入 x 的对应位置。
       • 公式化表示：
         $$x_{\text{new}}[i,j,k]=x_{\text{old}}[i,j,k]+\text{self.pe}[i,1,k]$$

3. 多头注意力类（MultiHeadAttention）

多头注意力将输入拆分为多个"头"并行计算注意力，再拼接结果，增强模型对不同子空间特征的捕捉能力。

class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        # 确保模型维度能被头数整除（每个头的维度需一致）
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model  # 模型总维度
        self.num_heads = num_heads  # 注意力头数
        self.d_k = d_model // num_heads  # 每个头的维度（d_model / num_heads）
        
        # 线性变换层：将输入映射到Q、K、V（共享权重）
        self.w_q = nn.Linear(d_model, d_model)  # Q的线性投影
        self.w_k = nn.Linear(d_model, d_model)  # K的线性投影
        self.w_v = nn.Linear(d_model, d_model)  # V的线性投影
        
        self.w_o = nn.Linear(d_model, d_model)  # 多头结果拼接后的线性投影
        self.dropout = nn.Dropout(p=dropout)  # 注意力权重的dropout
    
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)  # 获取批次大小
        
        # 1. 线性变换并拆分多头
        # (batch_size, seq_len, d_model) → (batch_size, num_heads, seq_len, d_k)
        q = self.w_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 2. 计算注意力分数：Q·K^T / √d_k（缩放点积注意力）
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # 3. 应用掩码（如填充掩码或序列掩码，防止关注无效位置）
        if mask is not None:
            # 掩码为0的位置设为负无穷，softmax后权重接近0
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 4. 计算注意力权重（softmax归一化）
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)  # 对注意力权重应用dropout
        
        # 5. 应用注意力到V：加权求和
        output = torch.matmul(attn, v)  # (batch_size, num_heads, seq_len, d_k)
        
        # 6. 拼接多头结果（恢复为d_model维度）
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        # 7. 最终线性投影
        return self.w_o(output), attn

• 核心逻辑：通过拆分多头并行计算注意力，再合并结果，既保持计算效率又增强模型表达能力。
• 缩放因子√d_k：防止分数过大导致softmax梯度消失。
• 掩码（mask）：用于屏蔽填充符号（PAD）或未来序列（解码器自注意力）。
• 代码解释
  • nn.Linear(d_model, d_model) 是 PyTorch 中用于构建线性变换层的类，其中两个参数分别表示输入特征维度和输出特征维度，此处均为 d_model，意味着输入和输出维度相同。
    1. 其数学原理是对输入向量 x 执行线性变换：y = x·W + b，其中 W 是形状为 (d_model, d_model) 的可学习权重矩阵，b 是形状为 (d_model,) 的可学习偏置向量，运算结果 y 与输入 x 保持相同维度。
    2. 该层常用于神经网络中对特征进行线性映射，在 Transformer 等模型中，常被用于多头注意力机制中对 Q、K、V 向量的线性投影等场景，在不改变特征维度的同时调整特征分布。
  • 线性投影是将高维向量通过线性变换映射到另一维度空间的操作，核心是利用矩阵乘法实现。
    1. 在神经网络中，通常通过线性层（如PyTorch的nn.Linear）完成，公式为 y = x·W + b，其中x是输入向量，W是可学习的权重矩阵，b是偏置项（可选）。
    2. 其作用包括：
       • 维度转换：将输入从d_in维映射到d_out维（W的形状为[d_out, d_in]）。
       • 特征提取：通过学习权重矩阵，捕捉输入数据中的线性特征关联。
       • 适配网络结构：在不同模块（如注意力机制中的Q、K、V转换）之间充当“接口”，统一特征维度。
    3. 例如在多头注意力中，线性投影将输入分别映射为查询（Q）、键（K）、值（V），确保它们能进行后续的注意力计算。
  • def forward(self, query, key, value, mask=None):是PyTorch中多头注意力机制（MultiHeadAttention）的前向传播方法定义，用于实现注意力计算的核心逻辑：
    1. query、key、value：注意力机制的三个核心输入向量（通常形状为(batch_size, seq_len, d_model)），分别对应查询、键、值
    2. mask：可选参数，用于屏蔽无效位置（如填充符或未来序列），避免模型关注这些位置的信息
    3. 该方法的主要作用是：
       • 对query、key、value进行线性变换并拆分多头
       • 计算注意力分数（通过Q·K^T / √d_k）
       • 应用掩码（mask）屏蔽无效位置
       • 通过softmax计算注意力权重
       • 加权聚合value并拼接多头结果，输出最终注意力特征
    4. 这是Transformer中注意力机制的核心实现，支撑着模型对序列中不同位置信息的关注能力。
  • (batch_size, seq_len, d_model) → (batch_size, num_heads, seq_len, d_k)是多头注意力机制中对输入的维度转换，具体解析如下：
    1. 原始形状(batch_size, seq_len, d_model)：表示输入为批次数据，包含batch_size个样本，每个样本是长度为seq_len的序列，每个序列元素的特征维度为d_model（模型总维度）。
    2. 转换后形状(batch_size, num_heads, seq_len, d_k)：通过线性投影和维度拆分，将总特征维度d_model按头数num_heads均分（d_k = d_model // num_heads），得到num_heads个并行的注意力头。每个头处理的特征维度为d_k，保持序列长度seq_len和批次batch_size不变，便于多头并行计算注意力。
    3. 本质是将高维特征拆分为多个低维子空间，让不同头关注不同特征，提升模型表达能力。
  • .view() 是张量（Tensor）的方法，用于重塑张量的形状，但不改变张量的元素数量和数据顺序。

    1. 其核心特点：

       • 要求重塑前后的元素总数必须相同（例如，形状为 (2, 4) 的张量可重塑为 (1, 8) 或 (8,) 等）。
       • 返回的是原张量的视图（view），即与原张量共享数据存储空间，修改视图会影响原张量，反之亦然。
       • 常用于调整张量维度以适配后续运算（如神经网络层的输入形状要求）。

    2. 示例：

       x = torch.randn(2, 4)  # 形状为 (2, 4)，共 8 个元素
       y = x.view(1, 8)       # 重塑为 (1, 8)，与 x 共享数据
       

  • q = self.w_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)是多头注意力机制中对查询向量（query）进行处理的关键步骤，主要实现将输入向量拆分到多个注意力头的功能，具体解析如下：
    1. 线性变换：self.w_q(query) 通过线性层 w_q 对输入的 query 向量进行线性投影，将其映射到与模型维度 d_model 一致的空间。
    2. 维度重塑：.view(batch_size, -1, self.num_heads, self.d_k) 将线性变换后的向量重塑维度。其中：
       • batch_size 是批次大小
       • -1 表示自动计算该维度（通常是序列长度 seq_len）
       • self.num_heads 是注意力头的数量
       • self.d_k 是每个头的维度（d_model // num_heads）
       • 重塑后形状为 (batch_size, seq_len, num_heads, d_k)
    3. 维度转置：.transpose(1, 2) 交换序列长度维度和头数维度，最终形状变为 (batch_size, num_heads, seq_len, d_k)，便于后续每个头独立计算注意力。
    4. 这一步骤的核心是将输入向量拆分到多个并行的注意力头，让模型能够同时关注不同子空间的特征。
  • torch.matmul 是 PyTorch 中用于矩阵乘法的核心函数，支持多维张量的矩阵乘法运算，行为会根据输入张量的维度自动调整：
    1. 对于 2 维张量（矩阵），执行标准的矩阵乘法（即 (n×m) × (m×p) → (n×p)）。
    2. 对于高维张量（维度 > 2），将前 k-2 维视为“批次维度”，对最后两维执行矩阵乘法，且批次维度需满足广播规则。
    3. 支持标量与张量、向量与矩阵等特殊情况的乘法（如向量点积、矩阵与向量的乘法等）。
    4. 例如：
       • torch.matmul(torch.rand(2,3), torch.rand(3,4)) 输出形状为 (2,4) 的矩阵；
       • torch.matmul(torch.rand(5,2,3), torch.rand(5,3,4)) 输出形状为 (5,2,4) 的张量（对 5 个批次分别做 2×3 与 3×4 的矩阵乘法）。
  • .masked_fill 是 PyTorch 中张量（Tensor）的方法，用于根据掩码（mask）对张量中的元素进行填充。
    1. 其基本用法为 tensor.masked_fill(mask, value)，作用是：当掩码 mask 中元素为 True 时，将张量 tensor 对应位置的元素替换为指定值 value。
    2. 例如，在注意力机制中，常使用 scores.masked_fill(mask == 0, -1e9)，通过掩码将无效位置（如填充符、未来序列）的注意力分数设为极小值，确保这些位置在经过 softmax 后权重接近 0，不影响模型决策。
  • attn = F.softmax(scores, dim=-1) 是注意力机制中的关键操作，作用是将注意力分数（scores）转换为注意力权重：
    1. scores 是注意力分数矩阵，形状通常为 (batch_size, num_heads, seq_len_q, seq_len_k)，表示查询（Q）与键（K）的匹配程度。
    2. F.softmax(..., dim=-1) 对最后一个维度（即 seq_len_k 维度）执行 softmax 运算，将分数归一化为 [0, 1] 区间的概率分布，且该维度上所有元素之和为 1。
    3. 结果 attn 即为注意力权重，反映了每个查询位置对不同键位置的关注程度，后续会与值（V）相乘实现加权求和。
  • attn = self.dropout(attn) 是对注意力权重（attn）应用 dropout 正则化的操作。
    1. 在训练时，该操作会以预设概率（如 0.1）随机将 attn 中的部分元素置为 0，同时将剩余元素按 1/(1-概率) 比例放大，以保持整体能量不变。
    2. 其作用是通过随机丢弃部分注意力权重，防止模型过度依赖某些特定位置的关联信息，增强模型的泛化能力，减少过拟合风险。在推理（测试）阶段，dropout 会自动关闭，不进行任何处理。
  • output = torch.matmul(attn, v) 是多头注意力机制中应用注意力权重的核心操作，功能是将注意力权重与值向量（Value）进行矩阵乘法，得到加权求和后的特征输出。具体来说：
    1. attn 是注意力权重矩阵，形状通常为 (batch_size, num_heads, seq_len_q, seq_len_k)，表示每个查询位置对键位置的关注程度。
    2. v 是值向量矩阵，形状通常为 (batch_size, num_heads, seq_len_v, d_k)（其中 seq_len_k = seq_len_v），包含每个位置的特征信息。
    3. 矩阵乘法后，输出 output 的形状为 (batch_size, num_heads, seq_len_q, d_k)，其每个元素是对应查询位置下，所有值向量按注意力权重加权求和的结果，体现了模型对不同位置信息的关注度整合。
  • 在 PyTorch 中，.contiguous() 是用于确保张量在内存中连续存储的方法。
    1. 当张量经过转置（如 .transpose()）、切片等操作后，其数据在内存中的存储顺序可能变得不连续（即逻辑形状与内存布局不一致）。此时某些依赖连续内存的操作（如 .view() 重塑形状）会报错。
    2. 调用 .contiguous() 会通过复制数据，使张量在内存中恢复连续存储状态，确保后续操作能正常执行。它不改变张量的数值和形状，仅调整内存布局。
  • output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)的作用是将多头注意力计算后的结果重新拼接为原始维度，具体解析如下：
    1. transpose(1, 2)：交换维度1和维度2。多头注意力计算后的数据形状为(batch_size, num_heads, seq_len, d_k)，交换后变为(batch_size, seq_len, num_heads, d_k)，将序列长度维度提前，为后续拼接做准备。
    2. contiguous()：确保张量在内存中连续存储。维度交换后张量可能变得不连续，此操作重新排列内存中的数据，保证后续view操作能正常执行。
    3. view(batch_size, -1, self.d_model)：重塑张量形状。-1表示自动计算该维度大小（结果为seq_len），最终将num_heads和d_k这两个维度合并为d_model（因d_model = num_heads * d_k），得到形状为(batch_size, seq_len, d_model)的输出，与多头注意力输入的维度保持一致，便于后续的残差连接和层归一化操作。

4. 位置感知的前馈网络（PositionWiseFeedForward）

每个位置独立的前馈网络，用于对注意力输出进行非线性变换。

class PositionWiseFeedForward(nn.Module):
    """位置wise前馈网络"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)  # 升维：d_model → d_ff
        self.fc2 = nn.Linear(d_ff, d_model)  # 降维：d_ff → d_model
        self.dropout = nn.Dropout(p=dropout)
        self.activation = nn.GELU()  # 激活函数（比ReLU更优，原始论文用ReLU）
    
    def forward(self, x):
        # 前馈网络流程：线性变换→激活→dropout→线性变换
        return self.fc2(self.dropout(self.activation(self.fc1(x))))

• 作用：对每个位置的特征进行独立的非线性变换，增强模型的非线性拟合能力。
• d_ff：通常设为4*d_model（如2048对应512的d_model），扩大特征空间。
• 代码解释
  • nn.Linear(d_model, d_ff) 是 PyTorch 中用于实现线性变换的层，作用是将输入特征从 d_model 维度映射到 d_ff 维度。
  1. 其数学原理为：对输入张量 x（形状为 (..., d_model)）执行运算 y = x · W + b，其中：
     • W 是可学习的权重矩阵（形状为 (d_model, d_ff)）
     • b 是可学习的偏置向量（形状为 (d_ff)）
     • 输出 y 形状为 (..., d_ff)
  2. 在 Transformer 等模型中，常作为前馈网络的第一层用于特征升维（如 d_model=512 映射到 d_ff=2048），配合后续的非线性激活函数增强模型表达能力。
  • 在Transformer的前馈网络中，“升维：d_model → d_ff”指的是将输入的特征维度从模型的基础维度（d_model）提升到更高的维度（d_ff）。
    1. 具体来说，前馈网络通过线性层（如nn.Linear(d_model, d_ff)）对输入进行变换，让特征在更高维度的空间中进行非线性映射（如通过GELU等激活函数）。这样做的目的是扩大模型的表达空间，增强对复杂特征的捕捉能力，随后再通过另一个线性层将维度降回d_model，既保留了高维空间的特征提取优势，又保证输出维度与模型其他部分兼容。
    2. 通常d_ff设置为d_model的4倍（如d_model=512时，d_ff=2048），这是实践中验证的有效设计，能在计算效率和表达能力间取得平衡。通过升维将特征空间扩大 4 倍，再通过降维压缩回原维度，本质是一种 “特征扩展 - 提炼” 的过程。

5. 编码器层（EncoderLayer）

编码器的基本单元，包含多头自注意力和前馈网络，并带残差连接和层归一化。

class EncoderLayer(nn.Module):
    """编码器层: 包含多头自注意力和前馈网络，以及残差连接和层归一化"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)  # 多头自注意力
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)  # 前馈网络
        
        # 层归一化：稳定训练，加速收敛
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        # dropout层：应用于子层输出
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
    
    def forward(self, x, mask):
        # 1. 多头自注意力 + 残差连接 + 层归一化
        attn_output, _ = self.self_attn(x, x, x, mask)  # Q=K=V（自注意力）
        # 残差连接：x（输入） + 子层输出（带dropout），再归一化
        x = self.norm1(x + self.dropout1(attn_output))
        
        # 2. 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        # 残差连接：x（上一步输出） + 前馈网络输出（带dropout），再归一化
        x = self.norm2(x + self.dropout2(ff_output))
        
        return x

• 残差连接：x + sublayer_output，缓解深层网络梯度消失问题，让模型更容易训练。
• 层归一化：对每个样本的特征维度做归一化（均值0，方差1），放在残差连接后（现代改进，原始论文放前）。

6. 解码器层（DecoderLayer）

解码器的基本单元，包含掩码自注意力、编码器-解码器注意力和前馈网络。

class DecoderLayer(nn.Module):
    """解码器层: 包含掩码多头自注意力、编码器-解码器注意力和前馈网络"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)  # 掩码自注意力（防未来信息）
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)  # 编码器-解码器注意力
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)
        
        # 3个层归一化（对应3个子层）
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        
        # 3个dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
    
    def forward(self, x, enc_output, self_mask, cross_mask):
        # 1. 掩码自注意力（解码器输入序列内部的注意力，屏蔽未来位置）
        attn_output, _ = self.self_attn(x, x, x, self_mask)
        x = self.norm1(x + self.dropout1(attn_output))  # 残差+归一化
        
        # 2. 编码器-解码器注意力（用解码器输出作为Q，编码器输出作为K/V）
        attn_output, _ = self.cross_attn(x, enc_output, enc_output, cross_mask)
        x = self.norm2(x + self.dropout2(attn_output))  # 残差+归一化
        
        # 3. 前馈网络
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout3(ff_output))  # 残差+归一化
        
        return x

• 掩码自注意力（self_attn）：通过掩码防止解码器关注"未来"的序列（如翻译时，生成第i个词不能看i+1及以后的词）。
• 编码器-解码器注意力（cross_attn）：让解码器关注编码器输出的关键信息（如翻译中源语言和目标语言的对齐）。

7. 编码器（Encoder）

由多个编码器层堆叠而成，负责将输入序列编码为上下文特征。

class Encoder(nn.Module):
    """编码器: 由多个编码器层堆叠而成"""
    def __init__(self, input_vocab_size, d_model, num_layers, num_heads, 
                 d_ff, max_len, dropout=0.1):
        super(Encoder, self).__init__()
        self.d_model = d_model
        
        # 嵌入层：将输入词索引映射为d_model维度的向量
        self.embedding = nn.Embedding(input_vocab_size, d_model)
        # 位置编码层
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 堆叠num_layers个编码器层（这里num_layers=10）
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x, mask):
        # 1. 嵌入层 + 缩放（原始论文：嵌入向量乘以√d_model，增强梯度）
        x = self.embedding(x) * math.sqrt(self.d_model)
        # 2. 添加位置编码（先转置适应位置编码的维度，再转置回来）
        x = self.pos_encoding(x.transpose(0, 1)).transpose(0, 1)
        
        # 3. 通过所有编码器层（10层）
        for layer in self.layers:
            x = layer(x, mask)
        
        return x  # 输出编码器最终特征

• 嵌入层（Embedding）：将离散的词索引转换为连续的向量表示。
• 缩放嵌入：* math.sqrt(d_model) 平衡嵌入向量和位置编码的量级（原始论文设计）。
• 代码解释：
  • nn.Embedding 是 PyTorch 中用于将离散整数（如词索引）映射为连续低维向量的层，主要用于处理文本等离散输入。
    1. 其核心功能是：通过一个可学习的嵌入矩阵，将输入的整数序列（形状为 (batch_size, seq_len)）转换为向量序列（形状为 (batch_size, seq_len, embedding_dim)），其中 embedding_dim 是指定的输出向量维度（如 Transformer 中的 d_model）。
    2. 升维过程中的参数（即嵌入矩阵）是可学习的，会在反向传播中根据损失函数进行更新，以学习到更贴合任务的语义向量表示。例如，语义相近的词经过嵌入后，其向量在空间中的距离也会较近。
  • nn.ModuleList 是 PyTorch 中用于管理多个子模块的容器类，主要作用是将多个 nn.Module 实例（如层、子模型等）组织成列表形式，方便批量管理和迭代访问。它的核心特点：
    1. 会自动注册列表中的子模块，使这些子模块的参数被纳入整个模型的参数管理（可通过 parameters() 访问，参与反向传播更新）。
    2. 支持像普通列表一样通过索引访问子模块，也支持迭代操作（如 for layer in layers）。
    3. 不定义前向传播逻辑，仅用于模块的存储和管理，需在 forward 方法中手动实现子模块的调用顺序。
    4. 示例：

       self.layers = nn.ModuleList([
           EncoderLayer(d_model, num_heads) 
           for _ in range(num_layers)  # 批量创建多个编码器层
       ])
       

       这里的每个 EncoderLayer 都是独立子模块，其参数会被自动管理，在反向传播时会被更新。

8. 解码器（Decoder）

由多个解码器层堆叠而成，负责将编码器的上下文特征解码为目标序列。

class Decoder(nn.Module):
    """解码器: 由多个解码器层堆叠而成"""
    def __init__(self, target_vocab_size, d_model, num_layers, num_heads, 
                 d_ff, max_len, dropout=0.1):
        super(Decoder, self).__init__()
        self.d_model = d_model
        
        # 目标序列的嵌入层
        self.embedding = nn.Embedding(target_vocab_size, d_model)
        # 位置编码层（与编码器共享逻辑）
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 堆叠num_layers个解码器层（这里num_layers=10）
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x, enc_output, self_mask, cross_mask):
        # 1. 嵌入层 + 缩放
        x = self.embedding(x) * math.sqrt(self.d_model)
        # 2. 添加位置编码
        x = self.pos_encoding(x.transpose(0, 1)).transpose(0, 1)
        
        # 3. 通过所有解码器层（10层）
        for layer in self.layers:
            x = layer(x, enc_output, self_mask, cross_mask)
        
        return x  # 输出解码器最终特征

• 结构与编码器类似，但输入是目标序列，且每层需接收编码器的输出（用于交叉注意力）。

9. 完整Transformer模型

整合编码器和解码器，输出最终预测结果。

class Transformer(nn.Module):
    """完整的Transformer模型"""
    def __init__(self, input_vocab_size, target_vocab_size, 
                 d_model=512, num_layers=10, num_heads=8, 
                 d_ff=2048, max_len=5000, dropout=0.1):
        super(Transformer, self).__init__()
        
        # 初始化编码器（10层）
        self.encoder = Encoder(
            input_vocab_size, d_model, num_layers, num_heads, d_ff, max_len, dropout
        )
        
        # 初始化解码器（10层）
        self.decoder = Decoder(
            target_vocab_size, d_model, num_layers, num_heads, d_ff, max_len, dropout
        )
        
        # 最终输出层：将解码器输出映射到目标词汇表大小
        self.fc_out = nn.Linear(d_model, target_vocab_size)
        
        self.init_weights()  # 初始化模型权重
    
    def init_weights(self):
        """初始化模型权重（Xavier均匀初始化，适合激活函数为tanh/sigmoid的场景）"""
        for p in self.parameters():
            if p.dim() > 1:  # 只初始化矩阵参数（如权重），忽略偏置等1维参数
                nn.init.xavier_uniform_(p)
    
    def forward(self, src, trg, src_mask, trg_mask, cross_mask):
        # 1. 编码器输出：(batch_size, src_seq_len, d_model)
        enc_output = self.encoder(src, src_mask)
        
        # 2. 解码器输出：(batch_size, trg_seq_len, d_model)
        dec_output = self.decoder(trg, enc_output, trg_mask, cross_mask)
        
        # 3. 最终预测：(batch_size, trg_seq_len, target_vocab_size)
        output = self.fc_out(dec_output)
        
        return output

• 权重初始化：Xavier_uniform_ 确保前向和反向传播中信号的方差一致，避免梯度消失/爆炸。
• 输出层：将解码器的特征向量映射到目标词汇表，用于后续计算交叉熵损失。
• 代码解释
  • self.parameters() 是 nn.Module 类的一个方法，用于返回当前模块及其所有子模块中所有可学习参数（如权重、偏置等）的迭代器。
    1. 这些参数是模型在训练过程中需要通过反向传播更新的张量（如线性层的 weight 和 bias、卷积层的卷积核等）。通过该方法，优化器（如 torch.optim.Adam）可以便捷地获取并管理所有需要优化的参数，实现对模型的训练更新。
    2. 例如，在定义优化器时常用 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)，其中 model.parameters() 就为优化器提供了所有需更新的参数。
  • .dim() 是张量（Tensor）的一个方法，用于返回该张量的维度数量（即阶数）。例如：
    1. 标量（0 维张量）调用 .dim() 返回 0
    2. 向量（1 维张量）调用 .dim() 返回 1
    3. 矩阵（2 维张量张量）调用 .dim() 返回 2
    4. 形状为 (3, 4, 5) 的三维张量调用 .dim() 返回 3
  • 在神经网络参数初始化中，“只初始化矩阵参数（如权重），忽略偏置等1维参数”是常见做法，原因如下：
    1. 矩阵参数（权重）的重要性：权重矩阵（如线性层的weight）是模型学习特征映射的核心，其值直接影响输入与输出的非线性关系。若初始化不当（如全为0），会导致模型无法学习有效特征（对称权重问题），因此需要通过Xavier、Kaiming等方法精心初始化，确保信号在正向/反向传播中稳定传递。
    2. 1维参数（偏置）的特性：偏置（如线性层的bias）是单维度参数，作用是给输出添加一个常量偏移，不参与复杂的特征映射。即使初始化为0，也能通过梯度下降快速学习到合适的值，且不会引发对称权重等问题，因此常默认初始化为0而无需额外处理。
    3. 这种做法既保证了核心参数的有效初始化，又简化了流程，是兼顾效率与效果的实践选择。
  • nn.init.xavier_uniform_ 是 PyTorch 中一种参数初始化方法，用于初始化神经网络中的权重矩阵。
    1. 其核心思想是通过均匀分布初始化权重，使前向传播时输入信号的方差和反向传播时梯度的方差尽可能一致，避免因权重值过大或过小导致的信号衰减或爆炸问题。
    2. 具体实现上，它会根据权重矩阵的输入维度 fan_in 和输出维度 fan_out，在区间 [-a, a] 内生成均匀分布的随机数，其中 a = gain * sqrt(6 / (fan_in + fan_out))（gain 为增益系数，通常根据激活函数类型设置）。
    3. 该初始化方法适用于 tanh、sigmoid 等激活函数，能有效帮助深层网络稳定训练，其参数在训练过程中会通过反向传播进行更新。

10. 示例：创建Transformer实例并测试

if __name__ == "__main__":
    # 超参数设置
    input_vocab_size = 5000  # 输入词汇表大小（如源语言词典大小）
    target_vocab_size = 5000  # 目标词汇表大小（如目标语言词典大小）
    d_model = 512  # 模型维度（所有子层输出维度）
    num_layers = 10  # 编码器和解码器层数
    num_heads = 8  # 注意力头数（512/8=64，每个头维度64）
    d_ff = 2048  # 前馈网络隐藏层维度（通常为4*d_model）
    max_len = 100  # 最大序列长度（位置编码支持的最长序列）
    dropout = 0.1  # dropout率（防止过拟合）
    
    # 创建Transformer模型实例
    transformer = Transformer(
        input_vocab_size, target_vocab_size,
        d_model, num_layers, num_heads,
        d_ff, max_len, dropout
    )
    
    # 打印模型结构（可选）
    print(transformer)
    
    # 测试模型输入输出
    batch_size = 32  # 批次大小
    src_seq_len = 20  # 输入序列长度
    trg_seq_len = 25  # 目标序列长度
    
    # 随机生成输入（模拟词索引序列）
    src = torch.randint(0, input_vocab_size, (batch_size, src_seq_len))  # (32, 20)
    trg = torch.randint(0, target_vocab_size, (batch_size, trg_seq_len))  # (32, 25)
    
    # 创建简化掩码（实际应用中需根据具体任务设计，如PAD掩码、序列掩码）
    src_mask = torch.ones(batch_size, 1, 1, src_seq_len)  # 源序列掩码（全1表示无屏蔽）
    trg_mask = torch.ones(batch_size, 1, trg_seq_len, trg_seq_len)  # 目标序列掩码
    cross_mask = torch.ones(batch_size, 1, trg_seq_len, src_seq_len)  # 编码器-解码器掩码
    
    # 前向传播
    output = transformer(src, trg, src_mask, trg_mask, cross_mask)
    # 输出形状应为 (batch_size, trg_seq_len, target_vocab_size) → (32, 25, 5000)
    print(f"输入形状: {src.shape}")
    print(f"目标形状: {trg.shape}")
    print(f"输出形状: {output.shape}")

• 测试逻辑：通过随机生成的输入序列验证模型前向传播的正确性，确保输出形状符合预期。
• 掩码说明：示例中使用全1掩码（无屏蔽），实际任务中需根据填充符（PAD）和序列顺序设计真实掩码。
• 代码解释：
  • print(transformer) 会在控制台输出 Transformer 类实例的结构信息，包括其内部包含的所有子模块（如编码器、解码器、各层组件等）及其层次关系。
    1. 输出内容会按层级展示模型的组成，例如：
    2. 编码器（encoder）及其包含的嵌入层（embedding）、位置编码（pos_encoding）、10 个编码器层（layers 列表中的 EncoderLayer）
    3. 解码器（decoder）及其包含的对应组件和 10 个解码器层（DecoderLayer）
    4. 最终输出层（fc_out）
    5. 同时会显示各层的关键参数（如输入输出维度、头数等），帮助开发者直观了解模型的整体架构和模块构成，便于调试和确认模型结构是否符合设计预期。

总结

该代码完整实现了Transformer的核心结构，包括：

• 10层编码器和10层解码器
• 每个层包含残差连接（x + sublayer_output）和层归一化（nn.LayerNorm）
• 多头注意力、位置编码、前馈网络等关键组件

通过逐行解析可以看到，Transformer的设计核心是并行化处理（替代RNN的序列依赖）和注意力机制（捕捉长距离依赖），而残差连接和层归一化则是实现深层网络稳定训练的关键。

完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class PositionalEncoding(nn.Module):
    """位置编码，为输入序列添加位置信息"""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        # 计算位置编码
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数索引使用正弦函数
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数索引使用余弦函数
        
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)  # 非可学习参数
    
    def forward(self, x):
        # x: (seq_len, batch_size, d_model)
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class MultiHeadAttention(nn.Module):
    """多头注意力机制"""
    def __init__(self, d_model, num_heads, dropout=0.1):
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 每个头的维度
        
        # 线性变换层
        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)
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 线性变换并分成多头
        # (batch_size, seq_len, d_model) -> (batch_size, num_heads, seq_len, d_k)
        q = self.w_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        k = self.w_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        v = self.w_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # 应用掩码
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 注意力权重
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        
        # 应用注意力到值
        output = torch.matmul(attn, v)
        
        # 拼接多头结果
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        # 最终线性变换
        return self.w_o(output), attn

class PositionWiseFeedForward(nn.Module):
    """位置wise前馈网络"""
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(p=dropout)
        self.activation = nn.GELU()  # 使用GELU激活函数，比原始论文的ReLU更优
    
    def forward(self, x):
        # 前馈网络: Linear -> Activation -> Dropout -> Linear
        return self.fc2(self.dropout(self.activation(self.fc1(x))))

class EncoderLayer(nn.Module):
    """编码器层: 包含多头自注意力和前馈网络，以及残差连接和层归一化"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff, dropout)
        
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        #  dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
    
    def forward(self, x, mask):
        # 多头自注意力 + 残差连接 + 层归一化
        attn_output, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + self.dropout1(attn_output))  # 残差连接
        
        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout2(ff_output))  # 残差连接
        
        return x

class DecoderLayer(nn.Module):
    """解码器层: 包含掩码多头自注意力、编码器-解码器注意力和前馈网络"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)  # 掩码自注意力
        self.cross_attn = MultiHeadAttention(d_model, num_heads, 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)
        
        # dropout层
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
    
    def forward(self, x, enc_output, self_mask, cross_mask):
        # 掩码自注意力 + 残差连接 + 层归一化
        attn_output, _ = self.self_attn(x, x, x, self_mask)
        x = self.norm1(x + self.dropout1(attn_output))
        
        # 编码器-解码器注意力 + 残差连接 + 层归一化
        attn_output, _ = self.cross_attn(x, enc_output, enc_output, cross_mask)
        x = self.norm2(x + self.dropout2(attn_output))
        
        # 前馈网络 + 残差连接 + 层归一化
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout3(ff_output))
        
        return x

class Encoder(nn.Module):
    """编码器: 由多个编码器层堆叠而成"""
    def __init__(self, input_vocab_size, d_model, num_layers, num_heads, 
                 d_ff, max_len, dropout=0.1):
        super(Encoder, self).__init__()
        self.d_model = d_model
        
        # 嵌入层
        self.embedding = nn.Embedding(input_vocab_size, d_model)
        # 位置编码
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 10个编码器层
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x, mask):
        # 嵌入层 + 位置编码
        x = self.embedding(x) * math.sqrt(self.d_model)  # 缩放嵌入
        x = self.pos_encoding(x.transpose(0, 1)).transpose(0, 1)
        
        # 通过所有编码器层
        for layer in self.layers:
            x = layer(x, mask)
        
        return x

class Decoder(nn.Module):
    """解码器: 由多个解码器层堆叠而成"""
    def __init__(self, target_vocab_size, d_model, num_layers, num_heads, 
                 d_ff, max_len, dropout=0.1):
        super(Decoder, self).__init__()
        self.d_model = d_model
        
        # 嵌入层
        self.embedding = nn.Embedding(target_vocab_size, d_model)
        # 位置编码
        self.pos_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        # 10个解码器层
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout) 
            for _ in range(num_layers)
        ])
        
        self.dropout = nn.Dropout(p=dropout)
    
    def forward(self, x, enc_output, self_mask, cross_mask):
        # 嵌入层 + 位置编码
        x = self.embedding(x) * math.sqrt(self.d_model)  # 缩放嵌入
        x = self.pos_encoding(x.transpose(0, 1)).transpose(0, 1)
        
        # 通过所有解码器层
        for layer in self.layers:
            x = layer(x, enc_output, self_mask, cross_mask)
        
        return x

class Transformer(nn.Module):
    """完整的Transformer模型"""
    def __init__(self, input_vocab_size, target_vocab_size, 
                 d_model=512, num_layers=10, num_heads=8, 
                 d_ff=2048, max_len=5000, dropout=0.1):
        super(Transformer, self).__init__()
        
        # 编码器（10层）
        self.encoder = Encoder(
            input_vocab_size, d_model, num_layers, num_heads, d_ff, max_len, dropout
        )
        
        # 解码器（10层）
        self.decoder = Decoder(
            target_vocab_size, d_model, num_layers, num_heads, d_ff, max_len, dropout
        )
        
        # 最终输出层
        self.fc_out = nn.Linear(d_model, target_vocab_size)
        
        self.init_weights()  # 初始化权重
    
    def init_weights(self):
        """初始化模型权重"""
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)  # 使用Xavier均匀初始化
    
    def forward(self, src, trg, src_mask, trg_mask, cross_mask):
        # 编码器输出
        enc_output = self.encoder(src, src_mask)
        
        # 解码器输出
        dec_output = self.decoder(trg, enc_output, trg_mask, cross_mask)
        
        # 最终输出
        output = self.fc_out(dec_output)
        
        return output

# 示例：创建一个Transformer实例
if __name__ == "__main__":
    # 超参数设置
    input_vocab_size = 5000  # 输入词汇表大小
    target_vocab_size = 5000  # 目标词汇表大小
    d_model = 512  # 模型维度
    num_layers = 10  # 编码器和解码器层数
    num_heads = 8  # 注意力头数
    d_ff = 2048  # 前馈网络隐藏层维度
    max_len = 100  # 最大序列长度
    dropout = 0.1  # dropout率
    
    # 创建Transformer模型
    transformer = Transformer(
        input_vocab_size, target_vocab_size,
        d_model, num_layers, num_heads,
        d_ff, max_len, dropout
    )
    
    # 打印模型结构
    print(transformer)
    
    # 测试模型输入输出
    batch_size = 32
    src_seq_len = 20
    trg_seq_len = 25
    
    # 随机生成输入
    src = torch.randint(0, input_vocab_size, (batch_size, src_seq_len))
    trg = torch.randint(0, target_vocab_size, (batch_size, trg_seq_len))
    
    # 创建掩码（实际应用中需要根据具体需求创建）
    src_mask = torch.ones(batch_size, 1, 1, src_seq_len)  # 简化的源掩码
    trg_mask = torch.ones(batch_size, 1, trg_seq_len, trg_seq_len)  # 简化的目标掩码
    cross_mask = torch.ones(batch_size, 1, trg_seq_len, src_seq_len)  # 简化的交叉掩码
    
    # 前向传播
    output = transformer(src, trg, src_mask, trg_mask, cross_mask)
    print(f"输入形状: {src.shape}")
    print(f"目标形状: {trg.shape}")
    print(f"输出形状: {output.shape}")  # 应输出 (batch_size, trg_seq_len, target_vocab_size)