PyTorch Transformer从入门到实战:手把手教你搭建中英翻译系统(附完整代码)

0 阅读18分钟

注:接上一篇Transformer理论之后的实战

很多人学Transformer只看理论,一写代码就懵。今天我就用PyTorch官方API,带你从头撸一个中英翻译项目。代码全公开,注释比代码还多,保证你跑得起来、看得懂、能复用。


一、PyTorch的Transformer模块:别再自己造轮子了

PyTorch从1.2版本开始就提供了torch.nn.Transformer等官方实现,封装了论文《Attention Is All You Need》中的完整架构。你不需要手写多头注意力、残差连接、前馈网络这些基础组件,直接调用即可。

但是,官方模块不包含位置编码,也不包含词嵌入层和输出层,这些需要我们自己补上。另外,掩码的生成也有现成方法。

下面我们先过一遍核心类,然后直接进入项目实战。

1.1 五大核心类

类名

作用

nn.Transformer

完整编码器-解码器容器,最外层接口

nn.TransformerEncoder

编码器,由多个

TransformerEncoderLayer

堆叠

nn.TransformerDecoder

解码器,由多个

TransformerDecoderLayer

堆叠

nn.TransformerEncoderLayer

单个编码器层(自注意力 + FFN + 残差+LayerNorm)

nn.TransformerDecoderLayer

单个解码器层(掩码自注意力 + 编码器-解码器注意力 + FFN + 残差+LayerNorm)

使用的时候,最省事的是直接实例化nn.Transformer,它会自动创建编码器和解码器。你也可以单独使用nn.TransformerEncoder,比如做文本分类(只需要编码器)。

1.2 构造参数详解

torch.nn.Transformer(
    d_model=512,           # 词向量维度,所有输入输出统一用这个维度
    nhead=8,               # 多头注意力的头数,必须能整除d_model
    num_encoder_layers=6,  # 编码器堆叠层数
    num_decoder_layers=6,  # 解码器堆叠层数
    dim_feedforward=2048,  # FFN隐藏层维度,一般是d_model的4倍
    dropout=0.1,           # 所有dropout层共用这个概率
    activation='relu',     # FFN的激活函数,可选'relu'或'gelu'
    batch_first=False,     # 输入形状是否为(batch, seq, feature)?推荐设为True
    norm_first=False,      # True表示先LayerNorm再子层(Pre-LN),False表示先子层再Norm(Post-LN)
    ...
)

注意:batch_first建议设为True,这样张量形状就是(batch_size, seq_len, d_model),符合大多数人的习惯。默认为False时形状是(seq_len, batch_size, d_model),容易搞混。

1.3 forward方法的输入输出

nn.Transformer.forward()接收以下关键参数:

参数

形状

含义

src

(batch, src_len, d_model)

源序列嵌入(词嵌入+位置编码)

tgt

(batch, tgt_len, d_model)

目标序列嵌入(解码器输入,通常是右移一位的)

src_key_padding_mask

(batch, src_len)

True

的位置表示该token是

<pad>

,会被忽略

tgt_key_padding_mask

(batch, tgt_len)

同上,用于解码器自注意力

tgt_mask

(tgt_len, tgt_len)

上三角掩码,防止解码器看到未来token

memory_key_padding_mask

(batch, src_len)

用于编码器-解码器注意力,屏蔽源端的

<pad>

输出形状:(batch, tgt_len, d_model),即解码器每个位置的输出表示。

重要:tgt_mask一般用nn.Transformer.generate_square_subsequent_mask()生成,它会返回一个上三角为-inf、下三角为0的矩阵。这样在softmax之后,未来位置的权重就变成0了。


二、中英翻译实战:从数据到部署

理论说完了,我们直接开干。目标是训练一个中文→英文的翻译模型。数据集使用经典的cmn.txt(中英文句子对),你可以从网上下载,格式类似:

Hi.	你好.
Hello.	你好.
How are you?	你好吗?
I'm fine.	我很好.

也可以直接下载我下载好的:pan.baidu.com/s/1As2fpzjO…

2.1 项目结构(推荐)

