中英翻译之05 Transformer 训练脚本

MAX。

266人浏览 · 2026-05-25 21:44:00

MAX。 · 2026-05-25 21:44:00 发布

中英翻译之05 Transformer 训练脚本

一、这个脚本到底是干嘛的？

一句话说清楚：这就是 AI 翻译官的老师！

前面我们准备好了：

课本：预处理好的中英文数字数据
学生：空的 Transformer 模型

现在这个脚本就是老师，负责：

把课本一页一页地拿给学生看
让学生试着翻译
批改学生的作业，指出哪里错了
让学生改正错误，不断进步
记录学生的成绩，最后选出学得最好的那个学生保存下来

这是整个机器翻译项目的最后一步，跑完这个脚本，你就得到了一个能真正翻译句子的 AI 模型。

二、代码逐行大白话解释

2.1 开头导入工具

# 导入时间模块：用来生成时间戳，给每次训练的日志文件夹起名字
import time

# 导入PyTorch核心库
import torch
# 导入神经网络模块和优化器模块
from torch import nn, optim

# 导入配置文件
from config import *
# 导入数据加载器
from dataset import get_dataloader
# 导入翻译模型
from model import TranslationModel
# 导入分词器
from tokenizer import ZhTokenizer, EnTokenizer

# 导入tqdm：用来在命令行显示漂亮的进度条，告诉你训练到哪了
from tqdm import tqdm
# 导入TensorBoard：用来记录训练过程，画损失曲线，可视化训练
from torch.utils.tensorboard import SummaryWriter

2.2 单轮训练函数 train_one_epoch

这是整个脚本最核心的部分！ 它定义了老师怎么教学生学完一整本课本（遍历完所有训练数据一次）。

# 训练一轮：教学生学完所有训练数据一次
def train_one_epoch(model, train_loader, loss_fn, optimizer, device):
    # 把模型设置为"上课模式"
    # 训练模式下，Dropout（随机让一些神经元睡觉，防止死记硬背）会生效
    model.train()
    # 初始化总损失：用来记录这一轮学生总共错了多少
    total_loss = 0

    # 遍历每一批数据，tqdm会显示进度条
    for input, target in tqdm(train_loader, desc='Train'):
        # 把数据搬到GPU上（如果有GPU的话）
        # 模型在哪个设备上，数据就必须在哪个设备上，不然会报错
        input = input.to(device)
        target = target.to(device)

最关键的一步：拆分解码器的输入和目标

90% 的新手都会在这里卡壳！ 这是 “教师强制训练” 的核心。

        # ============ 最重要的一步：拆分解码器输入和目标 ============
        # 我们的目标序列是这样的：[<sos>, 我, 爱, 你, <eos>]
        # 解码器输入：去掉最后一个token → [<sos>, 我, 爱, 你]
        decoder_input = target[:, :-1]
        # 解码器目标：去掉第一个token → [我, 爱, 你, <eos>]
        decoder_target = target[:, 1:]

超级形象的比喻：

老师教学生造句：“我爱你。”

老师说：“我” → 学生应该说：“爱”
老师说：“我爱” → 学生应该说：“你”
老师说：“我爱你” → 学生应该说：“。”

所以：

老师说的话（解码器输入）：[我, 我爱, 我爱你] → 对应去掉最后一个 token
学生应该说的话（解码器目标）：[爱, 你, 。] → 对应去掉第一个 token

这样模型就能学会：根据前面已经生成的词，预测下一个词。

生成三个掩码

        # 生成中文填充掩码：告诉模型哪些位置是0（填充的），不用管
        src_padding_mask = (input == model.src_padding_idx)
        # 生成英文填充掩码
        tgt_padding_mask = (decoder_input == model.tgt_padding_idx)

        # 生成因果掩码：告诉模型不能看未来的词
        # 比如预测"爱"的时候，不能看后面的"你"和"。"
        tgt_mask = model.transformer.generate_square_subsequent_mask(
            decoder_input.shape[1]
        ).bool().to(device)

前向传播：让学生试着翻译

        # 把数据传给模型，让模型输出预测结果
        # decoder_output形状：(32, 17, 1500) → 32个句子，每个句子17个词，每个词1500个可能
        decoder_output = model(input, decoder_input, src_padding_mask, tgt_mask, tgt_padding_mask)

计算损失：批改学生的作业

        # 计算损失：学生的预测和正确答案差多少
        # 这里有个坑！必须转置最后两个维度！
        loss = loss_fn(decoder_output.mT, decoder_target)

为什么要转置？

PyTorch 的交叉熵损失函数CrossEntropyLoss要求输入形状是(批次大小, 类别数, 序列长度)
但我们的模型输出形状是(批次大小, 序列长度, 类别数)
所以要用.mT把最后两个维度转置一下，变成(批次大小, 类别数, 序列长度)
忘了转置会报形状不匹配的错误，非常难排查！

反向传播 + 更新参数：让学生改正错误

        # 反向传播：计算每个参数应该改多少
        loss.backward()

        # 更新参数：根据计算出来的梯度，调整模型的权重
        optimizer.step()

        # 梯度清零：把这次的梯度擦掉，不然下次会和这次的加起来
        optimizer.zero_grad()

        # 累加这次的损失
        total_loss += loss.item()

    # 返回这一轮的平均损失
    return total_loss / len(train_loader)

