用 Fairseq 微调 RoBERTa 完成自定义文本分类任务:IMDB 情感分类全流程实战

发布时间:2026/9/14 19:02:27
用 Fairseq 微调 RoBERTa 完成自定义文本分类任务:IMDB 情感分类全流程实战 用 Fairseq 微调 RoBERTa 完成自定义文本分类任务IMDB 情感分类全流程实战【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本文基于 unilm 仓库中 kosmos-2 附带的 fairseq 示例文档kosmos-2/fairseq/examples/roberta/README.custom_classification.md完整讲解如何把 RoBERTa 预训练模型微调到一个自定义文本分类任务上以 IMDB 电影评论情感二分类为具体实例。读完后你可以独立走通“原始文本数据 → 格式化 → BPE 编码 → fairseq 二进制预处理 → 训练命令 → 用 hub 接口加载 checkpoint 推理”的完整链路并理解每一步背后的 fairseq 源码实现sentence_prediction任务、分类头注册与 hub 推理接口。一、这条自定义分类流水线要解决什么问题fairseq 中 RoBERTa 的微调示例默认面向 GLUE、WSC 等标准基准而实际业务中更多是“自己手里有一批文本 标签”。这份文档给出的方案是把任意分类数据组织成 fairseq 的sentence_prediction任务所需的目录结构然后复用fairseq-train的通用训练框架用--restore-file加载 RoBERTa 预训练权重只训练一个随机初始化的分类头classification head。整条链路由 6 个阶段构成阶段命令 / 工具产物1. 获取数据wget/taraclImdb/原始目录2. 格式化文档中的 Python 脚本train.input0/train.label/dev.input0/dev.label3. BPE 编码multiprocessing_bpe_encoder*.input0.bpetoken id 文本4. 二进制预处理fairseq-preprocess两次调用IMDB-bin/input0/与IMDB-bin/label/的.bin/.idx5. 训练fairseq-traincheckpoints/checkpoint_best.pt6. 推理RobertaModel.from_pretrainedpredict分类预测下面逐步展开。二、阶段 1获取并解压 IMDB 数据IMDB 数据集来自 Stanford 的 aclImdb 发布包下载后是一个“一条样本一个文件”的目录树train/pos、train/neg、test/pos、test/negwget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz tar zxvf aclImdb_v1.tar.gz解压后的结构为aclImdb/{train,test}/{pos,neg}/*.txt每个.txt文件包含一条影评所在目录名即类别标签。这种“文件即样本”的组织方式正是下一步格式化脚本要处理的输入形态。三、阶段 2把“文件即样本”转成两个对齐的文本文件fairseq 的sentence_prediction任务要求输入是行对齐的文本文件或后续的二进制索引文件*.input0存放样本文本*.label存放对应的标签整数且两文件行数一一对应。文档给出的转换脚本如下完整保留import argparse import os import random from glob import glob random.seed(0) def main(args): for split in [train, test]: samples [] for class_label in [pos, neg]: fnames glob(os.path.join(args.datadir, split, class_label) /*.txt) for fname in fnames: with open(fname) as fin: line fin.readline() samples.append((line, 1 if class_label pos else 0)) random.shuffle(samples) out_fname train if split train else dev f1 open(os.path.join(args.datadir, out_fname .input0), w) f2 open(os.path.join(args.datadir, out_fname .label), w) for sample in samples: f1.write(sample[0] \n) f2.write(str(sample[1]) \n) f1.close() f2.close() if __name__ __main__: parser argparse.ArgumentParser() parser.add_argument(--datadir, defaultaclImdb) args parser.parse_args() main(args)几个细节值得注意标签映射pos → 1neg → 0写入.label文件的是纯整数。后续 fairseq 会把这些行当作“词典 token”处理见第五节label词典的加载方式。随机打乱random.seed(0)固定随机种子后用random.shuffle(samples)打乱样本保证每次运行结果可复现。split 重命名原始数据里的test在这里被重命名为dev验证集。这是 fine-tuning 阶段的惯例——dev集用于训练中挑选最佳 checkpoint--best-checkpoint-metric accuracy真正的test评测留到最后做。产出文件aclImdb/train.input0、aclImdb/train.label、aclImdb/dev.input0、aclImdb/dev.label。注意每条样本只读取文件第一行fin.readline()因为 IMDB 的每条影评实际只占一行。换自己的数据集时只需把这里的读取与标签映射逻辑替换掉输出保持“一行一样本、input0 与 label 行对齐”即可。四、阶段 3用 GPT-2 BPE 对文本做分词编码RoBERTa 的分词遵循 GPT-2 BPE必须使用官方发布的encoder.json与vocab.bpe两份文件。fairseq 提供了多线程编码工具examples/roberta/multiprocessing_bpe_encoder.py它把--inputs指定的原始文本逐行转成以空格分隔的 token id 序列对应实现见 multiprocessing_bpe_encoder.py 中的MultiprocessingEncoder.encode_lines。# Download encoder.json and vocab.bpe wget -N https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json wget -N https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe for SPLIT in train dev; do python -m examples.roberta.multiprocessing_bpe_encoder \ --encoder-json encoder.json \ --vocab-bpe vocab.bpe \ --inputs aclImdb/$SPLIT.input0 \ --outputs aclImdb/$SPLIT.input0.bpe \ --workers 60 \ --keep-empty done命令行参数说明均可在 multiprocessing_bpe_encoder.py 的argparse定义中确认--encoder-json/--vocab-bpeGPT-2 BPE 词表文件路径--inputs/--outputs支持多个输入输出对要求数量一致-表示标准输入输出--workers多进程 worker 数默认 20文档示例用 60--keep-empty保留空行。不加该选项时遇到空行该样本会被标记为EMPTY并过滤掉见 源码 L108-L119——分类任务中一旦过滤了某一行input 与 label 就会行错位因此这里必须加--keep-empty。文档也指出BPE 编码完全可以合并进上一步的格式化脚本中对每条样本直接执行只是那样更慢所以实践中把编码独立成一步、用多进程批量跑更快。注意本步只编码input0label文件里的整数不需要 BPE 编码直接进入下一步的预处理。五、阶段 4fairseq-preprocess 生成二进制数据这一步把文本文件转成 fairseq 的二进制索引格式.bin/.idx训练时才能被高效加载。需要两次调用分别处理输入文本和标签并各自建立词典# Download fairseq dictionary. wget -N https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/dict.txt fairseq-preprocess \ --only-source \ --trainpref aclImdb/train.input0.bpe \ --validpref aclImdb/dev.input0.bpe \ --destdir IMDB-bin/input0 \ --workers 60 \ --srcdict dict.txt fairseq-preprocess \ --only-source \ --trainpref aclImdb/train.label \ --validpref aclImdb/dev.label \ --destdir IMDB-bin/label \ --workers 60两次调用的差异第一次IMDB-bin/input0--srcdict dict.txt指定使用随 RoBERTa 发布的官方 GPT-2 词典49152 词元。这样预处理后的 token id 与预训练模型 embedding 的 id 空间完全一致--restore-file加载权重时才能对得上无需重建词表。第二次IMDB-bin/label不指定词典由 fairseq 从train.label/dev.label中自动建一个小词典对 IMDB 而言就是0和1两个符号。从源码可以确认这个目录布局的要求sentence_prediction.py 的setup_task会分别从data/input0/dict.txt和data/label/dict.txt加载输入词典与标签词典并要求--num-classes 0load_dataset则依次加载input0必需、input1可选用于句子对和label三个数据集L132-L256。因此IMDB-bin/下必须有input0/与label/两个子目录且label的标签值必须是可被其词典表示的 token。六、阶段 5运行训练这是整条流水线的核心命令文档完整给出如下含全部环境变量TOTAL_NUM_UPDATES7812 # 10 epochs through IMDB for bsz 32 WARMUP_UPDATES469 # 6 percent of the number of updates LR1e-05 # Peak LR for polynomial LR scheduler. HEAD_NAMEimdb_head # Custom name for the classification head. NUM_CLASSES2 # Number of classes for the classification task. MAX_SENTENCES8 # Batch size. ROBERTA_PATH/path/to/roberta.large/model.pt CUDA_VISIBLE_DEVICES0 fairseq-train IMDB-bin/ \ --restore-file $ROBERTA_PATH \ --max-positions 512 \ --batch-size $MAX_SENTENCES \ --max-tokens 4400 \ --task sentence_prediction \ --reset-optimizer --reset-dataloader --reset-meters \ --required-batch-size-multiple 1 \ --init-token 0 --separator-token 2 \ --arch roberta_large \ --criterion sentence_prediction \ --classification-head-name $HEAD_NAME \ --num-classes $NUM_CLASSES \ --dropout 0.1 --attention-dropout 0.1 \ --weight-decay 0.1 --optimizer adam --adam-betas (0.9, 0.98) --adam-eps 1e-06 \ --clip-norm 0.0 \ --lr-scheduler polynomial_decay --lr $LR --total-num-update $TOTAL_NUM_UPDATES --warmup-updates $WARMUP_UPDATES \ --fp16 --fp16-init-scale 4 --threshold-loss-scale 1 --fp16-scale-window 128 \ --max-epoch 10 \ --best-checkpoint-metric accuracy --maximize-best-checkpoint-metric \ --shorten-method truncate \ --find-unused-parameters \ --update-freq 4关键参数逐项解读调度与批量参数取值作用TOTAL_NUM_UPDATES7812多项式衰减的总优化步数等于“10 个 epoch × IMDB 训练集 25000 条 ÷ 有效 batch 32”取整WARMUP_UPDATES469warmup 步数约为总步数的 6%--lr1e-05多项式调度的峰值学习率--batch-size/--update-freq8 / 4有效 batch 为 8×4 32 句--max-tokens4400单步最大 token 数防止长句把显存打爆--max-epoch10最多训练 10 个 epoch任务与模型参数取值作用--tasksentence_prediction使用上文分析的分类任务数据目录为第一个位置参数IMDB-bin/--restore-fileRoBERTa 预训练权重加载 encoder 权重重置优化器/数据加载器/计步器--reset-optimizer --reset-dataloader --reset-meters以支持从头开始微调--archroberta_large355M 参数的大模型架构--num-classes2分类类别数IMDB 二分类SentencePredictionTask.setup_task强制要求num_classes 0--classification-head-nameimdb_head分类头注册名推理时roberta.predict(imdb_head, ...)用它寻址--init-token/--separator-token0 / 2句首加sid 0句子对之间以/sid 2分隔IMDB 是单句任务separator 实际不参与--max-positions512序列上限配合--shorten-method truncate对超长样本做截断优化与数值--criterion sentence_prediction与 task 配套的分类损失--dropout 0.1 --attention-dropout 0.1、--weight-decay 0.1Adam 参数为(0.9, 0.98)、eps1e-06沿用 RoBERTa 官方配置风格--clip-norm 0.0不做梯度裁剪--fp16系列混合精度训练及损失缩放窗口参数--fp16-init-scale 4 --threshold-loss-scale 1 --fp16-scale-window 128--find-unused-parameters分类头之外的部分参数在某些 step 可能不参与梯度DistributedDataParallel 需要该标志避免报错--best-checkpoint-metric accuracy --maximize-best-checkpoint-metric以验证集 accuracy 为指标保留最优 checkpoint。显存不足时的调整方法文档给出的建议是调小--batch-size、同时调大--update-freq以保持有效 batch size 不变例如--batch-size 4 --update-freq 8。预期结果文档报告RoBERTa-large 在有效 batch 32 下训练 10 个 epoch 后best-validation-accuracy约96.5%IMDB dev 集。这是文档原文明确给出的参考值可作为你环境是否配置正确的对照基线。训练时分类头是如何被创建并训练的从源码链路看fairseq-train启动时sentence_prediction.py 的 build_model 在构建完 RoBERTa encoder 后调用model.register_classification_head(head_name, num_classes...)注册一个随机初始化的RobertaClassificationHead该头以nn.ModuleDict形式挂载在模型上model.py L51注册逻辑见 register_classification_head。训练时每个 batch 的目标是 OffsetTokensDataset 提供的标签 idlabel - nspecial即0/1映射回0/1损失由sentence_predictioncriterion 在分类头输出上计算——encoder 主体与分类头一起被端到端微调。七、阶段 6用 hub 接口加载 checkpoint 并推理训练结束后checkpoint 位于checkpoints/目录。此时可以用 RoBERTa 的 hub 接口把“task 模型 分词器”打包加载接口实现见 hub_interface.pyfrom fairseq.models.roberta import RobertaModel roberta RobertaModel.from_pretrained( checkpoints, checkpoint_filecheckpoint_best.pt, data_name_or_pathIMDB-bin ) roberta.eval() # disable dropoutcheckpoint_filecheckpoint_best.pt即第五步按 accuracy 指标挑选出的最优 checkpointdata_name_or_pathIMDB-bin指向预处理好的数据目录接口用它重建sentence_predictiontask从而拿到source_dictionary、label_dictionary和 BPE 编码器。然后用--classification-head-name设置的名字imdb_head做预测并把预测 id 通过 label 词典翻译回字符串hub_interface.py 的 predict 会对特征提取结果过分类头并返回log_softmax值argmax即预测类别label_fn lambda label: roberta.task.label_dictionary.string( [label roberta.task.label_dictionary.nspecial] ) tokens roberta.encode(Best movie this year) pred label_fn(roberta.predict(imdb_head, tokens).argmax().item()) assert pred 1 # positive tokens roberta.encode(Worst movie ever) pred label_fn(roberta.predict(imdb_head, tokens).argmax().item()) assert pred 0 # negative两处容易踩坑的细节roberta.encode(...)会自动补s ... /s并在 BPE 前加空格处理encode 实现与预处理阶段用dict.txt编码的 token 空间一致无需手动分词label_fn中的 nspecial是逆操作训练时目标做了id - nspecial见 OffsetTokensDataset推理翻译回字符串时就要把特殊符号的偏移加回来label_dictionary.string([label nspecial])才能得到0/1这样的字符串。八、复用到你自己的分类数据集改造清单把这条流水线迁移到其他分类任务多分类、短文本/长文本均可时需要改动的点非常集中阶段 2 脚本按你的数据源改读取逻辑保证输出train.input0/train.label/dev.input0/dev.label且行对齐标签可以是任意整数集合0..K-1。阶段 4 第二次预处理label词典仍自动生成无需干预但如果你的任务需要句子对如 NLI则需准备input0input1两份文本任务会按--separator-token拼接ConcatSentencesDataset 逻辑。阶段 5 命令只改NUM_CLASSES、TOTAL_NUM_UPDATES按你的训练集大小与有效 batch 重算、HEAD_NAME和ROBERTA_PATH其余超参可作为起点。阶段 6 推理predict的第一个参数换成你的HEAD_NAME。九、涉及的关键文件索引内容路径本文对应的原始示例文档kosmos-2/fairseq/examples/roberta/README.custom_classification.md多进程 BPE 编码工具kosmos-2/fairseq/examples/roberta/multiprocessing_bpe_encoder.pysentence_prediction 任务实现kosmos-2/fairseq/fairseq/tasks/sentence_prediction.pyRoBERTa 模型与分类头注册kosmos-2/fairseq/fairseq/models/roberta/model.pyhub 加载/编码/预测接口kosmos-2/fairseq/fairseq/models/roberta/hub_interface.pyRoBERTa 主 README模型下载与用法总览kosmos-2/fairseq/examples/roberta/README.md适用前提与限制以上流程基于仓库中 kosmos-2 目录下捆绑的这一版 fairseqhydra 化之前的fairseq-train命令行风格要求能拿到roberta.large的model.pt与 GPT-2 BPE 词表文件--restore-file指向的路径需替换为你本地的实际 checkpoint 位置。IMDB 的 96.5% 验证准确率是文档给出的参考基线复现时应以同配置对照。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考