translation-transformer/
├── data/
│   ├── raw/           # 存放原始cmn.txt
│   └── processed/     # 预处理后保存的索引文件和词表
├── models/            # 保存训练好的模型参数
├── logs/              # TensorBoard日志
├── src/
│   ├── config.py      # 所有超参数
│   ├── tokenizer.py   # 中英文分词器(含词表构建)
│   ├── process.py     # 数据预处理脚本
│   ├── dataset.py     # PyTorch Dataset和DataLoader
│   ├── model.py       # 位置编码 + 完整翻译模型
│   ├── train.py       # 训练循环
│   ├── predict.py     # 推理脚本(交互式)
│   └── evaluate.py    # 评估BLEU分数

2.2 配置参数(config.py)

把所有可调整的参数集中管理,方便调参。

from pathlib import Path    # 路径定义
 
# 1. 目录路径
# 项目根目录
ROOT_DIR = Path(__file__).parent.parent
# 数据目录
RAW_DATA_DIR = ROOT_DIR / 'data' / 'raw'
PROCESSED_DATA_DIR = ROOT_DIR / 'data' / 'processed'
# 模型目录
MODEL_DIR = ROOT_DIR / 'models'
# 日志目录
LOG_DIR = ROOT_DIR / 'logs'
 
# 2. 文件
RAW_DATA_FILE = 'cmn.txt'
TRAIN_DATA_FILE = 'train.jsonl'
TEST_DATA_FILE = 'test.jsonl'
# VOCAB_FILE = 'vocab.txt'    # 词表文件
EN_VOCAB_FILE = 'en_vocab.txt'    # 英文词表
CN_VOCAB_FILE = 'cn_vocab.txt'    # 中文文件
BEST_MODEL = 'best_model.pt'    # 最优模型参数文件
 
# 3. 特殊token
UNK_TOKEN = '<unk>'     # 未登录词
PAD_TOKEN = '<pad>'     # 填充词
START_TOKEN = '<sos>'   # 起始标记
END_TOKEN = '<eos>'     # 结束标记
 
# 4. 训练超参数
SEQ_LEN = 128     # 序列(最大)长度
BATCH_SIZE = 64
 
LEARNING_RATE = 1e-3
EPOCHS = 70
 
# 5. 模型结构参数
DIM_MODEL = 128
NUM_HEADS = 4
NUM_ENCODER_LAYERS = 2
NUM_DECODER_LAYERS = 2

说明:DIM_MODEL=128是为了在普通CPU/GPU上快速跑通,实际应用可以调大。NUM_ENCODER_LAYERS=2也是减少计算量,原论文用6层。

2.3 分词器(tokenizer.py)

我们需要分别处理中文和英文。中文按字切分(最简单),英文用NLTK的word_tokenize。同时要构建词表,并实现encode(句子→索引列表)和decode(索引→句子)方法。

import jieba
from config import *
 
from nltk import TreebankWordTokenizer, TreebankWordDetokenizer     # 英文分词器
 
class BaseTokenizer():
    unk_token = UNK_TOKEN   # 类属性
    pad_token = PAD_TOKEN
    start_token = START_TOKEN
    end_token = END_TOKEN
 
    # 初始化
    def __init__(self, vocab_list):
        self.vocab_list = vocab_list
        self.vocab_size = len(vocab_list)   # 词表大小
        self.word2id = { word : id for id, word in enumerate(vocab_list) }
        self.id2word = { id : word for id, word in enumerate(vocab_list) }
        # self.unk_token = UNK_TOKEN
        self.unk_id = self.word2id[self.unk_token]
        self.pad_id = self.word2id[self.pad_token]
        self.start_id = self.word2id[self.start_token]
        self.end_id = self.word2id[self.end_token]
 
    # 分词,类方法接口
    @classmethod
    def tokenize(cls, text) -> list[str]:
        pass
 
    # 编码(将文本分词、id化),并指定序列长度
    def encode(self, text, mark=False):
        # 分词
        tokens = self.tokenize(text)
 
        # 如果是目标序列,就在前后加入标记
        if mark:
            tokens = [self.start_token] + tokens + [self.end_token]
 
        # id化
        ids = [self.word2id.get(token, self.unk_id) for token in tokens]
        return ids
 
    # 构建词表,并保存到文件
    @classmethod
    def build_vocab(cls, sentences, vocab_file_path):
        # 1. 针对训练集分词,构建词表
        vocab_set = set()  # 利用集合做token去重
        for sentence in sentences:
            vocab_set.update( cls.tokenize(sentence) )
        # 转换成列表(词表,id2word),并处理未登录词和填充词
        vocab_list = [cls.pad_token, cls.unk_token, cls.start_token, cls.end_token] + list(vocab_set)
 
        print("词表大小:", len(vocab_list))
 
        # 2. 保存词表到文件
        with open( vocab_file_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(vocab_list))
 
    # 从文件加载词表,并创建分词器对象实例
    @classmethod
    def from_vocab(cls, vocab_file_path):
        # 1. 获取词表
        with open( vocab_file_path, 'r', encoding='utf-8') as f:
            # 读取每一行
            vocab_list = [token.strip() for token in f.readlines()]
        # 2. 构建分词器对象
        tokenizer = cls(vocab_list)
        return tokenizer
 
