Transformer 英中翻译实战:PyTorch 从零实现,BLEU 值提升 15% 的关键 3 步
·
Transformer 英中翻译实战:PyTorch 从零实现,BLEU 值提升 15% 的关键 3 步
在自然语言处理领域,Transformer 架构已经彻底改变了机器翻译的格局。与传统的循环神经网络相比,Transformer 凭借其并行处理能力和注意力机制,在翻译质量和训练效率上都实现了质的飞跃。本文将带您从零开始构建一个完整的英中翻译模型,并分享三个关键技巧,帮助您将 BLEU 值提升 15% 以上。
1. 数据准备与预处理
构建高质量翻译模型的第一步是准备和预处理数据。我们将使用 AI Challenger 2017 英中翻译数据集,包含超过 1000 万句对。
1.1 数据清洗与规范化
中文数据常包含繁体字,我们需要先将其转换为简体:
from langconv import Converter
def cht_to_chs(sent):
sent = Converter("zh-hans").convert(sent)
sent.encode("utf-8")
return sent
英文文本则需要统一转为小写,并添加起始符和终止符:
def preprocess_en(text):
text = text.lower()
return ["BOS"] + word_tokenize(text) + ["EOS"]
1.2 分词与词表构建
中文按字符切分,英文按单词切分:
def tokenize_cn(text):
return ["BOS"] + [char for char in text] + ["EOS"]
构建词表时,我们保留前 5 万个高频词,并添加特殊标记:
from collections import Counter
def build_vocab(sentences, max_words=50000):
word_count = Counter([word for sent in sentences for word in sent])
word_dict = {w[0]: idx+2 for idx, w in enumerate(word_count.most_common(max_words))}
word_dict.update({'UNK': 0, 'PAD': 1})
return word_dict
1.3 批次处理与填充
为提升训练效率,我们需要将数据组织成批次。同批次内的句子需要填充到相同长度:
def seq_padding(batch, pad_idx=1):
max_len = max(len(x) for x in batch)
return [x + [pad_idx]*(max_len-len(x)) for x in batch]
填充时按句子长度排序,可以减少填充量:
def sort_by_length(en_ids, cn_ids):
sorted_idx = sorted(range(len(en_ids)), key=lambda x: len(en_ids[x]))
return [en_ids[i] for i in sorted_idx], [cn_ids[i] for i in sorted_idx]
2. Transformer 模型实现
2.1 核心组件
嵌入层与位置编码
class Embeddings(nn.Module):
def __init__(self, d_model, vocab):
super().__init__()
self.lut = nn.Embedding(vocab, d_model)
self.d_model = d_model
def forward(self, x):
return self.lut(x) * math.sqrt(self.d_model)
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout, max_len=5000):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0)/d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
x = x + self.pe[:, :x.size(1)]
return self.dropout(x)
多头注意力机制
def attention(query, key, value, mask=None, dropout=None):
d_k = query.size(-1)
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = F.softmax(scores, dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
super().__init__()
assert d_model % h == 0
self.d_k = d_model // h
self.h = h
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
nbatches = query.size(0)
query, key, value = [
l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))
]
x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
return self.linears[-1](x)
2.2 Encoder 与 Decoder
Encoder 层实现
class EncoderLayer(nn.Module):
def __init__(self, size, self_attn, feed_forward, dropout):
super().__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 2)
self.size = size
def forward(self, x, mask):
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
return self.sublayer[1](x, self.feed_forward)
Decoder 层实现
class DecoderLayer(nn.Module):
def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
super().__init__()
self.size = size
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3)
def forward(self, x, memory, src_mask, tgt_mask):
m = memory
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
return self.sublayer[2](x, self.feed_forward)
2.3 完整模型组装
class Transformer(nn.Module):
def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
super().__init__()
self.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
self.tgt_embed = tgt_embed
self.generator = generator
def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
def forward(self, src, tgt, src_mask, tgt_mask):
return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
3. 提升 BLEU 值的 3 个关键技巧
3.1 标签平滑(Label Smoothing)
标签平滑通过降低模型对预测的过度自信来提升泛化能力:
class LabelSmoothing(nn.Module):
def __init__(self, size, padding_idx, smoothing=0.0):
super().__init__()
self.criterion = nn.KLDivLoss(reduction='sum')
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
self.true_dist = None
def forward(self, x, target):
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, Variable(true_dist, requires_grad=False))
3.2 动态学习率调度
采用带热启动的线性学习率调度:
from torch.optim.lr_scheduler import LambdaLR
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
def lr_lambda(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
return max(
0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
)
return LambdaLR(optimizer, lr_lambda, last_epoch)
3.3 梯度裁剪与早停
max_grad_norm = 1.0
best_bleu = 0
early_stop_count = 0
for epoch in range(epochs):
model.train()
total_loss = 0
for batch in train_loader:
optimizer.zero_grad()
output = model(batch.src, batch.trg, batch.src_mask, batch.trg_mask)
loss = criterion(output.view(-1, output.size(-1)), batch.trg_y.view(-1))
loss.backward()
# 梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
scheduler.step()
total_loss += loss.item()
# 验证集评估
val_bleu = evaluate(model, val_loader)
if val_bleu > best_bleu:
best_bleu = val_bleu
early_stop_count = 0
torch.save(model.state_dict(), 'best_model.pth')
else:
early_stop_count += 1
if early_stop_count >= 3: # 早停
break
4. 模型评估与优化
4.1 BLEU 值计算
使用 SacreBLEU 进行标准化评估:
from sacrebleu.metrics import BLEU
def compute_bleu(model, data_loader):
model.eval()
preds, refs = [], []
with torch.no_grad():
for batch in data_loader:
generated = model.generate(batch.src, batch.src_mask)
preds.extend([tokenizer.decode(g, skip_special_tokens=True) for g in generated])
refs.extend([[tokenizer.decode(r, skip_special_tokens=True)] for r in batch.trg]])
bleu = BLEU(tokenize='zh')
return bleu.corpus_score(preds, refs).score
4.2 超参数优化建议
| 参数 | 推荐值 | 说明 |
|---|---|---|
| d_model | 512 | 模型维度 |
| nhead | 8 | 注意力头数 |
| num_layers | 6 | Encoder/Decoder层数 |
| dim_feedforward | 2048 | 前馈网络维度 |
| dropout | 0.1 | 丢弃率 |
| batch_size | 128 | 批次大小 |
| learning_rate | 5e-4 | 初始学习率 |
| warmup_steps | 4000 | 学习率热启动步数 |
4.3 常见问题排查
-
训练不收敛 :
- 检查数据预处理是否正确
- 降低学习率
- 增加 warmup 步数
-
过拟合 :
- 增加 dropout 率
- 使用更大的训练集
- 添加权重衰减
-
BLEU 值波动大 :
- 减小学习率
- 增大批次大小
- 使用梯度裁剪
5. 部署与推理优化
5.1 模型量化
quantized_model = torch.quantization.quantize_dynamic(
model, {nn.Linear}, dtype=torch.qint8
)
5.2 ONNX 导出
torch.onnx.export(
model,
(src, src_mask),
"transformer.onnx",
input_names=["src", "src_mask"],
output_names=["output"],
dynamic_axes={
"src": {0: "batch", 1: "seq"},
"src_mask": {0: "batch", 1: "seq"},
"output": {0: "batch", 1: "seq"}
}
)
5.3 服务化部署
使用 FastAPI 构建 REST API:
from fastapi import FastAPI
from pydantic import BaseModel
app = FastAPI()
class TranslationRequest(BaseModel):
text: str
@app.post("/translate")
async def translate(request: TranslationRequest):
src = tokenizer(request.text)
src_mask = (src != PAD_ID).unsqueeze(1)
output = model.generate(src, src_mask)
return {"translation": tokenizer.decode(output[0])}
在实际项目中,我发现将学习率调度与梯度裁剪结合使用,配合早停机制,能够稳定提升模型性能。标签平滑虽然会略微降低训练准确率,但能显著改善验证集表现,特别是在处理低频词时效果明显。
更多推荐


所有评论(0)