这五步是深度学习训练的标准流程，永远不变：

前向传播 → 得到预测结果
计算损失 → 知道错了多少
反向传播 → 知道哪里错了
更新参数 → 改正错误
梯度清零 → 准备下一次

2.3 整体训练主函数 train

这是整个训练流程的总指挥，负责把所有东西串起来。

def train():
    # ============ 步骤1：选择用什么设备训练 ============
    # 有GPU就用GPU，没有就用CPU
    # GPU训练比CPU快10-100倍！
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # ============ 步骤2：获取训练数据加载器 ============
    train_loader = get_dataloader()

    # ============ 步骤3：加载分词器 ============
    # 从之前保存的词表文件中加载分词器
    zh_tokenizer = ZhTokenizer.from_vocab(MODEL_DIR / ZH_VOCAB_FILE)
    en_tokenizer = EnTokenizer.from_vocab(MODEL_DIR / EN_VOCAB_FILE)

    # ============ 步骤4：创建模型 ============
    # 传入中英文词表大小和填充ID，创建模型并搬到GPU上
    model = TranslationModel(
        zh_tokenizer.vocab_size,
        en_tokenizer.vocab_size,
        zh_tokenizer.pad_id,
        en_tokenizer.pad_id,
    ).to(device)

    # ============ 步骤5：定义损失函数和优化器 ============
    # 交叉熵损失函数：专门用来做多分类问题（预测下一个词就是多分类）
    # ignore_index=en_tokenizer.pad_id：忽略填充位置的损失
    # 因为填充的位置不是真正的词，不应该让模型为这些位置的错误负责
    loss_fn = nn.CrossEntropyLoss(ignore_index=en_tokenizer.pad_id)
    
    # Adam优化器：目前最好用的自适应学习率优化器
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)

    # 创建TensorBoard日志记录器：记录训练过程
    # 每次训练都会生成一个以时间戳命名的文件夹，方便对比不同次的训练
    writer = SummaryWriter(log_dir=LOG_DIR / time.strftime("%Y-%m-%d_%H-%M-%S"))

训练循环：教学生学很多轮

    # ============ 步骤6：开始训练循环 ============
    # 初始化最小损失为无穷大
    min_loss = float('inf')
    
    # 训练EPOCHS轮
    for epoch in range(EPOCHS):
        print(f'\nEpoch {epoch + 1}/{EPOCHS}')
        
        # 训练一轮，得到平均损失
        train_loss = train_one_epoch(model, train_loader, loss_fn, optimizer, device)
        print(f'Train Loss: {train_loss:.4f}')

        # 把损失值记录到TensorBoard，这样就能看到损失曲线了
        writer.add_scalar('Train Loss', train_loss, epoch + 1)

        # 保存最优模型：如果这一轮的损失比之前所有轮都小，就保存这个模型
        if train_loss < min_loss:
            min_loss = train_loss
            # 只保存模型的参数权重，不保存整个模型对象
            # 这样保存的文件更小，加载更快
            torch.save(model.state_dict(), MODEL_DIR / BEST_MODEL)
            print("✅ Best model saved!")

    # 训练结束，关闭日志记录器
    writer.close()
    print("\nTraining finished!")

为什么要保存最优模型而不是最后一个模型？

模型可能会 “过拟合”：学了太多课本上的细节，反而不会做新题了
表现最好的模型通常出现在训练的中间，而不是最后
所以我们只保存损失最小的那个模型，这是泛化能力最好的模型

2.4 程序入口

if __name__ == '__main__':
    train()

直接运行这个脚本，就会开始训练模型了。

三、特别重要的说明

整个训练流程一句话总结：

加载数据 → 创建模型 → 循环很多轮：
    遍历每一批数据 → 拆分解码器输入和目标 → 前向传播 → 计算损失 → 反向传播 → 更新参数
保存最优模型

教师强制训练 (Teacher Forcing)：
- 训练的时候，我们不把模型上一步预测的词作为下一步的输入
- 而是直接把正确的词作为输入
- 这样训练更快，更稳定，不会出现错误累积的问题
- 推理的时候才会用模型自己预测的词作为下一步的输入
为什么要忽略填充位置的损失？
- 我们为了让所有句子长度一样，在短句子后面补了很多 0
- 这些 0 不是真正的词，模型不需要预测它们
- 如果不忽略这些位置的损失，模型会把大部分精力都用在预测这些没用的 0 上
- 导致真正的词预测得很差
TensorBoard 怎么用？
- 训练的时候，打开一个新的命令行，输入：tensorboard --logdir=logs
- 然后在浏览器里打开http://localhost:6006
- 就能看到实时的损失曲线了
新手最容易踩的坑：
- 数据没有移到 GPU 上：input.to(device)和target.to(device)
- 掩码没有移到 GPU 上：tgt_mask.to(device)
- 损失计算时忘了转置：decoder_output.mT
- 梯度没有清零：optimizer.zero_grad()
- 忘了设置ignore_index：损失会非常大，模型学不会
- 模型没有移到 GPU 上：model.to(device)
- 训练的时候忘了设置model.train()
- 保存模型的时候保存了整个模型而不是state_dict