# 定义子类
# 中文分词器
class ChineseTokenizer(BaseTokenizer):
    @classmethod
    def tokenize(cls, text) -> list[str]:
        return list(text)
 
# 英文分词器
class EnglishTokenizer(BaseTokenizer):
    tokenizer = TreebankWordTokenizer()
    detokenizer = TreebankWordDetokenizer()
 
    @classmethod
    def tokenize(cls, text) -> list[str]:
        return cls.tokenizer.tokenize(text)
 
    # 解码:传入一个id列表,返回原始英文句子
    def decode(self, ids):
        # 将id转换为 token
        tokens = [ self.id2word[id] for id in ids ]
        return self.detokenizer.detokenize(tokens)
 
if __name__ == '__main__':
    en_tokenizer = EnglishTokenizer.from_vocab( MODEL_DIR / EN_VOCAB_FILE )
    cn_tokenizer = ChineseTokenizer.from_vocab( MODEL_DIR / CN_VOCAB_FILE )
 
    print("中文词表大小:", cn_tokenizer.vocab_size)
    print("英文词表大小:", en_tokenizer.vocab_size)
    print("特殊符号UNK:", cn_tokenizer.unk_token)
    print("特殊符号PAD ID:", en_tokenizer.pad_id)
    print("特殊符号START:", en_tokenizer.start_token)
    print("特殊符号END ID:", cn_tokenizer.end_id)
 
    print( cn_tokenizer.encode("自然语言处理") )
    print( en_tokenizer.encode("hello world!", mark=True) )

2.4 数据预处理(preprocess.py)

读取原始cmn.txt,划分训练集/测试集,构建词表,并将句子转成索引序列,保存为JSON Lines格式。

# 数据预处理
import pandas as pd
from sklearn.model_selection import train_test_split    # 划分数据集
 
from config import *
from tokenizer import ChineseTokenizer, EnglishTokenizer    # 中英文分词器
 
 
def preprocess():
    print("-------开始数据预处理...-------")
 
    # 1. 以csv格式读取txt文件,得到DataFrame;并提取两列,去除缺失值
    df = pd.read_csv(RAW_DATA_DIR / RAW_DATA_FILE, sep='\t', usecols=[0, 1], names=['en', 'cn'], encoding='utf-8').dropna()
 
    # 3. 对原始语料做划分
    train_df, test_df = train_test_split(df, test_size=0.2)
 
    # 4. 分词并构建词表、保存到文件
    ChineseTokenizer.build_vocab(train_df['cn'].tolist(), MODEL_DIR/CN_VOCAB_FILE)
    EnglishTokenizer.build_vocab(train_df['en'].tolist(), MODEL_DIR/EN_VOCAB_FILE)
 
    # 5. 创建分词器
    cn_tokenizer = ChineseTokenizer.from_vocab(MODEL_DIR/CN_VOCAB_FILE)
    en_tokenizer = EnglishTokenizer.from_vocab(MODEL_DIR/EN_VOCAB_FILE)
 
    # 6. 构建数据集
    train_df['cn'] = train_df['cn'].apply( lambda text: cn_tokenizer.encode(text, mark=False) )
    train_df['en'] = train_df['en'].apply(lambda text: en_tokenizer.encode(text, mark=True))
 
    test_df['cn'] = test_df['cn'].apply( lambda text: cn_tokenizer.encode(text, mark=False) )
    test_df['en'] = test_df['en'].apply(lambda text: en_tokenizer.encode(text, mark=True))
 
 
    # 7. 保存数据集到文件
    train_df.to_json(PROCESSED_DATA_DIR/TRAIN_DATA_FILE, orient='records', lines=True)
    test_df.to_json(PROCESSED_DATA_DIR/TEST_DATA_FILE, orient='records', lines=True)
 
    print("-------数据预处理完成-------")
 
 
if __name__ == '__main__':
    preprocess()

