1. 从“黑盒”到“白盒”为什么我们需要从零开始理解Transformer如果你在过去几年里接触过任何与深度学习、自然语言处理或者计算机视觉相关的内容那么“Transformer”这个词对你来说一定不陌生。它从一个2017年论文里的模型架构迅速演变成了驱动当今几乎所有主流AI模型从ChatGPT到Stable Diffusion的核心引擎。网上充斥着各种“一文读懂Transformer”、“Transformer图解”的文章它们用精美的图示和类比试图让你在十分钟内理解这个复杂的概念。但作为一个真正想动手实践、甚至想基于此进行创新的开发者或研究者你是否有过这样的感觉看完那些文章好像懂了但让你自己从头写一个或者当模型输出奇怪结果时让你去调试大脑却一片空白这就是“从零开始”的价值所在。市面上大多数教程包括那篇著名的《The Illustrated Transformer》都更像是在给你展示一辆已经组装好的、光鲜亮丽的跑车告诉你引擎、变速箱、车轮各自是干什么的。这很有用能帮你建立宏观认知。但“从零开始”意味着我们要从一堆散落的零件开始自己动手把每一个螺丝拧到位把每一根线缆接正确。在这个过程中你会被迫去理解每一个设计决策背后的“为什么”为什么用自注意力而不是循环为什么要有残差连接和层归一化位置编码为什么那样设计这些“为什么”才是你真正掌握Transformer并能在其基础上进行修改、优化甚至发明新架构的关键。所以这篇内容不是另一篇“图解”。我们将使用PyTorch从最基础的张量操作开始一步步地、一行代码一行代码地构建出一个完整的、可训练的Transformer模型。我们会先搭建最核心的自注意力机制然后组装成编码器层和解码器层最后将它们整合成一个完整的序列到序列Seq2Seq模型并用一个简单的任务比如复制序列或机器翻译的玩具数据集来验证它的确能工作。我的目标是当你跟着走完这一程你不仅能“读懂”Transformer更能“写出”和“调通”一个Transformer。你会发现那些看似高深的概念拆解到底层无非是一些矩阵乘法和非线性变换的组合。2. 环境搭建与核心依赖避开版本依赖的“暗礁”在动手写代码之前一个稳定、兼容的环境是重中之重。根据网络上的高频问题环境配置尤其是PyTorch与CUDA版本的匹配是新手的第一道坎也是最容易让人沮丧的环节。我们这里会提供一个清晰、可复现的路径。2.1 核心工具选型为什么是PyTorch Anaconda我们的选择很明确PyTorch和Anaconda或Miniconda。这不是随大流而是基于实践效率的考量。PyTorch以其动态计算图和直观的Pythonic接口著称非常适合研究和原型开发。它的torch.nn.Module和torch.nn.functional模块设计得非常优雅让我们可以像搭积木一样构建模型。相比于TensorFlow早期的静态图PyTorch的调试体验要好得多——你可以随时打印中间变量的值这对于理解Transformer内部的数据流动至关重要。Anaconda/Miniconda这是管理Python环境和包依赖的“瑞士军刀”。深度学习项目常常对特定版本的库如PyTorch、CUDA工具包、NumPy有苛刻要求。Conda可以创建独立的虚拟环境完美隔离不同项目间的依赖冲突。网上很多“InvalidArchiveError”或“DLL load failed”错误根源就是环境混乱。2.2 一步步搭建无坑环境我们以Windows系统、配备NVIDIA GPU例如搜索词中提到的RTX 5060为例。如果你的显卡是Intel Arc或AMD或者使用Mac步骤在PyTorch安装环节有所不同但逻辑相通。步骤1安装Anaconda去Anaconda官网下载并安装最新的Individual Edition。安装时务必勾选“Add Anaconda to my PATH environment variable”添加Anaconda到系统PATH这能让你在命令行中直接使用conda命令。步骤2创建并激活专属环境打开Anaconda PromptWindows或终端Mac/Linux执行以下命令# 创建一个名为‘transformer_zoo’的Python 3.9环境3.8-3.11通常都兼容 conda create -n transformer_zoo python3.9 -y # 激活这个环境 conda activate transformer_zoo激活后你的命令行提示符前面应该会显示(transformer_zoo)表示你正在这个独立环境中操作。步骤3安装匹配的PyTorch最关键的一步这是最多坑的地方。核心原则是你的PyTorch版本必须与你的CUDA版本匹配而CUDA版本又需要你的NVIDIA显卡驱动支持。确定CUDA版本在命令行输入nvidia-smi。顶部会显示“CUDA Version: 12.4”之类的信息。这表示你的驱动最高支持CUDA 12.4。但你不一定需要安装这个版本的CUDA Toolkit。前往PyTorch官网获取安装命令打开 pytorch.org 你会看到一个配置器。根据你的情况选择Stable(稳定版)你的操作系统(Windows/Linux/Mac)Package(Conda通常更省心)Language(Python)Compute Platform(CUDA 12.1, 11.8, 或CPU)。这里的选择决定了你将安装的PyTorch预编译包绑定了哪个CUDA版本。我建议选择比nvidia-smi显示版本低一点的稳定版本例如驱动支持12.4可以安装CUDA 12.1的PyTorch兼容性更好。执行安装命令假设我们选择CUDA 12.1官网可能会给出conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia在激活的transformer_zoo环境中执行它。Conda会自动解决依赖安装PyTorch、TorchVision、TorchAudio以及对应的CUDA运行时库。注意如果你使用的是Intel Arc GPU目前PyTorch通过Intel Extension for PyTorch (IPEX) 提供优化支持。你需要安装CPU版本的PyTorch再额外安装intel-extension-for-pytorch。对于Jetson Orin Nano等嵌入式平台则需要按照NVIDIA官方提供的JetPack SDK中的说明安装为其ARM架构预编译的PyTorch wheel包强制安装不兼容的版本如搜索词中提到的torchvision0.20.0必然会导致RuntimeError。步骤4验证安装在Python交互环境中执行以下代码import torch print(torch.__version__) # 打印PyTorch版本 print(torch.cuda.is_available()) # 如果使用GPU应返回True print(torch.cuda.get_device_name(0)) # 打印你的GPU型号如果一切顺利你将看到版本号、True和你的显卡名称。步骤5安装其他辅助库我们还需要一些辅助工具conda install numpy matplotlib tqdm jupyter -y # 或者用pip pip install numpy matplotlib tqdm jupyter至此一个坚实的开发环境就搭建好了。这个环境是后续所有代码运行的基础。3. 自注意力机制Transformer的灵魂与数学本质现在我们进入核心。抛开一切外部包装Transformer最革命性的创新就是自注意力机制。它解决了RNN和LSTM的序列建模痛点难以并行化、长程依赖衰减。我们来亲手实现它。3.1 直觉理解从“查字典”到“动态词典”想象一下你在读一句话“The animal didnt cross the street because it was too tired.” 这里的“it”指代什么是“animal”还是“street”人类会根据语义将注意力更多地放在“animal”上。自注意力机制就是让模型学会这种能力。在模型中句子中的每个词或更一般地每个“token”都会被转换成三个向量查询向量Query、键向量Key和值向量Value。你可以这样理解Query (Q)当前词发出的“提问”。比如“it”在问“谁是我指代的对象”Key (K)序列中每个词提供的“标签”或“标识”。比如“animal”的标签是“名词、主语、生物”“street”的标签是“名词、宾语、地点”。Value (V)每个词所代表的实际“信息”或“内容”。自注意力就是让“it”的Query去和所有词的Key进行匹配计算相似度得到一个权重分数。这个分数决定了在合成“it”的新表示时每个词的Value应该贡献多少。显然“animal”的Key应该与“it”的Query更匹配因此权重更高“animal”的Value对“it”新表示的贡献就更大。3.2 数学过程与代码实现让我们用PyTorch实现一个最基础的单头自注意力。第一步生成Q, K, V假设输入序列的维度是(batch_size, seq_len, d_model)其中d_model是模型的特征维度例如512。我们通过三个不同的线性变换层nn.Linear来生成Q, K, V。import torch import torch.nn as nn import torch.nn.functional as F import math class SelfAttention(nn.Module): def __init__(self, d_model): super().__init__() self.d_model d_model # 通常Q, K, V的维度等于 d_model 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) # 值变换 def forward(self, x): # x shape: (batch_size, seq_len, d_model) batch_size, seq_len, d_model x.shape Q self.w_q(x) # (batch_size, seq_len, d_model) K self.w_k(x) # (batch_size, seq_len, d_model) V self.w_v(x) # (batch_size, seq_len, d_model)第二步计算注意力分数缩放点积注意力注意力分数的计算公式是Attention(Q, K, V) softmax(Q * K^T / sqrt(d_k)) * V这里d_k是Key向量的维度通常等于d_model。除以sqrt(d_k)是为了在d_k较大时防止点积结果过大导致softmax梯度消失。def forward(self, x): # ... 生成Q, K, V的代码同上 ... # 计算 Q * K^T # 我们需要将K转置以便进行矩阵乘法。维度变化 # Q: (batch, seq_len_q, d_model) # K: (batch, seq_len_k, d_model) - 转置后两个维度: (batch, d_model, seq_len_k) # 相乘后得到: (batch, seq_len_q, seq_len_k) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_model) # 应用softmax得到注意力权重维度不变 (batch, seq_len_q, seq_len_k) attn_weights F.softmax(scores, dim-1) # 在最后一个维度(seq_len_k)上做softmax # 用注意力权重加权求和Value向量 # attn_weights: (batch, seq_len_q, seq_len_k) # V: (batch, seq_len_v, d_model) 注意seq_len_k 必须等于 seq_len_v # 相乘后得到: (batch, seq_len_q, d_model) output torch.matmul(attn_weights, V) return output, attn_weights # 返回输出和注意力权重用于可视化这就是最核心的自注意力计算。它允许序列中的任意两个位置直接交互无论它们相距多远且整个计算过程是高度可并行化的矩阵运算。实操心得在调试时务必关注张量的维度。torch.matmul在不同维度下的行为不同。使用.transpose(-2, -1)是交换最后两个维度的安全做法。打印中间变量如scores,attn_weights的shape是排查维度错误的最有效手段。3.3 多头注意力为什么需要多个“视角”单头注意力就像只用一种标准去衡量词与词之间的关系。多头注意力则是并行地运行多个独立的注意力“头”每个头都有自己的Q、K、V变换矩阵可以学习到在不同子空间subspace下的不同依赖关系。例如一个头可能专注于语法指代另一个头可能专注于情感关联。实现上我们不是将d_model直接用于单个头而是将其分割成h个头每个头的维度是d_k d_v d_model / h。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads 0, “d_model must be divisible by num_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) # 输出维度是 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 forward(self, x): batch_size, seq_len, d_model x.shape # 1. 线性投影并分头 Q self.w_q(x) # (batch, seq_len, d_model) K self.w_k(x) V self.w_v(x) # 重塑张量将“头”的维度分离出来 # 目标形状: (batch, num_heads, seq_len, d_k) Q Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 现在 Q, K, V 形状: (batch, num_heads, seq_len, d_k) # 2. 在每个头上独立计算缩放点积注意力 # 我们复用之前的SelfAttention计算逻辑但需要处理多出来的“头”维度 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # (batch, num_heads, seq_len, seq_len) attn_weights F.softmax(scores, dim-1) attn_output torch.matmul(attn_weights, V) # (batch, num_heads, seq_len, d_k) # 3. 合并多头 # 将“头”的维度移回并拼接 attn_output attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model) # contiguous()确保内存连续view操作更安全 # 4. 最终输出投影 output self.w_o(attn_output) # (batch, seq_len, d_model) return output, attn_weights通过多头机制模型能够同时关注来自不同表示子空间的信息增强了其表达能力。4. 位置编码与层归一化为无序的注意力注入顺序与稳定自注意力机制本身是“排列不变”的即打乱输入序列的顺序输出只是相应位置被打乱但每个位置的内容不变。这显然不符合语言等有序序列的需求。同时深度网络训练中的梯度问题也需要被解决。4.1 位置编码告诉模型“谁在前谁在后”Transformer使用位置编码来将序列的顺序信息注入到输入嵌入中。最经典的是使用正弦和余弦函数PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是位置索引i是维度索引。这种编码的特点是对于固定的偏移量kPE(posk)可以表示为PE(pos)的线性函数这使得模型能够轻松学习到相对位置关系。class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() self.dropout nn.Dropout(pdropout) # 创建位置编码矩阵 (max_len, d_model) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # (max_len, 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) # (1, max_len, d_model) 方便广播 self.register_buffer(‘pe’, pe) # 将其注册为缓冲区不参与训练但会随模型保存/加载 def forward(self, x): # x shape: (batch_size, seq_len, d_model) x x self.pe[:, :x.size(1), :] # 只取前seq_len个位置 return self.dropout(x)注意位置编码是在输入嵌入词向量之后相加而不是拼接。register_buffer是关键它告诉PyTorch这个张量是模型状态的一部分但不属于可训练参数parameters。4.2 层归一化与残差连接训练深度模型的“稳定器”Transformer的每个子层自注意力层、前馈网络层都遵循一个相同的模式LayerNorm(x Sublayer(x))。这就是残差连接和层归一化。残差连接将子层的输入直接加到其输出上x Sublayer(x)。这解决了深度网络中的梯度消失问题使得网络可以轻松地学习一个恒等映射确保至少不会比浅层网络更差。层归一化对单个样本的所有特征维度进行归一化与批归一化BN不同BN是在批次维度上。它稳定了激活值的分布加速训练收敛。公式为LayerNorm(x) γ * (x - μ) / √(σ² ε) β其中γ和β是可学习的缩放和偏移参数。class LayerNorm(nn.Module): “”“简化的层归一化实现”“” def __init__(self, features, eps1e-6): super().__init__() self.gamma nn.Parameter(torch.ones(features)) # 缩放参数 self.beta nn.Parameter(torch.zeros(features)) # 偏移参数 self.eps eps def forward(self, x): # x shape: (batch, seq_len, d_model) 或 (batch, d_model) mean x.mean(-1, keepdimTrue) # 在最后一个维度上求均值 std x.std(-1, keepdimTrue, unbiasedFalse) # 求标准差 return self.gamma * (x - mean) / (std self.eps) self.beta在实际中我们直接使用PyTorch内置的nn.LayerNorm它经过高度优化。子层连接包装器 为了方便我们创建一个通用的子层连接模块。class SublayerConnection(nn.Module): “”“残差连接 层归一化。注意先做归一化再经过子层是原始Transformer论文的做法Pre-LN。也有后置归一化Post-LN的变体。”“” def __init__(self, size, dropout): super().__init__() self.norm nn.LayerNorm(size) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer): “”“sublayer是一个函数例如一个自注意力层或前馈网络层。”“” # Pre-LN: 先归一化再经过子层再加残差 return x self.dropout(sublayer(self.norm(x)))5. 前馈网络与编码器层注意力之后的非线性变换自注意力层负责聚合信息而前馈网络则负责对每个位置的表示进行独立的、复杂的非线性变换。它是一个简单的两层全连接网络中间有一个ReLU激活函数。class PositionwiseFeedForward(nn.Module): “”“FFN(x) max(0, xW1 b1)W2 b2”“” def __init__(self, d_model, d_ff, dropout0.1): super().__init__() self.w_1 nn.Linear(d_model, d_ff) # 第一层通常扩大维度如 d_ff2048 self.w_2 nn.Linear(d_ff, d_model) # 第二层投影回 d_model self.dropout nn.Dropout(dropout) self.activation nn.ReLU() # 或者 GELU def forward(self, x): return self.w_2(self.dropout(self.activation(self.w_1(x))))现在我们可以组装一个完整的编码器层了。一个编码器层包含两个子层多头自注意力层和前馈网络层每个子层都被残差连接和层归一化包裹。class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() self.self_attn MultiHeadAttention(d_model, num_heads) self.feed_forward PositionwiseFeedForward(d_model, d_ff, dropout) self.sublayer nn.ModuleList([SublayerConnection(d_model, dropout) for _ in range(2)]) # 两个子层连接 def forward(self, x): # 第一个子层自注意力 x self.sublayer[0](x, lambda x: self.self_attn(x)[0]) # 我们只取注意力输出的第一个值 # 第二个子层前馈网络 x self.sublayer[1](x, self.feed_forward) return x6. 解码器与掩码生成任务的关键约束解码器负责根据编码器的输出和已生成的部分序列预测下一个词。它与编码器结构相似但有两个关键区别掩码多头注意力在解码器的第一个注意力层为了防止模型在训练时“偷看”未来的答案即当前位置只能关注它之前的位置需要使用注意力掩码。编码器-解码器注意力解码器的第二个注意力层其Query来自解码器自身而Key和Value来自编码器的最终输出。这允许解码器在生成每个词时有选择地关注输入序列的不同部分。6.1 实现掩码注意力我们需要一个掩码矩阵其形状为(seq_len, seq_len)。对于位置i它只能关注位置j i。通常将未来位置的值设为负无穷-1e9这样在softmax之后权重就为0。def subsequent_mask(size): “”“生成一个向后的掩码。用于解码器的自回归训练。”“” attn_shape (1, size, size) subsequent_mask torch.triu(torch.ones(attn_shape), diagonal1).bool() # triu返回上三角矩阵diagonal1表示不包括对角线。未来位置为True需要被掩码 return ~subsequent_mask # 取反未来位置为False过去和当前位置为True # 使用示例 mask subsequent_mask(5) print(mask) # tensor([[[ True, False, False, False, False], # [ True, True, False, False, False], # [ True, True, True, False, False], # [ True, True, True, True, False], # [ True, True, True, True, True]]])在多头注意力计算中我们需要将这个掩码加到注意力分数上在softmax之前# 在MultiHeadAttention.forward中计算scores之后 if mask is not None: # mask shape: (batch, 1, seq_len_q, seq_len_k) 或 (1, seq_len_q, seq_len_k) # 将需要掩码的位置maskFalse设置为一个很大的负数 scores scores.masked_fill(mask 0, -1e9) attn_weights F.softmax(scores, dim-1)6.2 构建解码器层class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super().__init__() # 三个子层 self.self_attn MultiHeadAttention(d_model, num_heads) self.cross_attn MultiHeadAttention(d_model, num_heads) # 编码器-解码器注意力 self.feed_forward PositionwiseFeedForward(d_model, d_ff, dropout) self.sublayer nn.ModuleList([SublayerConnection(d_model, dropout) for _ in range(3)]) def forward(self, x, memory, src_mask, tgt_mask): “”“ x: 解码器输入 (batch, tgt_seq_len, d_model) memory: 编码器输出 (batch, src_seq_len, d_model) src_mask: 源序列掩码用于padding tgt_mask: 目标序列掩码用于自回归 ”“” # 第一子层掩码自注意力 x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)[0]) # 第二子层编码器-解码器注意力。Query来自解码器Key和Value来自编码器输出(memory) x self.sublayer[1](x, lambda x: self.cross_attn(x, memory, memory, src_mask)[0]) # 第三子层前馈网络 x self.sublayer[2](x, self.feed_forward) return x7. 组装完整Transformer与嵌入层现在我们将编码器层堆叠N次形成编码器解码器层堆叠N次形成解码器。同时我们需要嵌入层将输入的词索引整数转换为稠密向量。class Embeddings(nn.Module): def __init__(self, vocab_size, d_model): super().__init__() self.lut nn.Embedding(vocab_size, d_model) self.d_model d_model def forward(self, x): # 根据原始论文嵌入值需要乘以 sqrt(d_model) return self.lut(x) * math.sqrt(self.d_model) class Encoder(nn.Module): def __init__(self, layer, N): super().__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) # 最后的层归一化 def forward(self, x, mask): for layer in self.layers: x layer(x, mask) return self.norm(x) class Decoder(nn.Module): def __init__(self, layer, N): super().__init__() self.layers nn.ModuleList([copy.deepcopy(layer) for _ in range(N)]) self.norm nn.LayerNorm(layer.size) def forward(self, x, memory, src_mask, tgt_mask): for layer in self.layers: x layer(x, memory, src_mask, tgt_mask) return self.norm(x)最后组装完整的Transformer模型import copy class Transformer(nn.Module): def __init__(self, src_vocab, tgt_vocab, N6, d_model512, d_ff2048, num_heads8, dropout0.1): super().__init__() self.encoder_embed nn.Sequential(Embeddings(src_vocab, d_model), PositionalEncoding(d_model, dropout)) self.decoder_embed nn.Sequential(Embeddings(tgt_vocab, d_model), PositionalEncoding(d_model, dropout)) encoder_layer EncoderLayer(d_model, num_heads, d_ff, dropout) decoder_layer DecoderLayer(d_model, num_heads, d_ff, dropout) self.encoder Encoder(encoder_layer, N) self.decoder Decoder(decoder_layer, N) self.generator nn.Linear(d_model, tgt_vocab) # 最后的线性层softmax生成词概率 # 参数初始化很重要 for p in self.parameters(): if p.dim() 1: nn.init.xavier_uniform_(p) def forward(self, src, tgt, src_mask, tgt_mask): # src, tgt: (batch, seq_len) memory self.encode(src, src_mask) output self.decode(tgt, memory, src_mask, tgt_mask) return self.generator(output) def encode(self, src, src_mask): src_emb self.encoder_embed(src) return self.encoder(src_emb, src_mask) def decode(self, tgt, memory, src_mask, tgt_mask): tgt_emb self.decoder_embed(tgt) return self.decoder(tgt_emb, memory, src_mask, tgt_mask)8. 训练与验证让模型真正“动”起来构建模型只是第一步训练它并验证其有效性才是最终目标。我们设计一个简单的任务序列复制。即让模型学会将输入序列原样输出。这虽然简单但足以验证模型的前向传播、反向传播和基本学习能力。8.1 数据准备与掩码生成我们生成随机的整数序列作为数据。def data_gen(batch_size, seq_len, vocab_size): “”“生成随机数据”“” src torch.randint(1, vocab_size, (batch_size, seq_len)) # 假设0是填充符 tgt src.clone() # 目标就是复制源序列 # 对于复制任务我们通常将目标序列向右移动一位并在开头添加起始符 # 这里为了简化我们直接让模型学习映射 src - tgt return src, tgt def create_mask(src, tgt, pad_idx0): “”“创建源掩码和目标掩码”“” src_mask (src ! pad_idx).unsqueeze(-2) # (batch, 1, src_len) tgt_mask (tgt ! pad_idx).unsqueeze(-2) # padding mask # 结合自回归掩码 seq_len tgt.size(-1) subsequent_m subsequent_mask(seq_len).to(tgt.device) tgt_mask tgt_mask subsequent_m # 逻辑与既要非填充又要不能看未来 return src_mask, tgt_mask8.2 训练循环与损失函数我们使用标准的交叉熵损失和Adam优化器。def train_simple_copy_task(): vocab_size 50 d_model 64 # 为了快速演示使用小模型 model Transformer(src_vocabvocab_size, tgt_vocabvocab_size, N2, d_modeld_model, d_ff256, num_heads4) criterion nn.CrossEntropyLoss(ignore_index0) # 忽略填充符的损失 optimizer torch.optim.Adam(model.parameters(), lr0.0001, betas(0.9, 0.98), eps1e-9) model.train() for epoch in range(20): total_loss 0 for _ in range(100): # 每个epoch 100个batch src, tgt data_gen(batch_size32, seq_len10, vocab_sizevocab_size) # 目标输入是tgt目标输出是tgt对于复制任务 tgt_input tgt[:, :-1] # 解码器输入去掉最后一个词 tgt_output tgt[:, 1:] # 解码器目标去掉第一个词起始符 # 在实际任务中需要在tgt_input前加起始符这里简化处理 src_mask, tgt_mask create_mask(src, tgt_input) optimizer.zero_grad() out model(src, tgt_input, src_mask, tgt_mask) # out: (batch, seq_len-1, vocab_size) loss criterion(out.contiguous().view(-1, vocab_size), tgt_output.contiguous().view(-1)) loss.backward() optimizer.step() total_loss loss.item() print(f“Epoch {epoch}, Loss: {total_loss / 100:.4f}”) # 简单推理测试 model.eval() with torch.no_grad(): test_src torch.randint(1, vocab_size, (1, 5)) print(f“Input: {test_src}”) # 贪婪解码简化版实际应用需要beam search memory model.encode(test_src, (test_src ! 0).unsqueeze(-2)) ys torch.ones(1, 1).fill_(1).long() # 起始符假设为1 for i in range(4): # 生成4个词 tgt_mask subsequent_mask(ys.size(1)).type_as(test_src.data) out model.decode(ys, memory, None, tgt_mask) prob model.generator(out[:, -1]) # 取最后一个位置的输出 _, next_word torch.max(prob, dim-1) ys torch.cat([ys, next_word.unsqueeze(0)], dim1) print(f“Predicted: {ys}”)运行这段代码你会看到损失在下降。虽然这个复制任务很简单但它证明了我们从头搭建的Transformer模型骨架是正确且可训练的。在这个过程中你亲手触碰了自注意力、位置编码、层归一化、残差连接、掩码机制等每一个核心部件。现在你可以自信地说你理解了Transformer的“五脏六腑”而不仅仅是看过它的“外观”。这为你后续理解更复杂的变体如BERT的编码器、GPT的解码器、Vision Transformer等以及进行自己的修改和实验打下了最坚实的基础。
