注:接上一篇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 可以优化的方向
-
更大的模型
:增大DIM_MODEL(如256或512)、增加层数(6层)、增加头数(8头),能显著提升BLEU分数。
-
学习率调度
:使用NoamOpt(Transformer论文中的预热+衰减策略),训练更稳定。
-
标签平滑
:减轻模型过度自信,提高泛化。
-
更优的分词
:中文可以用jieba分词或BPE子词,英文用BPE减少未登录词。
-
束搜索(Beam Search)
:推理时不用贪心,保留多个候选路径,提高翻译质量。
4.3 写在最后
这篇文章从理论到代码,把Transformer在PyTorch中的使用讲透了。你可以把这份代码当成模板,轻松迁移到其他Seq2Seq任务(如摘要、对话生成、代码生成)。下一章我们会深入源码级优化,并尝试在更大数据集上训练出实用模型。
如果你跑通了代码,欢迎在评论区留下你的BLEU分数。有任何问题,我会尽量回复。
别忘了点赞、收藏、转发,让更多人看到这篇干货!