2.5 自定义Dataset(dataset.py)

加载预处理好的索引文件,返回(src, tgt)张量。

import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
 
from config import *
 
from torch.nn.utils.rnn import pad_sequence     # 序列填充
 
# 自定义数据集类
class TranslationDataset(Dataset):
    # 初始化
    def __init__(self, path):
        # 定义属性,保存所有数据的字典列表
        self.data = pd.read_json(path, lines=True, orient='records').to_dict(orient='records')
 
    # 获取长度
    def __len__(self):
        return len(self.data)
 
    # 根据index获取元素
    def __getitem__(self, index):
        input = torch.tensor(self.data[index]['cn'], dtype=torch.long)
        target = torch.tensor(self.data[index]['en'], dtype=torch.long)
        return input, target
 
# 定义一个整理函数,将一批数据长度对齐(填充)
def collate_fn(batch):
    # batch形如 [ (input0, target0), (input1, target1),...];先分成inputs和targets两个列表
    input_tensor_list = [ item[0] for item in batch ]
    target_tensor_list = [ item[1] for item in batch ]
    # 合并成长度对齐的一个batch tensor
    input_batch_tensor = pad_sequence(input_tensor_list, batch_first=True, padding_value=0)
    target_batch_tensor = pad_sequence(target_tensor_list, batch_first=True, padding_value=0)
 
    return input_batch_tensor, target_batch_tensor
 
# 获取DataLoader的函数
def get_dataloader(train=True):
    path = PROCESSED_DATA_DIR / (TRAIN_DATA_FILE if train else TEST_DATA_FILE)
    dataset = TranslationDataset(path)
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)
    return dataloader
 
if __name__ == '__main__':
    train_dataloader = get_dataloader(train=True)
    test_dataloader = get_dataloader(train=False)
 
    # for input, target in train_dataloader:
    #     print(input.shape, target.shape)
    #     break
 
    data_iter = iter(train_dataloader)
    input_batch, target_batch = next(data_iter)
    print(input_batch.shape)
    print(target_batch.shape)
 
    input_batch, target_batch = next(data_iter)
    print(input_batch.shape)
    print(target_batch.shape)

2.6 位置编码与模型(model.py)

这是核心:实现位置编码,然后搭建TranslationModel,内部使用nn.Transformer。

import torch
import torch.nn as nn
 
from config import *
 
import math
 
# 自定义位置编码层
class PositionEncoding(nn.Module):
    # 初始化,生成位置编码矩阵 (L, E)
    def __init__(self, max_len, d_model):
        super().__init__()
        # 定义编码矩阵
        pe = torch.zeros(size=(max_len, d_model), dtype=torch.float)
        # 遍历矩阵每一行(每个位置 pos)
        for pos in range(max_len):
            # 遍历当前位置向量的每两个特征,步长为 2
            for _2i in range(0, d_model, 2):
                # 按公式计算向量里的这两个特征
                pe[pos, _2i] = math.sin( pos / (10000 ** (_2i / d_model)) )
                pe[pos, _2i + 1] = math.cos( pos / (10000 ** (_2i / d_model)) )
        self.register_buffer("pe", pe)
 
    # 前向传播
    def forward(self, x):
        # x 形状:(N, L, E=d_model)
        seq_len = x.shape[1]    # 提取当前序列长度 L
        # 在位置编码矩阵中截取 L 个向量,形状 (L, E)
        part_pe = self.pe[0:seq_len]
        return x + part_pe  # 广播叠加
 
class PositionEncoding_Pro(nn.Module):
    # 初始化,生成位置编码矩阵 (L, E)
    def __init__(self, max_len, d_model):
        super().__init__()
        # 定义编码矩阵
        pe = torch.zeros(size=(max_len, d_model))
        pos = torch.arange(0, max_len).unsqueeze(1) # (L, 1)
        _2i = torch.arange(0, d_model, 2)   # (d_model/2, )
 
        # 计算所有系数 (L, d_model/2)
        div_term = torch.pow(10000, (_2i / d_model))
 
        # 按奇偶数维度计算位置编码值
        pe[:, 0::2] = torch.sin( pos / div_term)
        pe[:, 1::2] = torch.cos( pos / div_term)
 
        self.register_buffer("pe", pe)
 
    # 前向传播
    def forward(self, x):
        # x 形状:(N, L, E=d_model)
        seq_len = x.shape[1]    # 提取当前序列长度 L
        # 在位置编码矩阵中截取 L 个向量,形状 (L, E)
        part_pe = self.pe[0:seq_len]
        return x + part_pe  # 广播叠加
 
 
