PyTorch Fairseq 神经机器翻译实践指南

概述

Fairseq 是 PyTorch 生态中一个强大的序列建模工具包，特别擅长神经机器翻译(NMT)任务。本文将详细介绍如何使用 Fairseq 进行机器翻译，包括预训练模型的使用和新模型的训练方法。

预训练模型使用

Fairseq 提供了多种预训练的翻译模型，涵盖不同架构和语言对：

模型架构类型

1. 卷积神经网络(CNN)模型：

   • 基于 Gehring 等人 2017 年提出的架构
   • 适用于 WMT14 英法、英德等语言对

2. Transformer 模型：

   • 基于 Vaswani 等人提出的经典架构
   • 包含 WMT14-19 多个比赛获胜模型
   • 支持英德、德英、英俄、俄英等多种语言方向

使用预训练模型

使用 PyTorch Hub 可以轻松加载预训练模型：

import torch

# 加载 WMT16 英德 Transformer 模型
en2de = torch.hub.load('pytorch/fairseq', 'transformer.wmt16.en-de',
                      tokenizer='moses', bpe='subword_nmt')
en2de.eval().cuda()  # 启用评估模式并移到 GPU

# 单句翻译
print(en2de.translate('Hello world!'))  # 输出: Hallo Welt!

# 批量翻译
print(en2de.translate(['How are you?', 'This is a test.']))

对于 WMT19 模型，需要使用 fastBPE 而非 subword_nmt：

en2de = torch.hub.load('pytorch/fairseq', 'transformer.wmt19.en-de',
                      checkpoint_file='model1.pt:model2.pt:model3.pt:model4.pt',
                      tokenizer='moses', bpe='fastbpe')

训练新模型

数据准备

训练新模型前需要准备和预处理数据。Fairseq 提供了多个数据准备脚本：

1. IWSLT'14 德英数据集：

   cd examples/translation/
   bash prepare-iwslt14.sh
   cd ../..
   

2. WMT'14 英德数据集：

   cd examples/translation/
   bash prepare-wmt14en2de.sh  # 或添加 --icml17 参数使用 WMT14 原始数据
   cd ../..
   

数据预处理

使用 fairseq-preprocess 进行数据二进制化：

TEXT=examples/translation/iwslt14.tokenized.de-en
fairseq-preprocess --source-lang de --target-lang en \
    --trainpref $TEXT/train --validpref $TEXT/valid --testpref $TEXT/test \
    --destdir data-bin/iwslt14.tokenized.de-en \
    --workers 20

模型训练

Transformer 模型训练示例

fairseq-train data-bin/iwslt14.tokenized.de-en \
    --arch transformer_iwslt_de_en --share-decoder-input-output-embed \
    --optimizer adam --adam-betas '(0.9, 0.98)' --clip-norm 0.0 \
    --lr 5e-4 --lr-scheduler inverse_sqrt --warmup-updates 4000 \
    --dropout 0.3 --weight-decay 0.0001 \
    --criterion label_smoothed_cross_entropy --label-smoothing 0.1 \
    --max-tokens 4096 \
    --eval-bleu \
    --eval-bleu-args '{"beam": 5, "max_len_a": 1.2, "max_len_b": 10}' \
    --eval-bleu-detok moses \
    --eval-bleu-remove-bpe \
    --eval-bleu-print-samples \
    --best-checkpoint-metric bleu --maximize-best-checkpoint-metric

CNN 模型训练示例

fairseq-train data-bin/wmt17_en_de \
    --arch fconv_wmt_en_de \
    --dropout 0.2 \
    --criterion label_smoothed_cross_entropy --label-smoothing 0.1 \
    --optimizer nag --clip-norm 0.1 \
    --lr 0.5 --lr-scheduler fixed --force-anneal 50 \
    --max-tokens 4000 \
    --save-dir checkpoints/fconv_wmt_en_de

模型评估

训练完成后，使用以下命令评估模型：

fairseq-generate data-bin/iwslt14.tokenized.de-en \
    --path checkpoints/checkpoint_best.pt \
    --batch-size 128 --beam 5 --remove-bpe

多语言翻译

Fairseq 支持多语言翻译模型的训练。以下是一个德法到英语的多语言翻译示例：

数据准备

cd examples/translation/
bash prepare-iwslt17-multilingual.sh
cd ../..

二进制化处理

# 德英数据
TEXT=examples/translation/iwslt17.de_fr.en.bpe16k
fairseq-preprocess --source-lang de --target-lang en \
    --trainpref $TEXT/train.bpe.de-en \
    --validpref $TEXT/valid0.bpe.de-en,$TEXT/valid1.bpe.de-en,$TEXT/valid2.bpe.de-en,$TEXT/valid3.bpe.de-en,$TEXT/valid4.bpe.de-en,$TEXT/valid5.bpe.de-en \
    --destdir data-bin/iwslt17.de_fr.en.bpe16k \
    --workers 10

# 法英数据（重用英文字典）
fairseq-preprocess --source-lang fr --target-lang en \
    --trainpref $TEXT/train.bpe.fr-en \
    --validpref $TEXT/valid0.bpe.fr-en,$TEXT/valid1.bpe.fr-en,$TEXT/valid2.bpe.fr-en,$TEXT/valid3.bpe.fr-en,$TEXT/valid4.bpe.fr-en,$TEXT/valid5.bpe.fr-en \
    --tgtdict data-bin/iwslt17.de_fr.en.bpe16k/dict.en.txt \
    --destdir data-bin/iwslt17.de_fr.en.bpe16k \
    --workers 10

多语言模型训练

fairseq-train data-bin/iwslt17.de_fr.en.bpe16k/ \
    --max-epoch 50 \
    --ddp-backend=legacy_ddp \
    --task multilingual_translation --lang-pairs de-en,fr-en \
    --arch multilingual_transformer_iwslt_de_en \
    --share-decoders --share-decoder-input-output-embed \
    --optimizer adam --adam-betas '(0.9, 0.98)' \
    --lr 0.0005 --lr-scheduler inverse_sqrt \
    --warmup-updates 4000 --warmup-init-lr '1e-07' \
    --label-smoothing 0.1 --criterion label_smoothed_cross_entropy \
    --dropout 0.3 --weight-decay 0.0001 \
    --save-dir checkpoints/multilingual_transformer \
    --max-tokens 4000 \
    --update-freq 8

最佳实践

1. 硬件选择：Transformer 模型训练推荐使用 GPU，显存越大越好
2. 批量大小：根据显存调整 --max-tokens 参数
3. 学习率：Transformer 通常使用较小的学习率(5e-4)，CNN 可以使用较大学习率(0.5)
4. 正则化：合理使用 dropout(0.1-0.3)和权重衰减(0.0001)
5. 评估指标：训练时启用 BLEU 评估可以更好地监控模型性能

通过本文介绍的方法，您可以轻松使用 Fairseq 进行各种机器翻译任务的实验和部署。