简介这份资源面向希望深入理解神经机器翻译与Transformer架构的深度学习学习者与研究者提供一套基于PyTorch的完整实现代码。内容围绕编码器-解码器结构展开涵盖多头注意力、位置编码、前馈神经网络等核心模块并包含数据预处理、词表构建、模型训练与评估的完整流程适合具备一定PyTorch基础、想动手复现NMT模型的读者。压缩包共59个文件以53个Python脚本为主辅以C源码、编译产物及fast_align对齐工具整体约468KB目录中可见模型定义、优化器、指标计算与配置等模块结构清晰便于按功能查阅。目前已有1788人学习下载可作为理解Transformer内部机制、调试训练流程与二次开发的实践参考。1. 从零搭一套 PyTorch Transformer 神经机器翻译为什么值得亲手写一遍很多人第一次接触神经机器翻译是从nn.Transformer那一行调用开始的跑通了却说不清数据怎么对齐、mask 怎么加、推理为什么和训练不一样。真到要改结构、换分词、上自己的语料时立刻卡住。这个标题要解决的就是这件事用 PyTorch 从数据管道到训练、推理完整实现一套基于 Transformer 的神经机器翻译系统而不是调个库看 loss 下降。它适合已经会写基础 PyTorch 训练循环、想真正吃透 Transformer 在 seq2seq 任务里怎么落地的人。下面按我实际搭过几版的顺序讲参数、坑、验证方法都给到能直接抄的程度。2. 数据管道与词表Transformer 翻译的地基怎么打Transformer 不像 RNN 那样天然按时间步吃序列它对数据的整齐度要求更高同一批里所有句子必须 pad 到等长还要额外维护 padding mask 和 causal mask。数据这一层没搭好后面训练再调参都是玄学。2.1 平行语料的读取与清洗神经机器翻译的标准输入是平行语料一行源语言对一行目标语言。常见做法是先做长度过滤和去重把源/目标长度比过大的句子丢掉这类句子多半是语料对齐错误留着只会让模型学噪声。def load_parallel(src_path, tgt_path, max_len100, max_ratio2.0): pairs [] with open(src_path, encodingutf-8) as fs, open(tgt_path, encodingutf-8) as ft: for s, t in zip(fs, ft): s, t s.strip(), t.strip() if not s or not t: continue # 长度过滤过长句子显存吃不消比例失衡多为对齐错误 if len(s.split()) max_len or len(t.split()) max_len: continue ratio len(s.split()) / max(len(t.split()), 1) if ratio max_ratio or ratio 1 / max_ratio: continue pairs.append((s, t)) # 去重避免重复句主导梯度 return list(dict.fromkeys(pairs))max_len控制单句词数上限直接决定后面 batch 的显存占用max_ratio是源目标长度比阈值1.5 到 2.5 之间都算合理语料质量差就收紧到 1.5。清洗完记得打印一下保留率掉到 70% 以下说明阈值太狠或者语料本身有问题。2.2 用 BPE 还是词级词表选型与实现词级词表实现简单但 OOV 多、词表爆炸BPE 子词能把罕见词拆成子词单元是现在翻译任务的主流。我一般用 sentencepiece 或 tokenizers 训一个共享或独立的 BPE词表大小 8k 到 32k 之间。词表太小句子被切太碎序列变长太大又退化成词级。import sentencepiece as spm # 训练源语言 BPEvocab_size 是核心参数 spm.SentencePieceTrainer.train( inputtrain.src, model_prefixspm_src, vocab_size16000, model_typebpe, character_coverage1.0, pad_id0, unk_id1, bos_id2, eos_id3, ) sp spm.SentencePieceProcessor(model_filespm_src.model) ids sp.encode(this is a test, out_typeint) # 含 bos/eosvocab_size是这里最需要调的参数中英这类差异大的语言对建议源目标各训一个character_coverage对中文要设 1.0否则生僻字全变 unk。注意 pad/unk/bos/eos 的 id 要固定下来后面 mask 全靠这几个特殊符号。2.3 Dataset 与动态 padding 的 collate 函数固定长度 padding 浪费算力动态 padding 按 batch 内最长句补齐更划算。collate 函数负责把一批样本 pad 齐并生成后续要用的 mask。import torch from torch.nn.utils.rnn import pad_sequence PAD, BOS, EOS 0, 2, 3 class TranslationDataset(torch.utils.data.Dataset): def __init__(self, pairs, sp_src, sp_tgt): self.pairs pairs self.sp_src, self.sp_tgt sp_src, sp_tgt def __len__(self): return len(self.pairs) def __getitem__(self, idx): s, t self.pairs[idx] src self.sp_src.encode(s, out_typeint) # 已含 bos/eos tgt self.sp_tgt.encode(t, out_typeint) return torch.tensor(src), torch.tensor(tgt) def collate(batch): srcs, tgts zip(*batch) src pad_sequence(srcs, batch_firstTrue, padding_valuePAD) tgt pad_sequence(tgts, batch_firstTrue, padding_valuePAD) return src, tgtbatch_firstTrue让张量形状是[B, L]和nn.Transformer默认的[L, B, D]不同后面进模型前要转置这点很容易翻车。pad 值必须和词表里的 pad_id 一致否则 mask 会算错。3. 模型结构手写编码器解码器还是直接用 nn.Transformer选型上分两派直接用nn.Transformer省事手写 Encoder/Decoder 层能改结构、看得清每个 mask。想真正理解翻译任务里 mask 怎么配合我建议至少手写一遍再切回官方实现。3.1 位置编码与词嵌入的缩放关系Transformer 没有循环结构位置信息全靠位置编码注入。原始论文用正弦编码也可以换成可学习的位置嵌入。关键细节是词嵌入要乘sqrt(d_model)否则和位置编码量级不匹配训练初期梯度会不稳。import math import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len).unsqueeze(1).float() div torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) pe[:, 1::2] torch.cos(pos * div) self.register_buffer(pe, pe.unsqueeze(0)) # [1, max_len, d_model] self.dropout nn.Dropout(dropout) def forward(self, x): # x: [B, L, D]词嵌入已乘 sqrt(d_model) return self.dropout(x self.pe[:, :x.size(1)])d_model要和词嵌入维度一致常见 256/512max_len要大于你语料里最长句超了会索引越界。register_buffer保证 pe 跟着模型上 GPU 但不参与梯度。3.2 三种 mask 的生成padding、causal、cross翻译任务里 mask 是最容易出错的地方。编码器只需要 padding mask解码器自注意力要 padding causal交叉注意力用编码器的 padding mask。搞混一个模型要么看到未来词要么把 pad 当有效词。def make_pad_mask(seq, pad_id0): # seq: [B, L] - [B, 1, 1, L]True 表示被 mask return (seq pad_id).unsqueeze(1).unsqueeze(2) def make_causal_mask(size, device): # 上三角为 True屏蔽未来位置 return torch.triu(torch.ones(size, size, devicedevice), diagonal1).bool() def make_tgt_mask(tgt, pad_id0): pad_mask make_pad_mask(tgt, pad_id) # [B,1,1,L] causal make_causal_mask(tgt.size(1), tgt.device) # [L,L] return pad_mask | causal # 广播成 [B,1,L,L]nn.Transformer的src_mask形状是[L, L]或[B, L, L]tgt_mask同理src_key_padding_mask是[B, L]。不同版本 PyTorch 对 mask 语义True 是被屏蔽还是保留有过调整务必对着你装的版本确认这是血泪经验。3.3 组装完整模型与输出层输出层是一个线性层把d_model映射到目标词表大小配合交叉熵。注意解码器输出要经过这个线性层且和词嵌入是否共享权重可以作为一个调优点。class Seq2SeqTransformer(nn.Module): def __init__(self, src_vocab, tgt_vocab, d_model512, nhead8, num_enc6, num_dec6, dim_ff2048, dropout0.1): super().__init__() self.src_emb nn.Embedding(src_vocab, d_model, padding_idx0) self.tgt_emb nn.Embedding(tgt_vocab, d_model, padding_idx0) self.pos PositionalEncoding(d_model, dropoutdropout) self.transformer nn.Transformer( d_modeld_model, nheadnhead, num_encoder_layersnum_enc, num_decoder_layersnum_dec, dim_feedforwarddim_ff, dropoutdropout, batch_firstTrue) self.fc nn.Linear(d_model, tgt_vocab) self.d_model d_model def forward(self, src, tgt, src_pad_mask, tgt_mask): src self.pos(self.src_emb(src) * math.sqrt(self.d_model)) tgt self.pos(self.tgt_emb(tgt) * math.sqrt(self.d_model)) out self.transformer(src, tgt, tgt_masktgt_mask, src_key_padding_masksrc_pad_mask.squeeze(1).squeeze(1), tgt_key_padding_mask(tgt 0)) return self.fc(out)nhead必须能整除d_model512/8 是经典组合dim_ff一般是d_model的 4 倍。batch_firstTrue时输入是[B, L, D]和前面 collate 对齐省掉转置。padding_idx0让 pad 的词嵌入不更新。4. 训练循环标签平滑、学习率预热与梯度裁剪翻译任务的训练有几个和普通分类不一样的地方目标序列要错位一位做 teacher forcing损失常用标签平滑学习率要预热。这几项配错loss 会卡住或者震荡。4.1 目标错位与损失计算解码器输入是bos ...预测目标是... eos也就是输入去掉最后一位、目标去掉第一位。交叉熵要忽略 pad 位置。import torch.nn.functional as F def compute_loss(logits, tgt, pad_id0, smoothing0.1): # logits: [B, L, V]tgt: [B, L] logits logits[:, :-1].reshape(-1, logits.size(-1)) gold tgt[:, 1:].reshape(-1) return F.cross_entropy(logits, gold, ignore_indexpad_id, label_smoothingsmoothing)label_smoothing取 0.1 是翻译任务的常见值能缓解过拟合、提升 BLEUignore_indexpad_id保证 pad 不贡献梯度。错位这一步写错模型会学成复制输入。4.2 Noam 学习率调度与 warmupTransformer 原论文用带 warmup 的学习率先线性升到峰值再按步数平方根衰减。直接上固定学习率前期容易发散。class NoamScheduler: def __init__(self, optimizer, d_model, warmup4000): self.opt optimizer self.d_model d_model self.warmup warmup self.step_num 0 def step(self): self.step_num 1 lr self.d_model ** (-0.5) * min( self.step_num ** (-0.5), self.step_num * self.warmup ** (-1.5)) for g in self.opt.param_groups: g[lr] lr self.opt.step()warmup是预热步数小数据集 2000 到 4000 都行大数据集可以到 8000。注意这个调度器自己调opt.step()别在外面再调一次否则步数翻倍。4.3 梯度裁剪与训练循环骨架翻译模型对梯度爆炸敏感裁剪阈值一般 1.0 到 5.0。下面是一个能跑的最小训练循环。def train_epoch(model, loader, optimizer, scheduler, device, clip1.0): model.train() total 0 for src, tgt in loader: src, tgt src.to(device), tgt.to(device) src_pad make_pad_mask(src) tgt_mask make_tgt_mask(tgt) logits model(src, tgt[:, :-1], src_pad, tgt_mask[:, :, :-1, :-1]) loss compute_loss(logits, tgt) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), clip) scheduler.step() total loss.item() return total / len(loader)clip是梯度范数上限超过就等比缩放。注意tgt_mask要跟着tgt[:, :-1]一起裁否则形状对不上。训练时每轮存一次 checkpoint翻译任务训练周期长后悔药就是早存盘。5. 推理与评估贪心、束搜索和 BLEU 怎么算训练完直接拿模型生成会发现和训练时行为不一样没有 teacher forcing得自己一步步喂。推理策略和评估指标是这一章的重点。5.1 自回归解码与贪心搜索推理时从bos开始每次取概率最大的词接到序列后面直到eos或达到最大长度。torch.no_grad() def greedy_decode(model, src, sp_tgt, device, max_len100): model.eval() src src.to(device) src_pad make_pad_mask(src) ys torch.tensor([[BOS]], devicedevice) for _ in range(max_len): tgt_mask make_tgt_mask(ys) logits model(src, ys, src_pad, tgt_mask) next_id logits[:, -1].argmax(-1, keepdimTrue) ys torch.cat([ys, next_id], dim1) if next_id.item() EOS: break return sp_tgt.decode(ys[0, 1:].tolist())max_len要设得比训练最长句略大防止死循环。贪心快但容易局部最优短句够用长句质量一般。5.2 束搜索的实现要点束搜索同时保留 k 条候选按累计对数概率排序。k 取 4 到 10 之间太大速度掉得厉害。torch.no_grad() def beam_search(model, src, sp_tgt, device, beam5, max_len100): model.eval() src src.to(device) src_pad make_pad_mask(src) beams [(torch.tensor([[BOS]], devicedevice), 0.0)] for _ in range(max_len): cands [] for ys, score in beams: if ys[0, -1].item() EOS: cands.append((ys, score)) continue tgt_mask make_tgt_mask(ys) logits model(src, ys, src_pad, tgt_mask)[:, -1] logp torch.log_softmax(logits, dim-1) topk_p, topk_i logp.topk(beam) for p, i in zip(topk_p[0], topk_i[0]): cands.append((torch.cat([ys, i.view(1, 1)], dim1), score p.item())) beams sorted(cands, keylambda x: x[1], reverseTrue)[:beam] if all(b[0][0, -1].item() EOS for b in beams): break return sp_tgt.decode(beams[0][0][0, 1:].tolist())beam是束宽翻译任务常用 4 到 5。注意要对概率取 log 再累加直接乘概率会下溢。束搜索里长度归一化是个可选项长句不加归一化会吃亏。5.3 BLEU 评估与结果解读BLEU 是翻译的标准指标用 sacrebleu 或 nltk 都行。它看 n-gram 重合度对分词敏感评估时要用和训练一致的分词。from sacrebleu.metrics import BLEU def evaluate(model, loader, sp_tgt, device): refs, hyps [], [] for src, tgt in loader: for i in range(src.size(0)): hyps.append(beam_search(model, src[i:i1], sp_tgt, device)) refs.append(sp_tgt.decode(tgt[i].tolist())) bleu BLEU() return bleu.corpus_score(hyps, [refs]).scoreBLEU 分数受分词、大小写影响很大跨实验比较时保持评估脚本一致。单句 BLEU 波动大看语料级分数。分数低先别急着改模型回头查数据对齐和 mask。6. 避坑与排查翻译模型训练里最常见的 5 个翻车现场这一章是我踩过的坑里最费时间的几个按现象、原因、解决写遇到问题先对着查。6.1 loss 不降或直接变 nan现象训练几步后 loss 变 nan或者一直卡在 log(词表大小) 附近不动。原因学习率没预热、梯度爆炸或者 mask 语义搞反导致 softmax 全被屏蔽。解决先确认 warmup 生效加梯度裁剪到 1.0再打印一个 batch 的 mask确认被屏蔽位置是 True 且和 PyTorch 版本语义一致。nan 出现前把学习率降一个量级试。6.2 生成的句子全是重复词现象解码输出像「的的的的」或者一直重复同一短语。原因解码器看到了未来信息causal mask 没加对或者训练时目标错位写错导致模型学成复制。解决检查make_causal_mask是否真的屏蔽了上三角确认tgt_mask形状和输入对齐再核对compute_loss里的错位切片是不是logits[:, :-1]对tgt[:, 1:]。6.3 训练 loss 正常但 BLEU 极低现象loss 降到很低评估 BLEU 却只有个位数。原因评估时的分词和训练不一致或者推理时忘了把bos去掉、eos没截断。解决统一用同一个 sentencepiece 模型解码decode时跳过特殊符号确认推理输入是[BOS]开头而不是空序列。6.4 显存爆掉或 batch 上不去现象加大 batch 就 OOM或者长句一进来就崩。原因固定 padding 到全局最长句或者max_len设太大。解决用动态 padding 的 collate按 batch 内最长补齐把max_len从 100 降到 80 试开启混合精度torch.cuda.amp显存能省三成左右。6.5 推理速度慢到没法用现象束搜索一条句子要好几秒。原因束宽太大或者每步都重算整个编码器。解决编码器输出只算一次并缓存解码循环里复用束宽从 10 降到 4批量解码而不是一条条来。生产环境还可以考虑导出 ONNX 或 TorchScript 加速。7. 把模型推到能用混合精度、checkpoint 平均与一个验证习惯训练到能出像样 BLEU 只是第一步真正上线前还有几个提效技巧。混合精度是最划算的一个torch.cuda.amp几乎不改代码就能省显存、提速度。scaler torch.cuda.amp.GradScaler() def train_step(model, src, tgt, optimizer, scheduler, device): with torch.cuda.amp.autocast(): src_pad make_pad_mask(src) tgt_mask make_tgt_mask(tgt) logits model(src, tgt[:, :-1], src_pad, tgt_mask[:, :, :-1, :-1]) loss compute_loss(logits, tgt) optimizer.zero_grad() scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() scheduler.step() return loss.item()注意scaler.unscale_要在裁剪前调用否则裁的是缩放后的梯度阈值就失去意义。这是混合精度里最容易忽略的一步。另一个实用技巧是 checkpoint 平均把最后几个 epoch 的权重取平均往往比单点模型稳BLEU 能涨零点几个点。做法简单加载多个 state_dict 逐参数求平均即可。def average_checkpoints(paths, model): states [torch.load(p, map_locationcpu) for p in paths] avg {} for k in states[0]: avg[k] sum(s[k].float() for s in states) / len(states) model.load_state_dict(avg) return model参数上参与平均的 checkpoint 一般取验证集最好的 3 到 5 个太多会把早期欠拟合的权重也混进来。最后说一个我养成的验证习惯每次改完数据管道或 mask先拿 20 条句子做一次过拟合测试——用极小学习率把模型在这 20 条上训到 loss 接近 0再解码看能不能复现原句。如果连这 20 条都记不住问题一定在数据或 mask不在模型容量。这个习惯帮我省了无数次盲目调参的时间。翻译任务链路长任何一环出错都表现为 loss 不降与其猜不如用最小可复现的过拟合测试把问题锁死。希望帮到你。本文还有配套的精品资源点击获取
简介这份资源面向希望深入理解神经机器翻译与Transformer架构的深度学习学习者与研究者提供一套基于PyTorch的完整实现代码。内容围绕编码器-解码器结构展开涵盖多头注意力、位置编码、前馈神经网络等核心模块并包含数据预处理、词表构建、模型训练与评估的完整流程适合具备一定PyTorch基础、想动手复现NMT模型的读者。压缩包共59个文件以53个Python脚本为主辅以C源码、编译产物及fast_align对齐工具整体约468KB目录中可见模型定义、优化器、指标计算与配置等模块结构清晰便于按功能查阅。目前已有1788人学习下载可作为理解Transformer内部机制、调试训练流程与二次开发的实践参考。1. 从零搭一套 PyTorch Transformer 神经机器翻译为什么值得亲手写一遍很多人第一次接触神经机器翻译是从nn.Transformer那一行调用开始的跑通了却说不清数据怎么对齐、mask 怎么加、推理为什么和训练不一样。真到要改结构、换分词、上自己的语料时立刻卡住。这个标题要解决的就是这件事用 PyTorch 从数据管道到训练、推理完整实现一套基于 Transformer 的神经机器翻译系统而不是调个库看 loss 下降。它适合已经会写基础 PyTorch 训练循环、想真正吃透 Transformer 在 seq2seq 任务里怎么落地的人。下面按我实际搭过几版的顺序讲参数、坑、验证方法都给到能直接抄的程度。2. 数据管道与词表Transformer 翻译的地基怎么打Transformer 不像 RNN 那样天然按时间步吃序列它对数据的整齐度要求更高同一批里所有句子必须 pad 到等长还要额外维护 padding mask 和 causal mask。数据这一层没搭好后面训练再调参都是玄学。2.1 平行语料的读取与清洗神经机器翻译的标准输入是平行语料一行源语言对一行目标语言。常见做法是先做长度过滤和去重把源/目标长度比过大的句子丢掉这类句子多半是语料对齐错误留着只会让模型学噪声。def load_parallel(src_path, tgt_path, max_len100, max_ratio2.0): pairs [] with open(src_path, encodingutf-8) as fs, open(tgt_path, encodingutf-8) as ft: for s, t in zip(fs, ft): s, t s.strip(), t.strip() if not s or not t: continue # 长度过滤过长句子显存吃不消比例失衡多为对齐错误 if len(s.split()) max_len or len(t.split()) max_len: continue ratio len(s.split()) / max(len(t.split()), 1) if ratio max_ratio or ratio 1 / max_ratio: continue pairs.append((s, t)) # 去重避免重复句主导梯度 return list(dict.fromkeys(pairs))max_len控制单句词数上限直接决定后面 batch 的显存占用max_ratio是源目标长度比阈值1.5 到 2.5 之间都算合理语料质量差就收紧到 1.5。清洗完记得打印一下保留率掉到 70% 以下说明阈值太狠或者语料本身有问题。2.2 用 BPE 还是词级词表选型与实现词级词表实现简单但 OOV 多、词表爆炸BPE 子词能把罕见词拆成子词单元是现在翻译任务的主流。我一般用 sentencepiece 或 tokenizers 训一个共享或独立的 BPE词表大小 8k 到 32k 之间。词表太小句子被切太碎序列变长太大又退化成词级。import sentencepiece as spm # 训练源语言 BPEvocab_size 是核心参数 spm.SentencePieceTrainer.train( inputtrain.src, model_prefixspm_src, vocab_size16000, model_typebpe, character_coverage1.0, pad_id0, unk_id1, bos_id2, eos_id3, ) sp spm.SentencePieceProcessor(model_filespm_src.model) ids sp.encode(this is a test, out_typeint) # 含 bos/eosvocab_size是这里最需要调的参数中英这类差异大的语言对建议源目标各训一个character_coverage对中文要设 1.0否则生僻字全变 unk。注意 pad/unk/bos/eos 的 id 要固定下来后面 mask 全靠这几个特殊符号。2.3 Dataset 与动态 padding 的 collate 函数固定长度 padding 浪费算力动态 padding 按 batch 内最长句补齐更划算。collate 函数负责把一批样本 pad 齐并生成后续要用的 mask。import torch from torch.nn.utils.rnn import pad_sequence PAD, BOS, EOS 0, 2, 3 class TranslationDataset(torch.utils.data.Dataset): def __init__(self, pairs, sp_src, sp_tgt): self.pairs pairs self.sp_src, self.sp_tgt sp_src, sp_tgt def __len__(self): return len(self.pairs) def __getitem__(self, idx): s, t self.pairs[idx] src self.sp_src.encode(s, out_typeint) # 已含 bos/eos tgt self.sp_tgt.encode(t, out_typeint) return torch.tensor(src), torch.tensor(tgt) def collate(batch): srcs, tgts zip(*batch) src pad_sequence(srcs, batch_firstTrue, padding_valuePAD) tgt pad_sequence(tgts, batch_firstTrue, padding_valuePAD) return src, tgtbatch_firstTrue让张量形状是[B, L]和nn.Transformer默认的[L, B, D]不同后面进模型前要转置这点很容易翻车。pad 值必须和词表里的 pad_id 一致否则 mask 会算错。3. 模型结构手写编码器解码器还是直接用 nn.Transformer选型上分两派直接用nn.Transformer省事手写 Encoder/Decoder 层能改结构、看得清每个 mask。想真正理解翻译任务里 mask 怎么配合我建议至少手写一遍再切回官方实现。3.1 位置编码与词嵌入的缩放关系Transformer 没有循环结构位置信息全靠位置编码注入。原始论文用正弦编码也可以换成可学习的位置嵌入。关键细节是词嵌入要乘sqrt(d_model)否则和位置编码量级不匹配训练初期梯度会不稳。import math import torch.nn as nn class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len5000, dropout0.1): super().__init__() pe torch.zeros(max_len, d_model) pos torch.arange(0, max_len).unsqueeze(1).float() div torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(pos * div) pe[:, 1::2] torch.cos(pos * div) self.register_buffer(pe, pe.unsqueeze(0)) # [1, max_len, d_model] self.dropout nn.Dropout(dropout) def forward(self, x): # x: [B, L, D]词嵌入已乘 sqrt(d_model) return self.dropout(x self.pe[:, :x.size(1)])d_model要和词嵌入维度一致常见 256/512max_len要大于你语料里最长句超了会索引越界。register_buffer保证 pe 跟着模型上 GPU 但不参与梯度。3.2 三种 mask 的生成padding、causal、cross翻译任务里 mask 是最容易出错的地方。编码器只需要 padding mask解码器自注意力要 padding causal交叉注意力用编码器的 padding mask。搞混一个模型要么看到未来词要么把 pad 当有效词。def make_pad_mask(seq, pad_id0): # seq: [B, L] - [B, 1, 1, L]True 表示被 mask return (seq pad_id).unsqueeze(1).unsqueeze(2) def make_causal_mask(size, device): # 上三角为 True屏蔽未来位置 return torch.triu(torch.ones(size, size, devicedevice), diagonal1).bool() def make_tgt_mask(tgt, pad_id0): pad_mask make_pad_mask(tgt, pad_id) # [B,1,1,L] causal make_causal_mask(tgt.size(1), tgt.device) # [L,L] return pad_mask | causal # 广播成 [B,1,L,L]nn.Transformer的src_mask形状是[L, L]或[B, L, L]tgt_mask同理src_key_padding_mask是[B, L]。不同版本 PyTorch 对 mask 语义True 是被屏蔽还是保留有过调整务必对着你装的版本确认这是血泪经验。3.3 组装完整模型与输出层输出层是一个线性层把d_model映射到目标词表大小配合交叉熵。注意解码器输出要经过这个线性层且和词嵌入是否共享权重可以作为一个调优点。class Seq2SeqTransformer(nn.Module): def __init__(self, src_vocab, tgt_vocab, d_model512, nhead8, num_enc6, num_dec6, dim_ff2048, dropout0.1): super().__init__() self.src_emb nn.Embedding(src_vocab, d_model, padding_idx0) self.tgt_emb nn.Embedding(tgt_vocab, d_model, padding_idx0) self.pos PositionalEncoding(d_model, dropoutdropout) self.transformer nn.Transformer( d_modeld_model, nheadnhead, num_encoder_layersnum_enc, num_decoder_layersnum_dec, dim_feedforwarddim_ff, dropoutdropout, batch_firstTrue) self.fc nn.Linear(d_model, tgt_vocab) self.d_model d_model def forward(self, src, tgt, src_pad_mask, tgt_mask): src self.pos(self.src_emb(src) * math.sqrt(self.d_model)) tgt self.pos(self.tgt_emb(tgt) * math.sqrt(self.d_model)) out self.transformer(src, tgt, tgt_masktgt_mask, src_key_padding_masksrc_pad_mask.squeeze(1).squeeze(1), tgt_key_padding_mask(tgt 0)) return self.fc(out)nhead必须能整除d_model512/8 是经典组合dim_ff一般是d_model的 4 倍。batch_firstTrue时输入是[B, L, D]和前面 collate 对齐省掉转置。padding_idx0让 pad 的词嵌入不更新。4. 训练循环标签平滑、学习率预热与梯度裁剪翻译任务的训练有几个和普通分类不一样的地方目标序列要错位一位做 teacher forcing损失常用标签平滑学习率要预热。这几项配错loss 会卡住或者震荡。4.1 目标错位与损失计算解码器输入是bos ...预测目标是... eos也就是输入去掉最后一位、目标去掉第一位。交叉熵要忽略 pad 位置。import torch.nn.functional as F def compute_loss(logits, tgt, pad_id0, smoothing0.1): # logits: [B, L, V]tgt: [B, L] logits logits[:, :-1].reshape(-1, logits.size(-1)) gold tgt[:, 1:].reshape(-1) return F.cross_entropy(logits, gold, ignore_indexpad_id, label_smoothingsmoothing)label_smoothing取 0.1 是翻译任务的常见值能缓解过拟合、提升 BLEUignore_indexpad_id保证 pad 不贡献梯度。错位这一步写错模型会学成复制输入。4.2 Noam 学习率调度与 warmupTransformer 原论文用带 warmup 的学习率先线性升到峰值再按步数平方根衰减。直接上固定学习率前期容易发散。class NoamScheduler: def __init__(self, optimizer, d_model, warmup4000): self.opt optimizer self.d_model d_model self.warmup warmup self.step_num 0 def step(self): self.step_num 1 lr self.d_model ** (-0.5) * min( self.step_num ** (-0.5), self.step_num * self.warmup ** (-1.5)) for g in self.opt.param_groups: g[lr] lr self.opt.step()warmup是预热步数小数据集 2000 到 4000 都行大数据集可以到 8000。注意这个调度器自己调opt.step()别在外面再调一次否则步数翻倍。4.3 梯度裁剪与训练循环骨架翻译模型对梯度爆炸敏感裁剪阈值一般 1.0 到 5.0。下面是一个能跑的最小训练循环。def train_epoch(model, loader, optimizer, scheduler, device, clip1.0): model.train() total 0 for src, tgt in loader: src, tgt src.to(device), tgt.to(device) src_pad make_pad_mask(src) tgt_mask make_tgt_mask(tgt) logits model(src, tgt[:, :-1], src_pad, tgt_mask[:, :, :-1, :-1]) loss compute_loss(logits, tgt) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), clip) scheduler.step() total loss.item() return total / len(loader)clip是梯度范数上限超过就等比缩放。注意tgt_mask要跟着tgt[:, :-1]一起裁否则形状对不上。训练时每轮存一次 checkpoint翻译任务训练周期长后悔药就是早存盘。5. 推理与评估贪心、束搜索和 BLEU 怎么算训练完直接拿模型生成会发现和训练时行为不一样没有 teacher forcing得自己一步步喂。推理策略和评估指标是这一章的重点。5.1 自回归解码与贪心搜索推理时从bos开始每次取概率最大的词接到序列后面直到eos或达到最大长度。torch.no_grad() def greedy_decode(model, src, sp_tgt, device, max_len100): model.eval() src src.to(device) src_pad make_pad_mask(src) ys torch.tensor([[BOS]], devicedevice) for _ in range(max_len): tgt_mask make_tgt_mask(ys) logits model(src, ys, src_pad, tgt_mask) next_id logits[:, -1].argmax(-1, keepdimTrue) ys torch.cat([ys, next_id], dim1) if next_id.item() EOS: break return sp_tgt.decode(ys[0, 1:].tolist())max_len要设得比训练最长句略大防止死循环。贪心快但容易局部最优短句够用长句质量一般。5.2 束搜索的实现要点束搜索同时保留 k 条候选按累计对数概率排序。k 取 4 到 10 之间太大速度掉得厉害。torch.no_grad() def beam_search(model, src, sp_tgt, device, beam5, max_len100): model.eval() src src.to(device) src_pad make_pad_mask(src) beams [(torch.tensor([[BOS]], devicedevice), 0.0)] for _ in range(max_len): cands [] for ys, score in beams: if ys[0, -1].item() EOS: cands.append((ys, score)) continue tgt_mask make_tgt_mask(ys) logits model(src, ys, src_pad, tgt_mask)[:, -1] logp torch.log_softmax(logits, dim-1) topk_p, topk_i logp.topk(beam) for p, i in zip(topk_p[0], topk_i[0]): cands.append((torch.cat([ys, i.view(1, 1)], dim1), score p.item())) beams sorted(cands, keylambda x: x[1], reverseTrue)[:beam] if all(b[0][0, -1].item() EOS for b in beams): break return sp_tgt.decode(beams[0][0][0, 1:].tolist())beam是束宽翻译任务常用 4 到 5。注意要对概率取 log 再累加直接乘概率会下溢。束搜索里长度归一化是个可选项长句不加归一化会吃亏。5.3 BLEU 评估与结果解读BLEU 是翻译的标准指标用 sacrebleu 或 nltk 都行。它看 n-gram 重合度对分词敏感评估时要用和训练一致的分词。from sacrebleu.metrics import BLEU def evaluate(model, loader, sp_tgt, device): refs, hyps [], [] for src, tgt in loader: for i in range(src.size(0)): hyps.append(beam_search(model, src[i:i1], sp_tgt, device)) refs.append(sp_tgt.decode(tgt[i].tolist())) bleu BLEU() return bleu.corpus_score(hyps, [refs]).scoreBLEU 分数受分词、大小写影响很大跨实验比较时保持评估脚本一致。单句 BLEU 波动大看语料级分数。分数低先别急着改模型回头查数据对齐和 mask。6. 避坑与排查翻译模型训练里最常见的 5 个翻车现场这一章是我踩过的坑里最费时间的几个按现象、原因、解决写遇到问题先对着查。6.1 loss 不降或直接变 nan现象训练几步后 loss 变 nan或者一直卡在 log(词表大小) 附近不动。原因学习率没预热、梯度爆炸或者 mask 语义搞反导致 softmax 全被屏蔽。解决先确认 warmup 生效加梯度裁剪到 1.0再打印一个 batch 的 mask确认被屏蔽位置是 True 且和 PyTorch 版本语义一致。nan 出现前把学习率降一个量级试。6.2 生成的句子全是重复词现象解码输出像「的的的的」或者一直重复同一短语。原因解码器看到了未来信息causal mask 没加对或者训练时目标错位写错导致模型学成复制。解决检查make_causal_mask是否真的屏蔽了上三角确认tgt_mask形状和输入对齐再核对compute_loss里的错位切片是不是logits[:, :-1]对tgt[:, 1:]。6.3 训练 loss 正常但 BLEU 极低现象loss 降到很低评估 BLEU 却只有个位数。原因评估时的分词和训练不一致或者推理时忘了把bos去掉、eos没截断。解决统一用同一个 sentencepiece 模型解码decode时跳过特殊符号确认推理输入是[BOS]开头而不是空序列。6.4 显存爆掉或 batch 上不去现象加大 batch 就 OOM或者长句一进来就崩。原因固定 padding 到全局最长句或者max_len设太大。解决用动态 padding 的 collate按 batch 内最长补齐把max_len从 100 降到 80 试开启混合精度torch.cuda.amp显存能省三成左右。6.5 推理速度慢到没法用现象束搜索一条句子要好几秒。原因束宽太大或者每步都重算整个编码器。解决编码器输出只算一次并缓存解码循环里复用束宽从 10 降到 4批量解码而不是一条条来。生产环境还可以考虑导出 ONNX 或 TorchScript 加速。7. 把模型推到能用混合精度、checkpoint 平均与一个验证习惯训练到能出像样 BLEU 只是第一步真正上线前还有几个提效技巧。混合精度是最划算的一个torch.cuda.amp几乎不改代码就能省显存、提速度。scaler torch.cuda.amp.GradScaler() def train_step(model, src, tgt, optimizer, scheduler, device): with torch.cuda.amp.autocast(): src_pad make_pad_mask(src) tgt_mask make_tgt_mask(tgt) logits model(src, tgt[:, :-1], src_pad, tgt_mask[:, :, :-1, :-1]) loss compute_loss(logits, tgt) optimizer.zero_grad() scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() scheduler.step() return loss.item()注意scaler.unscale_要在裁剪前调用否则裁的是缩放后的梯度阈值就失去意义。这是混合精度里最容易忽略的一步。另一个实用技巧是 checkpoint 平均把最后几个 epoch 的权重取平均往往比单点模型稳BLEU 能涨零点几个点。做法简单加载多个 state_dict 逐参数求平均即可。def average_checkpoints(paths, model): states [torch.load(p, map_locationcpu) for p in paths] avg {} for k in states[0]: avg[k] sum(s[k].float() for s in states) / len(states) model.load_state_dict(avg) return model参数上参与平均的 checkpoint 一般取验证集最好的 3 到 5 个太多会把早期欠拟合的权重也混进来。最后说一个我养成的验证习惯每次改完数据管道或 mask先拿 20 条句子做一次过拟合测试——用极小学习率把模型在这 20 条上训到 loss 接近 0再解码看能不能复现原句。如果连这 20 条都记不住问题一定在数据或 mask不在模型容量。这个习惯帮我省了无数次盲目调参的时间。翻译任务链路长任何一环出错都表现为 loss 不降与其猜不如用最小可复现的过拟合测试把问题锁死。希望帮到你。本文还有配套的精品资源点击获取
汉明码在BSC与AWGN信道下的误码率仿真与性能分析)