# 自定义模型
class TranslationModel(nn.Module):
    # 初始化
    def __init__(self, cn_vocab_size, en_vocab_size, cn_padding_idx, en_padding_idx):
        super().__init__()
        # 词嵌入层
        self.cn_embedding = nn.Embedding(cn_vocab_size, embedding_dim=DIM_MODEL, padding_idx=cn_padding_idx)
        self.en_embedding = nn.Embedding(en_vocab_size, embedding_dim=DIM_MODEL, padding_idx=en_padding_idx)
        # 位置编码
        self.position_encoding = PositionEncoding(SEQ_LEN, DIM_MODEL)
        # Transformer层
        self.transformer = nn.Transformer(
            d_model=DIM_MODEL,
            nhead=NUM_HEADS,
            num_encoder_layers=NUM_ENCODER_LAYERS,
            num_decoder_layers=NUM_DECODER_LAYERS,
            batch_first=True,
        )
        # 输出线性层
        self.linear = nn.Linear(in_features=DIM_MODEL, out_features=en_vocab_size)
 
    # 前向传播,将 Transformer 需要的参数全部传入
    def forward(self, src, tgt, src_pad_mask, tgt_mask):
        # 输入源序列 (N, S),目标序列 (N, T)
        # 编码
        memory = self.encode(src, src_pad_mask)
        # 解码
        output = self.decode(tgt, memory, tgt_mask=tgt_mask, memory_pad_mask=src_pad_mask)
        return output
 
    # 编码方法
    def encode(self, src, src_pad_mask):
        # src形状: (N, S=src_len),src_pad_mask 形状:(N, S)
        # 1. 词嵌入
        embed = self.cn_embedding(src)
        # embed 形状:(N, S, E=d_model)
        # 2. 叠加位置编码
        input = self.position_encoding(embed)
        # input 形状:(N, S, E)
        # 3. Transformer编码器前向传播
        memory = self.transformer.encoder(src=input, src_key_padding_mask=src_pad_mask)
        # memory 形状:(N, S, E)
        return memory
 
    # 解码方法
    def decode(self, tgt, memory, tgt_mask=None, memory_pad_mask=None):
        # tgt形状: (N, T=tgt_len),tgt_mask 形状:(T, T)
        # 1. 词嵌入
        embed = self.en_embedding(tgt)
        # embed 形状:(N, T, E=d_model)
        # 2. 叠加位置编码
        input = self.position_encoding(embed)
        # input 形状:(N, T, E)
        # 3. Transformer解码器前向传播
        output = self.transformer.decoder(
            tgt = input,
            memory = memory,
            tgt_mask=tgt_mask,
            memory_key_padding_mask=memory_pad_mask,
        )
        # output 形状:(N, T, E)
        # 4. 经过输出线性层整合,得到预测输出
        output = self.linear(output)
        # output 形状:(N, T, en_vocab_size)
        return output
 
if __name__ == '__main__':
    # 定义模型
    model = TranslationModel(1000, 1024, 0, 0)
    print(model)

注意:训练时我们直接调用forward,它会自动执行编码器+解码器。推理时为了效率,我们先用encode得到memory,然后循环调用decode逐词生成。

2.7 训练脚本(train.py)

训练循环,包括生成掩码、计算损失、反向传播、保存最佳模型。

import torch
from torch import nn, optim
 
from tqdm import tqdm   # 进度条工具
 
from config import *
from dataset import get_dataloader  # 获取数据加载器
from model import TranslationModel  # 模型
 
from torch.utils.tensorboard import SummaryWriter   # 日志写入器
import time # 时间库
 
from tokenizer import ChineseTokenizer, EnglishTokenizer    # 分词器
 
