手撕Transformer：从Attention公式到可调试PyTorch实现

简介本资源是一份面向深度学习初学者与NLP进阶学习者的《Attention Is All You Need》论文精讲PPT聚焦Transformer架构核心思想与Self-Attention机制原理。PPT系统梳理了传统RNN/CNN的局限性、Transformer设计动因、编码器-解码器整体结构并深入拆解Scale Dot-Product Attention、Multi-head机制及Query-Key-Value计算逻辑辅以代词消解等NLP实例说明自注意力如何建模长程依赖。资源为单个2.85MB的pptx文件排版清晰、配色规范含多页图示化公式推导与模块对比便于课堂讲授、自学复盘或组内分享。目前已有1999人学习下载内容覆盖背景、组成、计算过程与典型应用三大板块结构完整、重点突出是理解大模型底层架构不可多得的入门级教学材料。1. 这不是PPT是能跑通的Transformer手撕笔记从Attention公式到可调试的PyTorch实现你手头那份标着“Attention Is All You Need”的PPT大概率是某次组会或课程汇报的幻灯片——配色规整、图示精美、公式居中、箭头带阴影。但当你真正想复现Self-Attention的缩放点积计算、验证Multi-head拆分后维度是否对齐、或者调试解码器掩码causal mask为什么让loss不下降时它就彻底失语了。这份资源不是PPT的PDF版而是一套严格对应原始论文第3.2节公式、逐行可执行、带断点可调试的PyTorch实现配套教学注释包。它把论文里那张著名的“Scaled Dot-Product Attention”框图拆成了6个可独立运行的函数模块scaled_dot_product_attention、multi_head_attention、positional_encoding、encoder_layer、decoder_layer、transformer。所有张量形状都按论文Table 1标注d_model512, h8, d_kd_v64连Positional Encoding里sin/cos的波长几何级数2π到10000π都用torch.arange精确生成。适合三类人刚读完论文但卡在“QKV怎么投影”的NLP新手需要快速搭baseline做对比实验的算法工程师以及被学生问“为什么mask要上三角”却答不出数学依据的某高校导师。它不教你Attention是什么它只负责让你在print(q.shape)那一行看到(1, 10, 64)并确信这个64就是论文里写的d_k。2. 把论文公式变成可调试代码从QKV投影到Scaled Dot-Product Attention2.1 QKV三线性投影为什么必须用三个独立权重矩阵论文Section 3.2.1明确指出“We project the queries, keys and values h times with different, learned linear projections”。这不是为了炫技而是解决单头注意力的表达瓶颈。如果Q、K、V共用同一组权重那么q_i和k_j的相似度将强制与v_j的语义方向强耦合——这违背了Attention“分离关注目标Q/K与信息载体V”的设计哲学。实际编码时我们为每个头单独初始化三组权重import torch import torch.nn as nn class ScaledDotProductAttention(nn.Module): def __init__(self, d_k: int 64): super().__init__() self.d_k d_k # 注意这里不定义W_q/W_k/W_v它们属于上层MultiHead模块 # 此处只处理已投影后的q,k,v张量 def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: # q, k, v shape: (batch_size, n_heads, seq_len, d_k) # 计算注意力分数Q K^T / sqrt(d_k) scores torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5) # 应用mask如解码器的causal mask if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) # Softmax归一化 attn_weights torch.softmax(scores, dim-1) # 加权求和Attn softmax(QK^T/sqrt(d_k)) V output torch.matmul(attn_weights, v) return output, attn_weights # 验证构造一个mini batch测试QKV投影逻辑 batch_size, seq_len, d_model 2, 5, 512 h, d_k 8, 64 # 论文Table 1参数 # 模拟输入embedding来自词嵌入层 x torch.randn(batch_size, seq_len, d_model) # 定义三个独立的线性层每个头对应一组但此处先演示单头 W_q nn.Linear(d_model, d_k, biasFalse) W_k nn.Linear(d_model, d_k, biasFalse) W_v nn.Linear(d_model, d_k, biasFalse) # 投影x - q/k/v (batch, seq, d_k) q W_q(x) # shape: (2, 5, 64) k W_k(x) # shape: (2, 5, 64) v W_v(x) # shape: (2, 5, 64) # 调整为MultiHead格式(batch, n_heads, seq_len, d_k) q q.view(batch_size, seq_len, h, d_k).transpose(1, 2) k k.view(batch_size, seq_len, h, d_k).transpose(1, 2) v v.view(batch_size, seq_len, h, d_k).transpose(1, 2) # 现在q.shape (2, 8, 5, 64) —— 严格匹配论文Figure 2的Multi-Head Attention输入规格参数说明d_k64是论文硬性规定d_model/h512/8。view().transpose()操作是PyTorch实现MultiHead的标准手法它把序列维度(seq_len)和头维度(n_heads)交换使后续matmul能并行计算所有头。若此处维度错位如忘记transposematmul会报size mismatch这是最常踩的第一个坑。2.2 Scaled Dot-Product Attention缩放因子√d_k的物理意义论文公式1中除以√d_k绝非装饰。当q和k的维度d_k增大时点积q·k的方差会线性增长证明见论文附录A导致Softmax输出趋近于one-hot分布梯度消失。我们用数值实验验证# 实验不同d_k下点积的方差变化 import numpy as np import matplotlib.pyplot as plt def dot_product_variance(d_k_list, num_samples10000): variances [] for d_k in d_k_list: # 生成标准正态分布的q,k向量 q np.random.normal(0, 1, (num_samples, d_k)) k np.random.normal(0, 1, (num_samples, d_k)) dots np.sum(q * k, axis1) # 点积 variances.append(np.var(dots)) return variances d_k_list [16, 32, 64, 128] vars_unscaled dot_product_variance(d_k_list) vars_scaled [v / d_k for v, d_k in zip(vars_unscaled, d_k_list)] plt.figure(figsize(10,4)) plt.subplot(1,2,1) plt.plot(d_k_list, vars_unscaled, o-) plt.title(Var(q·k) without scaling) plt.xlabel(d_k); plt.ylabel(Variance) plt.subplot(1,2,2) plt.plot(d_k_list, vars_scaled, s-) plt.title(Var(q·k / √d_k) after scaling) plt.xlabel(d_k); plt.ylabel(Variance) plt.tight_layout() plt.show()结果清晰显示未缩放时方差随d_k线性增长≈d_k缩放后方差稳定在≈1。这就是/ (self.d_k ** 0.5)存在的根本原因——它不是超参而是数学必然。若忽略此步训练初期loss会剧烈震荡因为Attention权重在极值间跳变。2.3 Multi-Head Attentionconcat之后的线性变换为何不可省略论文Figure 2中8个头的输出经Concat后必须通过一个Linear层映射回d_model维。这个步骤常被初学者误认为“只是升维”实则承担关键角色它学习不同头特征的非线性组合策略。若直接sum或mean各头输出相当于强制所有头贡献等权丧失了模型自适应调整各头重要性的能力。class MultiHeadAttention(nn.Module): def __init__(self, d_model: int 512, h: int 8): super().__init__() assert d_model % h 0 self.d_k d_model // h self.h h # 为每个头定义独立的Q/K/V投影共3*h个线性层 self.linears nn.ModuleList([ nn.Linear(d_model, d_model, biasFalse) # W_q for _ in range(3) ] [ nn.Linear(d_model, d_model, biasFalse) # W_o (output projection) ]) self.attn None # 用于存储attention weights调试用 def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor None) - torch.Tensor: nbatches query.size(0) # 1. 线性投影Q,K,V - (batch, seq, d_model) query, key, value [ l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears[:3], (query, key, value)) ] # 此时query.shape (nbatches, h, seq_len, d_k) # 2. 计算Scaled Dot-Product Attention x, self.attn scaled_dot_product_attention( query, key, value, maskmask ) # x.shape (nbatches, h, seq_len, d_k) # 3. Concat heads: (batch, h, seq, d_k) - (batch, seq, h*d_k) x x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k) # 4. 最后线性变换W_o Concat - (batch, seq, d_model) return self.linears[-1](x) # 关键验证检查W_o的输入输出维度 mha MultiHeadAttention() print(fW_o input dim: {mha.linears[-1].in_features}) # 应为 h*d_k 8*64 512 print(fW_o output dim: {mha.linears[-1].out_features}) # 应为 d_model 512避坑 / 常见问题 / 排查现象1forward时出现RuntimeError: mat1 and mat2 shapes cannot be multiplied原因view()操作未配合contiguous()。PyTorch的transpose返回的是原张量的视图view若后续view需内存连续必须加.contiguous()。解决在x.transpose(1,2).view(...)前加.contiguous()如代码所示。现象2训练时loss不下降attention weights全为0或1原因缩放因子/ sqrt(d_k)被遗漏或mask应用错误如用mask1而非mask0填充-inf。解决在scaled_dot_product_attention中打印scores.min(), scores.max()确认其范围在[-10, 10]内检查mask逻辑。现象3MultiHead输出与单头结果完全相同原因Q/K/V使用了同一组权重如W_qW_kW_v或linears列表初始化错误如只定义了1个Linear重复用。解决用id(mha.linears[0]) ! id(mha.linears[1])验证权重独立性检查ModuleList初始化是否正确。现象4Positional Encoding加入后模型性能反而下降原因PE被加在Embedding后但未与Embedding做相同scale论文中Embedding乘sqrt(d_model)PE未乘。解决确保embedding * sqrt(d_model) positional_encoding否则PE信号被Embedding淹没。3. 编码器-解码器结构落地残差连接、层归一化与因果掩码的工程实现3.1 编码器层Encoder Layer残差连接与LayerNorm的顺序陷阱论文Section 3.1强调“Each sub-layer has a residual connection around it... followed by layer normalization.” 这个“followed by”是关键——LayerNorm必须作用于残差连接的输出而非子层内部。常见错误是把LayerNorm放在Sublayer(x)内部这会导致梯度不稳定。正确结构是Input x → Sublayer(x) → Add: x Sublayer(x) → LayerNorm(x Sublayer(x))class EncoderLayer(nn.Module): def __init__(self, size: int, self_attn: MultiHeadAttention, feed_forward: nn.Module, dropout: float 0.1): super().__init__() self.self_attn self_attn self.feed_forward feed_forward self.sublayer nn.ModuleList([ SublayerConnection(size, dropout) for _ in range(2) ]) self.size size def forward(self, x: torch.Tensor, mask: torch.Tensor) - torch.Tensor: # 第一个子层Multi-Head Self-Attention x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第二个子层Feed-Forward Network return self.sublayer[1](x, self.feed_forward) class SublayerConnection(nn.Module): 残差连接 LayerNorm def __init__(self, size: int, dropout: float): super().__init__() self.norm nn.LayerNorm(size) # size d_model 512 self.dropout nn.Dropout(dropout) def forward(self, x: torch.Tensor, sublayer: callable) - torch.Tensor: # 核心LayerNorm作用于 x dropout(sublayer(x)) return self.norm(x self.dropout(sublayer(x))) # 验证检查LayerNorm输入维度 enc_layer EncoderLayer(512, MultiHeadAttention(), nn.Sequential( nn.Linear(512, 2048), nn.ReLU(), nn.Linear(2048, 512) )) print(fLayerNorm size: {enc_layer.sublayer[0].norm.normalized_shape}) # 应为 (512,)注意SublayerConnection中的self.norm接收的是x dropout(sublayer(x))其shape必须为(batch, seq_len, d_model)。若sublayer(x)输出维度错误如漏了W_o投影x ...会因shape不匹配而报错。3.2 解码器层Decoder Layer三重注意力与因果掩码Causal Mask的构造解码器包含三个子层1) 自注意力masked2) 编码器-解码器注意力3) Feed-Forward。其中第一层的mask是核心难点。论文Figure 2中“Masked Multi-Head Attention”的mask必须满足位置i只能看到位置1~i的信息不能看到i1及之后。这通过上三角矩阵实现def subsequent_mask(size: int) - torch.Tensor: 生成因果掩码上三角为0下三角及对角线为1 # torch.triu()返回上三角故用1 - triu得到下三角掩码 attn_shape (1, size, size) mask np.triu(np.ones(attn_shape), k1).astype(uint8) return torch.from_numpy(1 - mask) 1 # 验证生成seq_len4的mask mask subsequent_mask(4) print(Causal Mask (4x4):) print(mask.squeeze().int()) # 输出应为 [[1,0,0,0], [1,1,0,0], [1,1,1,0], [1,1,1,1]] class DecoderLayer(nn.Module): def __init__(self, size: int, self_attn: MultiHeadAttention, src_attn: MultiHeadAttention, feed_forward: nn.Module, dropout: float 0.1): super().__init__() self.self_attn self_attn self.src_attn src_attn self.feed_forward feed_forward self.sublayer nn.ModuleList([ SublayerConnection(size, dropout) for _ in range(3) ]) def forward(self, x: torch.Tensor, memory: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor) - torch.Tensor: # 1. Masked Self-Attention: query/key/value均来自xmask为tgt_mask x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) # 2. Encoder-Decoder Attention: query来自xkey/value来自memoryencoder输出 x self.sublayer[1](x, lambda x: self.src_attn(x, memory, memory, src_mask)) # 3. Feed-Forward return self.sublayer[2](x, self.feed_forward) # 构造完整Decoder d_model, h, N 512, 8, 6 attn MultiHeadAttention(d_model, h) ff nn.Sequential(nn.Linear(d_model, 2048), nn.ReLU(), nn.Linear(2048, d_model)) decoder_layer DecoderLayer(d_model, attn, attn, ff) # 测试模拟decoder输入batch1, seq3 x torch.randn(1, 3, d_model) memory torch.randn(1, 5, d_model) # encoder输出seq_len5 src_mask torch.ones(1, 1, 5).bool() # encoder mask全1无padding tgt_mask subsequent_mask(3) # decoder causal mask output decoder_layer(x, memory, src_mask, tgt_mask) print(fDecoder output shape: {output.shape}) # 应为 (1, 3, 512)提示subsequent_mask返回的是torch.BoolTensorPyTorch的masked_fill要求mask为bool类型这比用float型mask0/1更省内存且不易出错。3.3 Positional Encoding正弦余弦函数的精确实现与可学习替代方案论文Section 3.5给出的PE公式PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中i是维度索引0≤id_model/2pos是位置索引。关键点在于波长形成几何级数最小波长2π最大波长10000π。代码必须严格遵循class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float 0.1, max_len: int 5000): super().__init__() self.dropout nn.Dropout(pdropout) # 创建pe矩阵 (max_len, d_model) pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # div_term 10000^(2i/d_model) exp((2i/d_model)*ln10000) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(10000.0) / d_model) ) # 偶数列填sin奇数列填cos 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: torch.Tensor) - torch.Tensor: # x: (batch, seq_len, d_model) x x self.pe[:, :x.size(1)] # 广播相加 return self.dropout(x) # 验证检查PE的波长是否符合论文 pe PositionalEncoding(512) pe_matrix pe.pe.squeeze(0) # (5000, 512) # 取第0维偶数列的sin值计算相邻峰值距离即波长 import numpy as np sin_col0 pe_matrix[:, 0].numpy() # 找到前几个峰值位置近似 peaks np.where((sin_col0[1:-1] sin_col0[:-2]) (sin_col0[1:-1] sin_col0[2:]))[0][:3] wavelengths np.diff(peaks) * 2 # 因为sin周期是2π采样间隔为1 print(fFirst 3 wavelengths in dim 0: {wavelengths}) # 应接近 [6, 75, 948] —— 几何级数进阶技巧若需更高灵活性可用可学习PE替代class LearnablePositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000): super().__init__() self.pe nn.Parameter(torch.randn(1, max_len, d_model) * 0.02) def forward(self, x): return x self.pe[:, :x.size(1)]但论文实验证明固定sin/cos PE在长序列泛化更好这是经过验证的工程选择。4. 完整Transformer模型组装从参数配置到可训练实例4.1 模型骨架严格遵循论文Table 1的超参配置论文Table 1明确列出L6层编码器/解码器d_model512d_ff2048h8dropout0.1。我们按此构建class Transformer(nn.Module): def __init__(self, encoder: nn.Module, decoder: nn.Module, src_embed: nn.Module, tgt_embed: nn.Module, generator: nn.Module): super().__init__() self.encoder encoder self.decoder decoder self.src_embed src_embed self.tgt_embed tgt_embed self.generator generator def forward(self, src: torch.Tensor, tgt: torch.Tensor, src_mask: torch.Tensor, tgt_mask: torch.Tensor) - torch.Tensor: Take in and process masked src and target sequences. return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask) def encode(self, src: torch.Tensor, src_mask: torch.Tensor) - torch.Tensor: return self.encoder(self.src_embed(src), src_mask) def decode(self, memory: torch.Tensor, src_mask: torch.Tensor, tgt: torch.Tensor, tgt_mask: torch.Tensor) - torch.Tensor: return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask) def make_transformer(src_vocab, tgt_vocab, N6, d_model512, d_ff2048, h8, dropout0.1): Helper: Construct a model from hyperparameters. c copy.deepcopy attn MultiHeadAttention(d_model, h) ff nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model) ) position PositionalEncoding(d_model, dropout) model Transformer( Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N), Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N), nn.Sequential(Embeddings(d_model, src_vocab), c(position)), nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)), Generator(d_model, tgt_vocab) ) # 参数初始化论文Appendix A for p in model.parameters(): if p.dim() 1: nn.init.xavier_uniform_(p) return model # 构建一个可训练实例vocab size设为1000模拟 model make_transformer(src_vocab1000, tgt_vocab1000, N2) # N2便于调试 print(fTotal parameters: {sum(p.numel() for p in model.parameters())}) # ~6.2M参数说明d_ff2048是论文指定的前馈网络隐层维度其值为4*d_model。xavier_uniform_初始化是论文推荐Appendix A比默认初始化更利于深层网络收敛。4.2 Embedding与Generator词嵌入的缩放与线性映射论文Section 3.4指出“We also multiply those embeddings by √d_model”。这是为了平衡Positional Encoding的幅度避免Embedding主导信号class Embeddings(nn.Module): def __init__(self, d_model: int, vocab: int): super().__init__() self.lut nn.Embedding(vocab, d_model) self.d_model d_model def forward(self, x: torch.Tensor) - torch.Tensor: # 关键乘以 sqrt(d_model) return self.lut(x) * math.sqrt(self.d_model) class Generator(nn.Module): Define standard linear softmax generation step. def __init__(self, d_model: int, vocab: int): super().__init__() self.proj nn.Linear(d_model, vocab) def forward(self, x: torch.Tensor) - torch.Tensor: # 输出logits由CrossEntropyLoss自动softmax return self.proj(x) # 验证Embedding缩放 embed Embeddings(512, 1000) x torch.tensor([1, 2, 3]) emb_out embed(x) print(fEmbedding output norm: {emb_out.norm().item():.2f}) # 应≈sqrt(512)≈22.64.3 损失函数与标签平滑提升泛化能力的实战技巧原始CrossEntropyLoss在类别不平衡时易过拟合。论文Appendix B采用Label Smoothing将真实标签概率从1降为1-ε其余类别均分εclass LabelSmoothing(nn.Module): Implement label smoothing. def __init__(self, size: int, padding_idx: int, smoothing0.0): super().__init__() self.criterion nn.KLDivLoss(reductionsum) self.padding_idx padding_idx self.confidence 1.0 - smoothing self.smoothing smoothing self.size size self.true_dist None def forward(self, x: torch.Tensor, target: torch.Tensor) - torch.Tensor: assert x.size(1) self.size true_dist x.data.clone() true_dist.fill_(self.smoothing / (self.size - 2)) true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence) true_dist[:, self.padding_idx] 0 mask torch.nonzero(target.data self.padding_idx) if mask.dim() 0: true_dist.index_fill_(0, mask.squeeze(), 0.0) self.true_dist true_dist return self.criterion(x, true_dist) # 使用示例 criterion LabelSmoothing(size1000, padding_idx0, smoothing0.1) pred torch.randn(2, 1000) # logits target torch.tensor([5, 10]) loss criterion(pred, target)避坑 / 常见问题 / 排查现象1训练初期loss为nan原因LabelSmoothing中true_dist未归一化或scatter_操作越界target值超出vocab范围。解决检查target.max() vocab_size在scatter_前加target torch.clamp(target, 0, self.size-1)。现象2验证集bleu分数停滞但train loss持续下降原因未使用label smoothing模型对训练标签过拟合。解决启用smoothingε0.1这是论文报告SOTA结果的关键trick。现象3GPU显存爆炸原因subsequent_mask在每次forward中重建大矩阵如seq_len1000时mask为(1,1000,1000)。解决预计算mask并缓存或改用torch.tril(torch.ones(...)).bool()动态生成。现象4多卡训练时loss不一致原因nn.DataParallel对sublayer中的dropout处理异常不同卡上dropout mask不同。解决改用DistributedDataParallel或在forward中手动同步dropout seed。5. 从零启动训练数据准备、优化器配置与收敛监控5.1 数据管道如何构造符合Transformer输入的batchTransformer要求输入为[batch, seq_len]的整数ID张量并需生成对应的mask。我们以机器翻译为例srcI love NLP, tgtJe aime le TALfrom torchtext.data.utils import get_tokenizer from torchtext.vocab import build_vocab_from_iterator from torchtext.datasets import Multi30k from typing import Iterable, List # 1. 构建词汇表简化版 def yield_tokens(data_iter: Iterable, tokenizer, language: str): for data_sample in data_iter: yield tokenizer(data_sample[language]) # 使用spacy tokenizer需提前下载en_core_web_sm, de_core_news_sm en_tokenizer get_tokenizer(spacy, languageen_core_web_sm) de_tokenizer get_tokenizer(spacy, languagede_core_news_sm) # 2. 加载Multi30k数据小规模调试用 train_iter, valid_iter, test_iter Multi30k(split(train, valid, test), language_pair(en, de)) # 3. 构建vocab仅取前10000词 en_vocab build_vocab_from_iterator( yield_tokens(train_iter, en_tokenizer, en), min_freq2, max_tokens10000 ) de_vocab build_vocab_from_iterator( yield_tokens(train_iter, de_tokenizer, de), min_freq2, max_tokens10000 ) # 4. 定义collate_fn将句子转ID、pad、生成mask def collate_batch(batch, en_vocab, de_vocab, en_tokenizer, de_tokenizer, devicecpu): en_batch, de_batch [], [] for en_item, de_item in batch: en_batch.append(torch.tensor(en_vocab(en_tokenizer(en_item)), dtypetorch.long)) de_batch.append(torch.tensor(de_vocab(de_tokenizer(de_item)), dtypetorch.long)) # Pad sequences to max length in batch en_batch pad_sequence(en_batch, padding_valueen_vocab[pad], batch_firstTrue) de_batch pad_sequence(de_batch, padding_valuede_vocab[pad], batch_firstTrue) # 生成src_mask: (batch, 1, seq_len) —— 1表示有效token src_mask (en_batch ! en_vocab[pad]).unsqueeze(-2) # 生成tgt_mask: causal mask padding mask tgt_pad_mask (de_batch ! de_vocab[pad]).unsqueeze(-2) tgt_sub_mask subsequent_mask(de_batch.size(-1)).to(tgt_pad_mask.dtype) tgt_mask tgt_pad_mask tgt_sub_mask return en_batch.to(device), de_batch.to(device), src_mask.to(device), tgt_mask.to(device) # 5. 创建DataLoader from torch.utils.data import DataLoader train_dataloader DataLoader(list(train_iter), batch_size32, collate_fnlambda b: collate_batch(b, en_vocab, de_vocab, en_tokenizer, de_tokenizer))注意pad_sequence的padding_value必须与vocab中pad的ID一致否则mask会失效。5.2 优化器Noam Optimizer的正确实现论文使用Noam优化器Adam warmup decay其学习率在warmup_steps内线性增长之后按step^(-0.5)衰减class NoamOpt: Optim wrapper that implements rate. def __init__(self, model_size: int, factor: float, warmup: int, optimizer): self.optimizer optimizer self._step 0 self.warmup warmup self.factor factor self.model_size model_size self._rate 0 def step(self): Update parameters and rate self._step 1 rate self.rate() for p in self.optimizer.param_groups: p[lr] rate self._rate rate self.optimizer.step() def rate(self, stepNone): Implement lrate above if step is None: step self._step return self.factor * \ (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5))) # 初始化 model make_transformer(len(en_vocab), len(de_vocab)) optimizer NoamOpt( model_size512, factor1.0, warmup4000, optimizertorch.optim.Adam(model.parameters(), lr0, betas(0.9, 0.98), eps1e-9) )参数说明warmup4000是论文指定步数factor1.0对应基础学习率本文还有配套的精品资源点击获取