使用 fairseq 训练判别式重排序模型 DrNMT：数据准备、训练与推理完整指南

使用 fairseq 训练判别式重排序模型 DrNMT数据准备、训练与推理完整指南【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq导读本文基于 fairseq 仓库中的 examples/discriminative_reranking_nmt 示例系统讲解如何训练 DrNMTDiscriminative Reranking for Neural Machine Translation——一个基于 XLM-R 初始化、面向机器翻译输出的判别式重排序器。文章覆盖从数据准备、BPE 与打分预处理、fairseq 二值化、Hydra 配置训练到推理阶段的前向模型fw得分与重排序得分融合、权重搜索tuning与 BLEU/TER 评估的完整闭环并深入对应源码task / model / criterion / 预处理与重排脚本解释每一步的底层原理。读完本文你将能够基于任意基础 MT 模型复现 DrNMT 的生成 N 条候选 → 训练重排序器 → 融合得分重排完整流程。DrNMT 方法简介DrNMT 是发表于 ACL 2021Lee, Auli, Ranzato《Discriminative Reranking for Neural Machine Translation》见 README 顶部引用的判别式重排序方法。其核心思路是在基础 MT 模型fw model用 beam search 生成 N 个候选译文之后训练一个独立的判别式打分器对每个源句的 N 个候选按目标指标BLEU 或 TER的相对优劣排序从而选出比单纯 beam search 最优解更接近参考译文的候选。与传统的在大型单语/平行语料上训练的生成式重排序语言模型不同DrNMT 直接以候选译文相对参考译文的指标分数为监督信号用 KL 散度训练一个分类式打分模型。示例目录中的完整组件包括scripts/prep_data.py把原始文本 N-best 候选转成 BPE 序列与指标分数标签tasks/discriminative_reranking_task.py定义discriminative_reranking_nmt任务负责多分片shard数据加载、以 beam 为单位的分组打乱、验证时用目标指标评估models/discriminative_reranking_model.py定义discriminative_nmt_reranker模型即基于 XLM-R 的 BertRankercriterions/discriminative_reranking_criterion.py定义kl_divergence_rereanking损失config/deen.yaml论文中 De→En 实验的 Hydra 训练配置drnmt_rerank.py推理/重排与权重搜索脚本。各子模块在__init__.py中通过 fairseq 的注册机制导出模型、任务、损失分别注册为discriminative_nmt_reranker、discriminative_reranking_nmt、kl_divergence_rereanking训练时通过common.user_dir指向本目录即可被 fairseq 发现。一、数据准备1.1 准备三份输入文件首先按照 examples/translation 的说明训练一个基础 MT 模型。随后准备三份纯文本文件每行一个句子不做任何 sentencepiece 或 tokenization 预处理文件行数内容源句文件L每个源句一行source_sentence_1 ... source_sentence_L参考译文文件L每个源句对应的 ground truth 目标句一行target_sentence_1 ... target_sentence_L候选译文文件L × N每个源句的 N 条候选按序连续排列source_sentence_1_hypo_1 ... source_sentence_1_hypo_N然后是source_sentence_2_hypo_1 ...N 为每句候选数候选译文由基础 MT 模型以 beamN 生成。论文中使用 N50。1.2 下载 XLM-R 预训练模型DrNMT 重排序器以 XLM-R 为骨架初始化因此需要下载 XLM-R base 模型README 中指向 examples/xlmr 的预训练模型说明wget https://dl.fbaipublicfiles.com/fairseq/models/xlmr.base.tar.gz tar zxvf xlmr.base.tar.gz解压后目录应包含三个文件dict.txt词表、model.pt预训练权重、sentencepiece.bpe.modelSPM 模型。1.3 生成 BPE 数据与指标分数运行scripts/prep_data.py为每个数据切分train / valid / test 等生成 BPE 化的输入和指标分数标签。先定义公共变量N每个源句的候选数论文用 50SPLIT数据切分名即train、valid、test若一个切分有多个数据集则用split_name, split_name1, split_name2, ...例如train, train1, valid, valid1NUM_SHARDS分片数非训练切分必须设为 1METRICDrNMT 要优化的指标支持bleu或ter。SOURCE_FILE/path/to/source_sentence_file TARGET_FILE/path/to/target_sentence_file HYPO_FILE/path/to/hypo_file XLMR_DIR/path/to/xlmr OUTPUT_DIR/path/to/output python scripts/prep_data.py \ --input-source ${SOURCE_FILE} \ --input-target ${TARGET_FILE} \ --input-hypo ${HYPO_FILE} \ --output-dir ${OUTPUT_DIR} \ --split $SPLIT \ --beam $N \ --sentencepiece-model ${XLMR_DIR}/sentencepiece.bpe.model \ --metric $METRIC \ --num-shards ${NUM_SHARDS}脚本会在${OUTPUT_DIR}/$METRIC下创建${NUM_SHARDS}个分片每个分片下split*/input_src、split*/input_tgt、split*/$METRIC目录中分别生成$SPLIT.bpe与$SPLIT.$METRIC文件。从 scripts/prep_data.py 源码看其内部流程为校验输入一致性源句与参考句数必须相等候选总数必须能被--beam整除且等于源句数 × beam否则直接断言报错对每个 (源句, 参考句, N 条候选) 三元组用 sentencepiece 把源句与每条候选编码为 BPE 序列sp.EncodeAsPieces用sacrebleu计算每条候选相对参考句的分数bleu时输出bleu.score sys_len ref_len加 4 组 counts/totals共 11 个数ter时输出ter.score num_edits ref_length共 3 个数支持--n-proc默认 8多进程并行处理--num-shards通过range(ns, num_sents, num_shards)把句子均匀分片便于超大训练集多机并行读取。1.4 fairseq-preprocess 二值化用 XLM-R 的dict.txt作为共享词典把 BPE 文本转成 fairseq 的二进制 indexed dataset。多个 train/valid 集合之间用逗号分隔# use comma to separate if there are more than one train or valid set for suffix in src tgt ; do fairseq-preprocess --only-source \ --trainpref ${OUTPUT_DIR}/$METRIC/split1/input_${suffix}/train.bpe \ --validpref ${OUTPUT_DIR}/$METRIC/split1/input_${suffix}/valid.bpe \ --destdir ${OUTPUT_DIR}/$METRIC/split1/input_${suffix} \ --workers 60 \ --srcdict ${XLMR_DIR}/dict.txt done for i in seq 2 ${NUM_SHARDS}; do for suffix in src tgt ; do fairseq-preprocess --only-source \ --trainpref ${OUTPUT_DIR}/$METRIC/split${i}/input_${suffix}/train.bpe \ --destdir ${OUTPUT_DIR}/$METRIC/split${i}/input_${suffix} \ --workers 60 \ --srcdict ${XLMR_DIR}/dict.txt ln -s ${OUTPUT_DIR}/$METRIC/split1/input_${suffix}/valid* ${OUTPUT_DIR}/$METRIC/split${i}/input_${suffix}/. done ln -s ${OUTPUT_DIR}/$METRIC/split1/$METRIC/valid* ${OUTPUT_DIR}/$METRIC/split${i}/$METRIC/. done第一段处理 split1 的 train 与 valid--only-source且提供--validpref后续分片只处理各自的 train并把 split1 的 valid 通过软链接共享过来——这也解释了为何非训练切分要求NUM_SHARDS1。注意标签文件$METRIC目录下的train.bleu/valid.bleu等不需要二值化task 会直接以RawLabelDataset读取原始分数文本。二、训练2.1 启动训练训练命令基于fairseq-hydra-train配置来自示例自带的 config/deen.yamlEXP_DIR/path/to/exp # An example of training the model with the config for De-En experiment in the paper. # The config uses 16 GPUs and 50 hypotheses. # For training with fewer number of GPUs, set # distributed_training.distributed_world_sizek optimization.update_freq[x] where x 16/k # For training with fewer number of hypotheses, set # task.mt_beamN dataset.batch_sizeN dataset.required_batch_size_multipleN fairseq-hydra-train -m \ --config-dir config/ --config-name deen \ task.data${OUTPUT_DIR}/$METRIC/split1/ \ task.num_data_splits${NUM_SHARDS} \ model.pretrained_model${XLMR_DIR}/model.pt \ common.user_dir${FAIRSEQ_ROOT}/examples/discriminative_reranking_nmt \ checkpoint.save_dir${EXP_DIR}-m表示多机/多卡multinode运行配合distributed_world_size16task.data指向 split1注意路径末尾的/task 会按需替换成其他分片见下文task.num_data_splits与prep_data.py的--num-shards一致model.pretrained_model指向 XLM-R 的model.ptcommon.user_dir让 fairseq 加载示例目录中注册的自定义 task / model / criterion若 GPU 数量少于 16设distributed_training.distributed_world_sizek且optimization.update_freq[x]x 16/k保持等效 batch若候选数少于 50设task.mt_beamN dataset.batch_sizeN dataset.required_batch_size_multipleN。2.2 deen.yaml 配置逐项解读config/deen.yaml 是论文 De→En 实验的配置各模块含义如下common: fp16: true # 混合精度训练 log_format: json log_interval: 50 seed: 2 checkpoint: no_epoch_checkpoints: true # 只保留 best checkpoint best_checkpoint_metric: bleu # 按验证集 BLEU 选最优 maximize_best_checkpoint_metric: true task: _name: discriminative_reranking_nmt data: ??? # 命令行注入task.data num_data_splits: ??? # 命令行注入task.num_data_splits include_src: true # 模型输入包含源句src tgt 拼接 mt_beam: 50 # 候选数与生成时的 beam 一致 eval_target_metric: true # 验证时用目标指标评估 target_metric: bleu # 优化目标bleu / ter dataset: batch_size: 50 num_workers: 6 required_batch_size_multiple: 50 # 必须为 mt_beam 的倍数 valid_subset: ??? criterion: _name: kl_divergence_rereanking target_dist_norm: minmax # 对目标分数做 min-max 归一化 temperature: 0.5 # 目标分布 softmax 温度 optimization: max_epoch: 200 lr: [0.00005] update_freq: [32] # 梯度累积等效扩大 batch optimizer: _name: adam adam_betas: (0.9,0.98) adam_eps: 1e-06 lr_scheduler: _name: polynomial_decay warmup_updates: 8000 total_num_update: 320000 model: _name: discriminative_nmt_reranker pretrained_model: ??? # 命令行注入model.pretrained_model classifier_dropout: 0.2 distributed_training: ddp_backend: no_c10d distributed_world_size: 16关于update_freq当 GPU 数从 16 降到 k 时设optimization.update_freq[x]x 16/k可保持每步等效 batch 不变。2.3 训练数据的加载逻辑task 源码视角discriminative_reranking_task.py 中的load_dataset揭示了几个关键设计分片轮转if self.cfg.data.endswith(1): data_shard (epoch - 1) % num_data_splits 1即训练数据路径以split1结尾时每个 epoch 自动轮转到不同分片split1 → split2 → ... → splitN从而实现多分片数据并行使用多 train 集合并训练切分按train、train1、train2… 依次探测存在则全部加载并用ConcatDataset拼接对应 README 中多个数据集的命名约定标签读取np.loadtxt(label_path)读取$METRIC下的分数文件若target_metric ter则取负TER 越小越好取负后统一为越大越好输入组织当include_srctrue时源句与目标句通过ConcatSentencesDataset拼接用 BOS/EOS 分隔并各自PrependTokenDataset加 BOS、TruncateDataset截断到max_positions默认 512以 beam 为单位打乱训练集 shuffle 不是逐句打乱而是np.arange(0, len(dataset), mt_beam)先打乱句组再在同一组内保持 N 条候选相邻SortDataset 构造的 shuffle 索引保证一个源句的 N 个候选始终同批出现批次约束加载后断言len(dataset) % mt_beam 0与 criterion 中的检查呼应。2.4 模型结构基于 XLM-R 的 BertRankerdiscriminative_reranking_model.py 中BertRanker的核心设计初始化通过hub_utils.from_pretrained加载 XLM-Rmodel.pt构建TransformerSentenceEncodernum_segments2额外学习第二组语言/分段 embedding并以strictFalse载入权重——注意update_init_roberta_model_state会剔除lm_head与version键、把layernorm_embedding重命名为emb_layer_norm以对齐结构句级表示sentence_forward支持head取s位置默认、meanpool、maxpool三种池化方式配置项model.sentence_rep可选transform_layer把 768 维投影到自定义embed_dim联合分类joint classification当model.joint_classificationsent时先对 beam 维做joint_forward——用若干层TransformerSentenceEncoderLayer让 N 条候选的表示互相交互再送入分类头论文中用于联合打分默认none即逐候选独立打分分类头复用 fairseq 的RobertaClassificationHeadtanh 激活 classifier_dropout输出 1 维分数另有freeze_embeddings、n_trans_layers_to_freeze等微调策略选项。2.5 损失函数KL 散度重排序损失criterions/discriminative_reranking_criterion.py 实现了kl_divergence_rereanking前向切块由于 beam 可能很大如 50forward_batch_size默认 32可通过criterion.forward_batch_size调小把样本切成小块做模型前向避免大 beam 下 OOM目标分布对每组 beam 的目标分数做minmax归一化(target - min) / (max - min eps)再除以temperature配置 0.5做softmax得到理想重排序分布模型分布对打分器 logits 做log_softmaxKL 损失loss -(target_dist * model_dist - target_dist * target_dist.log()).sum()即让模型分布逼近目标分布批量校验断言sample_size % task.cfg.mt_beam 0否则提示设置--required-batch-size-multiple。验证阶段若开启eval_target_metrictask 的valid_step会取出每句得分最高的候选汇总其 BLEU/TER 统计量通过sacrebleu.compute_bleu或 edits/ref_len 之比在训练日志中实时输出bleu/ter指标——这就是checkpoint.best_checkpoint_metric: bleu能够按验证 BLEU 选 best checkpoint 的原因。三、推理与重排打分3.1 用基础 MT 模型生成 N-best 候选先用fairseq-interactive对验证集/测试集源句生成 N 条候选。源句文件需要先用基础 MT 模型所用的 sentencepiece 编码每行一句。# genrate N hypotheses with the base MT model (fw score) VALID_SOURCE_FILE/path/to/source_sentences # one sentence per line, converted to the sentencepiece used by the base MT model VALID_TARGET_FILE/path/to/target_sentences # one sentence per line in raw text, i.e. no sentencepiece and tokenization MT_MODEL/path/to/mt_model MT_DATA_PATH/path/to/mt_data cat ${VALID_SOURCE_FILE} | \ fairseq-interactive ${MT_DATA_PATH} \ --max-tokens 4000 --buffer-size 16 \ --num-workers 32 --path ${MT_MODEL} \ --beam $N --nbest $N \ --post-process sentencepiece valid-hypo.out关键点--beam $N --nbest $N使每个源句输出 N 条候选输出重定向到文件如valid-hypo.out该文件包含S-源句与D-候选及 fw 分数行是下一步drnmt_rerank.py的输入。3.2 在验证集上搜索最优融合权重DrNMT 重排采用加权融合final_score fw_weight × (fw_score / tgt_len^lenpen) reranker_score。--tune模式在验证集上随机搜索fw_weight与lenpen# replace bleu with ter to optimize for TER python drnmt_rerank.py \ ${OUTPUT_DIR}/$METRIC/split1/ \ --path ${EXP_DIR}/checkpoint_best.pt \ --in-text valid-hypo.out \ --results-path ${EXP_DIR} \ --gen-subset valid \ --target-text ${VALID_TARGET_FILE} \ --user-dir ${FAIRSEQ_ROOT}/examples/discriminative_reranking_nmt \ --bpe sentencepiece \ --sentencepiece-model ${XLMR_DIR}/sentencepiece.bpe.model \ --beam $N \ --batch-size $N \ --metric bleu \ --tune位置参数为数据目录split1/仅用于加载词典与 tokenizer/bpe 配置--target-text在 tune 模式下必须提供脚本中显式断言否则只能输出得分最高的候选而无法计算指标。3.3 用最优权重重排测试集# genrate N hypotheses with the base MT model (fw score) TEST_SOURCE_FILE/path/to/source_sentences # one sentence per line, converted to the sentencepiece used by the base MT model cat ${TEST_SOURCE_FILE} | \ fairseq-interactive ${MT_DATA_PATH} \ --max-tokens 4000 --buffer-size 16 \ --num-workers 32 --path ${MT_MODEL} \ --beam $N --nbest $N \ --post-process sentencepiece test-hypo.out # replace bleu with ter to evaluate TER # Add --target-text for evaluating BLEU/TER, # otherwise the script will only generate the hypotheses with the highest scores only. python drnmt_rerank.py \ ${OUTPUT_DIR}/$METRIC/split1/ \ --path ${EXP_DIR}/checkpoint_best.pt \ --in-text test-hypo.out \ --results-path ${EXP_DIR} \ --gen-subset test \ --user-dir ${FAIRSEQ_ROOT}/examples/discriminative_reranking_nmt \ --bpe sentencepiece \ --sentencepiece-model ${XLMR_DIR}/sentencepiece.bpe.model \ --beam $N \ --batch-size $N \ --metric bleu \ --fw-weight ${BEST_FW_WEIGHT} \ --lenpen ${BEST_LENPEN}测试阶段不再--tune而是直接传入验证阶段搜索到的最优--fw-weight与--lenpen。若提供--target-text脚本会同时打印重排前后的BLEU/TER对比否则仅把最高分候选写入${EXP_DIR}/generate-${gen_subset}.txt每行序号\t分数\t候选。3.4 drnmt_rerank.py 内部流程从 drnmt_rerank.py 源码看推理分四步解析 fairseq-interactive 输出parse_fairseq_gen解析S-/D-行得到源句列表、按句分组的候选列表与 fw 分数列表同时校验len(src) * beam len(hyp)重排序器打分make_batches把 (src, hypo) 对按include_src与否组织成task.build_dataset_for_inference的输入若包含源句用 EOS 分隔拼接对应模型中的get_segment_labels分段逻辑经task.inference_step(generator, ...)得到每条候选的重排序分数融合与重排get_best_hyps按fw_score / tgt_len^lenpen × fw_weight model_score计算每条候选的最终分每个 beam 组内取 argmax 作为最佳候选权重搜索--tune在--lower/upper-bound-fw-weight默认 0~3与--lower/upper-bound-lenpen默认 0~3范围内做--num-trials默认 1000次随机搜索用 32 进程multiprocessing.Pool并行评估BLEU 取最大值、TER 取最小值对应的权重即为最优。四、端到端流程小结与注意事项完整的 DrNMT 使用链路为训练基础 MT 模型见 examples/translation用 beamN 对 train/valid/test 生成 N-best 候选整理成源句、参考句、候选三份纯文本文件下载 XLM-R base运行scripts/prep_data.py生成 BPE 输入与 BLEU/TER 标签train 可多分片fairseq-preprocess二值化 src/tgt软链接共享 valid 数据fairseq-hydra-train加载config/deen.yaml训练重排序器注意task.data、task.num_data_splits、model.pretrained_model、common.user_dir四个必填项以及减少 GPU/beam 时的对应调整用fairseq-interactive生成验证集/测试集 N-bestdrnmt_rerank.py --tune在验证集搜权重再以最优权重对测试集重排并报告 BLEU/TER 提升。实操中的关键约束均有源码/脚本断言支撑候选文件行数必须等于源句数 × beam否则prep_data.py与drnmt_rerank.py都会直接报错dataset.batch_size与required_batch_size_multiple必须等于mt_beam训练切分可用多分片与train1/train2...命名扩展valid/test 必须NUM_SHARDS1TER 作为优化目标时标签在 task 中取负以统一越大越好的方向验证评估也相应以编辑数/参考长度计算。如需复现论文实验可参照论文配置 config/deen.yaml 中的超参16 GPU、beam 50、lr 5e-5、warmup 8000、polynomial_decay 至 320k 步、minmax 归一化 0.5 温度、classifier_dropout 0.2、seed 2。引用inproceedings{lee2021discriminative, title{Discriminative Reranking for Neural Machine Translation}, author{Lee, Ann and Auli, Michael and Ranzato, MarcAurelio}, booktitle{ACL}, year{2021} }【免费下载链接】fairseqFacebook AI Research Sequence-to-Sequence Toolkit written in Python.项目地址: https://gitcode.com/gh_mirrors/fa/fairseq创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考