# 定义训练引擎函数:训练一个epoch,返回平均损失
def train_one_epoch(model, train_loader, loss, optimizer, device):
    model.train()
 
    total_loss = 0
 
    # 按批次进行迭代
    for inputs, targets in tqdm(train_loader, desc='训练:'):
        inputs, targets = inputs.to(device), targets.to(device)    # 形状 (N=64, L)
        # 0. 准备参数
        # 0.1 基于目标序列,得到解码的输入和目标 (N, T=tgt_len)
        decoder_inputs = targets[:, :-1]
        decoder_targets = targets[:, 1:]
        # 0.2 源序列填充掩码,(N, S)
        src_pad_mask = (inputs == model.cn_embedding.padding_idx)
        # 0.3 目标序列自注意力掩码 (T, T)
        tgt_mask = model.transformer.generate_square_subsequent_mask( decoder_inputs.shape[1] )
 
        # 1. 前向传播,(N, T, en_vocab_size)
        decoder_outputs = model(src=inputs, tgt=decoder_inputs, src_pad_mask=src_pad_mask, tgt_mask=tgt_mask)
 
        # 2. 计算损失,输出形状 (N, vocab_size, L),目标形状 (N, L)
        loss_value = loss(decoder_outputs.transpose(1, 2), decoder_targets)
 
        # 3. 反向传播
        loss_value.backward()
        # 4. 更新参数
        optimizer.step()
        # 5. 梯度清零
        optimizer.zero_grad()
        # 累加损失
        total_loss += loss_value.item()
    return total_loss / len(train_loader)
 
# 训练整体流程
def train():
    # 1. 定义设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 2. 创建数据加载器
    train_loader = get_dataloader()
 
    # 3. 获取词表,创建分词器
    cn_tokenizer = ChineseTokenizer.from_vocab(MODEL_DIR/CN_VOCAB_FILE)
    en_tokenizer = EnglishTokenizer.from_vocab(MODEL_DIR/EN_VOCAB_FILE)
 
    # 4. 定义模型
    model = TranslationModel( cn_tokenizer.vocab_size, en_tokenizer.vocab_size, cn_tokenizer.pad_id, en_tokenizer.pad_id ).to(device)
 
    # 5. 定义损失函数
    loss = nn.CrossEntropyLoss(ignore_index=en_tokenizer.pad_id)
 
    # 6. 定义优化器
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
 
    # 7. 定义一个tensorboard写入器
    writer = SummaryWriter(log_dir=LOG_DIR / time.strftime('%Y-%m-%d_%H-%M-%S'))
 
    # 8. 核心训练流程,按epoch进行迭代
    min_loss = float('inf')     # 记录最小训练损失
    for epoch in range(EPOCHS):
        print("="*10, f"EPOCH:{epoch+1}", "="*10)
        this_loss = train_one_epoch(model, train_loader, loss, optimizer, device)
        print("本轮训练损失:", this_loss)
 
        # 将损失写入日志
        writer.add_scalar('loss', this_loss, epoch+1)
 
        # 判断损失是否下降,保存最优模型
        if this_loss < min_loss:
            min_loss = this_loss
            torch.save( model.state_dict(), MODEL_DIR / BEST_MODEL )
            print("模型保存成功!")
    # 关闭写入器
    writer.close()
 
if __name__ == '__main__':
    train()

2.8 推理脚本(predict.py)

实现交互式翻译。关键是用model.encode先编码源句子,然后循环调用model.decode,每次生成一个词,直到遇到或达到最大长度。

import torch
from config import *
from model import TranslationModel
from tokenizer import ChineseTokenizer, EnglishTokenizer
 
# 核心预测逻辑函数,返回一批数据的预测概率
def predict_batch(model, inputs, tokenizer, device):
    model.eval()
    # 前向传播
    with torch.no_grad():
        # 定义当前 batch_size (N)
        batch_size = inputs.shape[0]
 
        # 1. 前向传播
        # 1.1 编码
        src_pad_mask = (inputs == model.cn_embedding.padding_idx)
        memory = model.encode(inputs, src_pad_mask)
 
        # 1.2 解码:自回归生成
        # 1.2.1 构建第一时间步的输入,长度为 N 的向量 (N, T=1),内容全部为 <sos> 的 id
        decoder_input = torch.full(size=(batch_size, 1), fill_value=tokenizer.start_id).to(device)
 
        # 保存生成的id列表
        generated_ids = []
        # 定义一个长度为 N 的 tensor,保存每个样本是否已生成<eos>
        is_finished = torch.full(size=[batch_size], fill_value=False).to(device)
 
        # 1.2.2 循环迭代,自回归生成
        for i in range(SEQ_LEN):
            # (1) 解码,decoder_output 形状 (N, T, en_vocab_size)
            tgt_mask = model.transformer.generate_square_subsequent_mask(decoder_input.shape[1])
            decoder_output = model.decode(decoder_input, memory, tgt_mask=tgt_mask, memory_pad_mask=src_pad_mask)
            # (2) 词选择:贪心解码,得到预测下一个词的id (N, L=1)
            next_token_ids = torch.argmax(decoder_output[:, -1, :], dim=-1, keepdim=True)
            # (3) 保存预测id到生成列表中
            generated_ids.append(next_token_ids)
            # (4) 更新输入,形状增加 (N, T+1)
            decoder_input = torch.cat((decoder_input, next_token_ids), dim=-1)
            # (5) 判断是否生成 <eos>,如果一批全部生成<eos>则退出循环
            is_finished |= (next_token_ids.squeeze(1) == tokenizer.end_id)
            if is_finished.all():
                break
    # 处理生成结果
    # 基于生成列表 generated_ids: [ tensor(N, 1), tensor(N, 1), ...]
    # (1) 将列表转成 (N, L) 的张量
    generated_tensor =  torch.cat( generated_ids, dim=1 )
    # (2) 转换为二维列表
    generated_list = generated_tensor.tolist()
    # 形如:[ [*, *, *, eos], [*, *, eos, *], [*, *, *, *], ... ]
    # (3) 去掉每个元素(句子的id列表)中,eos之后的所有内容
    for i, sentence_ids in enumerate(generated_list):
        if tokenizer.end_id in sentence_ids:
            eos_pos = sentence_ids.index(tokenizer.end_id)
            generated_list[i] = sentence_ids[:eos_pos]
    # 形如:[ [*, *, *], [*, *], [*, *, *, *], ... ]
 
    return generated_list # 二维列表返回
 
def predict(text, model, cn_tokenizer, en_tokenizer, device):
    # 1. 准备数据:文本处理
    # 1.1/1.2 分词、id化
    ids = cn_tokenizer.encode(text)
    # 1.3 转换为 tensor,作为输入
    input = torch.tensor([ids], dtype=torch.long).to(device)
 
    # 2. 预测
    # 前向传播,得到预测概率
    result = predict_batch(model, input, en_tokenizer, device)
 
    return en_tokenizer.decode( result[0] )    # 只有唯一的一个数据,解码成英文句子
 
def run_predict():
    # 1. 确定设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
    # 2. 获取词表,得到分词器
    cn_tokenizer = ChineseTokenizer.from_vocab(MODEL_DIR / CN_VOCAB_FILE)
    en_tokenizer = EnglishTokenizer.from_vocab(MODEL_DIR / EN_VOCAB_FILE)
    print("词表加载成功!")
 
    # 3. 加载模型
    model = TranslationModel( cn_tokenizer.vocab_size, en_tokenizer.vocab_size, cn_tokenizer.pad_id, en_tokenizer.pad_id ).to(device)
 
    model.load_state_dict( torch.load( MODEL_DIR/BEST_MODEL ) )
    print("模型加载成功!")
 
    # 6. 程序运行流程
    print("欢迎使用中英翻译模型!输入q或者quit退出...")
    while True: # 核心:一个死循环
        # 捕获用户输入
        user_input = input("中文 > ")
        # 判断:如果是q或者quit,直接退出
        if user_input.strip() in ['q', 'quit']:
            print("欢迎下次再来!")
            break
        # 判断:如果是空白,提示信息后继续循环
        if user_input.strip() == '':
            print("请输入有效内容!")
            continue
 
        # 预测译文
        result = predict(user_input, model, cn_tokenizer, en_tokenizer, device)
        print("英文:", result)
 
if __name__ == '__main__':
    # text = "我们公司"
    # top5_tokens = predict(text)
    # print(top5_tokens)
    run_predict()

推理过程模拟了自回归生成:每一步把当前已生成的序列作为解码器输入,预测下一个词,然后拼接到输入中,直到遇到。

2.9 评估脚本(evaluate.py)

计算测试集上的BLEU分数,这是机器翻译的常用指标。

import torch
from tqdm import tqdm
from config import *
from model import TranslationModel
from dataset import get_dataloader
from predict import predict_batch    # 预测核心逻辑,得到批数据预测概率
from tokenizer import ChineseTokenizer, EnglishTokenizer
 
from nltk.translate.bleu_score import corpus_bleu   # 引入评价指标bleu
 
# 验证核心逻辑,返回评价指标(准确率)
def evaluate(model, dataloader, tokenizer, device):
    # 用列表记录参考译文和预测译文
    references = []
    predictions = []
 
    model.eval()
    with torch.no_grad():
        for inputs, targets in dataloader:
            inputs = inputs.to(device)
            targets = targets.tolist()  # 转成列表,方便计算
 
            # 前向传播,得到一批样本的预测结果
            batch_result = predict_batch(model, inputs, tokenizer, device)
            # 合并这一批结果到预测总列表
            predictions.extend( batch_result )
            # 合并这一批的目标值(参考译文)到总列表
            references.extend( [ [target[1:target.index(tokenizer.end_id)]] for target in targets ] )
 
    # 调库计算bleu评分
    bleu_score = corpus_bleu(references, predictions)
    return bleu_score
 
# 评估主流程
def run_evaluate():
    # 1. 确定设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
    # 2. 获取词表
    cn_tokenizer = ChineseTokenizer.from_vocab(MODEL_DIR / CN_VOCAB_FILE)
    en_tokenizer = EnglishTokenizer.from_vocab(MODEL_DIR / EN_VOCAB_FILE)
    print("词表加载成功!")
 
    # 3. 加载模型
    model = TranslationModel( cn_tokenizer.vocab_size, en_tokenizer.vocab_size, cn_tokenizer.pad_id, en_tokenizer.pad_id ).to(device)
 
    model.load_state_dict( torch.load( MODEL_DIR/BEST_MODEL ) )
    print("模型加载成功!")
 
    # 4. 获取测试数据集(加载器)
    test_dataloader = get_dataloader(train=False)
 
    # 5. 调用评估逻辑
    bleu = evaluate(model, test_dataloader, en_tokenizer, device)
 
    print("评估结果:")
    print("BLEU 评分: ", bleu)
 
if __name__ == '__main__':
    run_evaluate()

2.10 完整代码下载(包含数据集)

代码下载地址:pan.baidu.com/s/1yR5z1SDj…


三、运行指南

1、准备数据:下载cmn.txt放到data/raw/目录下(格式:英文\t中文)。

2、预处理

python preprocess.py

3、训练

python train.py

打印结果(EPOCHS我设置的70轮):

========== EPOCH:69 ==========
训练:: 100%|██████████| 365/365 [00:03<00:00, 92.58it/s]
本轮训练损失: 0.2390552392561142
模型保存成功!
========== EPOCH:70 ==========
训练:: 100%|██████████| 365/365 [00:03<00:00, 93.00it/s]
本轮训练损失: 0.23596478994578532
模型保存成功!

训练过程中可以用TensorBoard查看损失曲线:

tensorboard --logdir ./logs

4、交互式翻译

python predict.py

5、评估BLEU

python evaluate.py

打印结果:

词表加载成功!
模型加载成功!
评估结果:
BLEU 评分:  0.3229325229992989

四、总结与延伸

4.1 核心要点回顾

  • PyTorch的nn.Transformer

    :封装了完整的Encoder-Decoder,但需要自己提供位置编码、嵌入层和输出层。

  • 位置编码

    :使用正余弦函数,不参与训练,直接加到嵌入向量上。

  • 掩码机制

    :src_key_padding_mask屏蔽;tgt_mask屏蔽未来词;tgt_key_padding_mask同时屏蔽目标端的。

  • 训练与推理差异

    :训练时并行计算(一次输入整个目标序列),推理时自回归循环(每次只生成一个词)。

  • 数据预处理

    :英文句子加和,中文句子不加;统一固定长度,短填充长截断。

4.2 可以优化的方向

  1. 更大的模型

    :增大DIM_MODEL(如256或512)、增加层数(6层)、增加头数(8头),能显著提升BLEU分数。

  2. 学习率调度

    :使用NoamOpt(Transformer论文中的预热+衰减策略),训练更稳定。

  3. 标签平滑

    :减轻模型过度自信,提高泛化。

  4. 更优的分词

    :中文可以用jieba分词或BPE子词,英文用BPE减少未登录词。

  5. 束搜索(Beam Search)

    :推理时不用贪心,保留多个候选路径,提高翻译质量。

4.3 写在最后

这篇文章从理论到代码,把Transformer在PyTorch中的使用讲透了。你可以把这份代码当成模板,轻松迁移到其他Seq2Seq任务(如摘要、对话生成、代码生成)。下一章我们会深入源码级优化,并尝试在更大数据集上训练出实用模型。

如果你跑通了代码,欢迎在评论区留下你的BLEU分数。有任何问题,我会尽量回复。

别忘了点赞、收藏、转发,让更多人看到这篇干货!