unilm 仓库中的 fairseq 教程用字符级 RNN 实现姓名分类理解 fairseq 的 Model 与 Task 注册机制【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm本篇指南基于 unilm 仓库infoxlm/fairseq目录内置的 fairseq 文档完整讲解如何在 fairseq 框架上扩展一个分类任务以用字符级 RNN 判断人名所属民族Classifying Names with a Character-Level RNN为案例覆盖数据预处理fairseq-preprocess、注册新 Modelregister_model、注册新 Taskregister_task、使用fairseq-train训练以及编写交互式评估脚本eval_classifier.py五个步骤。读完本文你将掌握 fairseq 插件化扩展的完整链路——如何把一个普通的 PyTorch 模型封装进 fairseq 的训练/检查点/命令行体系并复用其批处理与推理基础设施。1. 案例背景把分类任务塞进 seq2seq 框架fairseq 的核心抽象围绕序列对source-target设计而人名 - 民族是一个典型的分类问题。本教程的解决思路是把标签当作长度为 1 的目标序列从而复用 fairseq 的预处理工具、LanguagePairDataset批处理机制和fairseq-train命令行工具无需重写训练循环。教程使用的数据是已按字符分词、并切分为 train/valid/test 三集的修改版数据集tutorial_names.tar.gz需从 fairseq 官方数据源下载解压仓库内不包含该数据文件。解压后目录结构为names/train.input、names/train.label等文件对。2. 数据预处理复用 fairseq-preprocess 构建词表虽然fairseq-preprocess主要面向序列到序列任务但这里通过把 label 当作 target 序列来复用它并用--dataset-impl raw输出raw格式文件以增强可读性 fairseq-preprocess \ --trainpref names/train --validpref names/valid --testpref names/test \ --source-lang input --target-lang label \ --destdir names-bin --dataset-impl raw命令执行后会在names-bin/目录下生成inputs与labels两个词表dictionary。这一步的关键点--source-lang input/--target-lang label决定了词表文件名dict.input.txt、dict.label.txt以及分片文件后缀后文 Task 中读取的train.input/train.label正来源于此该数据集中输入是逐字符空格分隔的序列target 只有一个符号因此后续 Task 只需声明 target 最大长度为 1。3. 注册新 ModelRNN 编码器 fairseq 封装类下一步在fairseq/models/rnn_classifier.py中注册一个用简单 RNN 编码输入句、预测输出标签的模型。与 PyTorch 原版教程相比fairseq 版本额外支持批量数据和 GPU Tensor。3.1 拷贝 PyTorch 教程中的 RNN 模块文件 rnn_classifier.py教程新建首先定义一个标准 nn.Moduleimport torch import torch.nn as nn class RNN(nn.Module): def __init__(self, input_size, hidden_size, output_size): super(RNN, self).__init__() self.hidden_size hidden_size self.i2h nn.Linear(input_size hidden_size, hidden_size) self.i2o nn.Linear(input_size hidden_size, output_size) self.softmax nn.LogSoftmax(dim1) def forward(self, input, hidden): combined torch.cat((input, hidden), 1) hidden self.i2h(combined) output self.i2o(combined) output self.softmax(output) return output, hidden def initHidden(self): return torch.zeros(1, self.hidden_size)每个字符步网络把当前 one-hot 输入与上一时刻隐状态拼接后同时产生新隐状态和带 LogSoftmax 的输出分布。3.2 用 register_model 装饰器接入 fairseq 注册表所有注册的模型都必须实现 BaseFairseqModel 接口。在同一个文件中教程定义了FairseqRNNClassifier并以名字rnn_classifier注册from fairseq.models import BaseFairseqModel, register_model # 注意register_model 装饰器必须紧贴 Model 类定义 register_model(rnn_classifier) class FairseqRNNClassifier(BaseFairseqModel): staticmethod def add_args(parser): # 模型可以覆写此方法来添加新的命令行参数。 # 这里添加一个 --hidden-dim 参数配置隐状态维度。 parser.add_argument( --hidden-dim, typeint, metavarN, helpdimensionality of the hidden state, ) classmethod def build_model(cls, args, task): # fairseq 通过调用 build_model() 来初始化模型。 # 这提供了更大灵活性返回的模型实例 # 可以不同于调用它的类。这里直接返回 # FairseqRNNClassifier 实例。 # 初始化 RNN 模块 rnn RNN( # Task 在下一节定义注意 task 持有 source # 输入句和 target标签的词表 input_sizelen(task.source_dictionary), hidden_sizeargs.hidden_dim, output_sizelen(task.target_dictionary), ) # 返回模块的包装版本 return FairseqRNNClassifier( rnnrnn, input_vocabtask.source_dictionary, ) def __init__(self, rnn, input_vocab): super(FairseqRNNClassifier, self).__init__() self.rnn rnn self.input_vocab input_vocab # 教程中的 RNN 模块期望 one-hot 输入因此预先计算 # 单位矩阵用于把 index 转换为 one-hot 向量。 # 注册为 buffer 以便在调用 cuda() 时随模型移动到 GPU。 self.register_buffer(one_hot_inputs, torch.eye(len(input_vocab))) def forward(self, src_tokens, src_lengths): # forward() 的输入由 Task 决定具体是 mini-batch 中 # net_input 键的内容。*src_tokens* 形状为 # (batch, src_len)*src_lengths* 形状为 (batch)。 bsz, max_src_len src_tokens.size() # 初始化 RNN 隐状态。相比原教程这里同时处理 # 批量输入并工作在 GPU 上。 hidden self.rnn.initHidden() hidden hidden.repeat(bsz, 1) # 扩展到 batch 维度 hidden hidden.to(src_tokens.device) # 移到 GPU for i in range(max_src_len): # 警告输入带有 padding此处应当对 padding # 做掩码避免影响结果。这留作读者练习。 # padding 符号由 self.input_vocab.pad() 给出 # 每条输入未填充的长度由 *src_lengths* 给出。 # 把一批输入字符 one-hot 编码。 input self.one_hot_inputs[src_tokens[:, i].long()] # 送入 RNN。 output, hidden self.rnn(input, hidden) # 返回最终输出状态用于预测 return output这段实现里有三个值得注意的 fairseq 习惯build_model(args, task)工厂模式fairseq 不是直接__init__模型而是通过 classmethodbuild_model构造因此返回实例的类型可以与被调用的类不同。词表大小从task.source_dictionary/task.target_dictionary获取即Task 是模型与词表之间的桥梁register_buffer(one_hot_inputs, ...)把单位矩阵注册为 buffer 而非普通属性保证调用cuda()时自动迁移设备且不参与优化器参数forward(src_tokens, src_lengths)的入参契约由 Task 的net_input键决定——后文的 Task 实现正是构造了这个键。3.3 从源码看注册机制MODEL_REGISTRY 与自动导入教程使用的register_model与register_model_architecture装饰器定义在 models/init.py其底层是三个模块级注册表MODEL_REGISTRY、ARCH_MODEL_REGISTRY、ARCH_CONFIG_REGISTRY。结合源码可以看到几个关键事实register_model会把类写入MODEL_REGISTRY并强制校验issubclass(cls, BaseFairseqModel)见 models/init.pyregister_model_architecture(model_name, arch_name)要求 model 已注册并把 arch 名映射到 model 类、把配置函数存入ARCH_CONFIG_REGISTRY训练入口通过build_model(args, task)即ARCH_MODEL_REGISTRY[args.arch].build_model(args, task)完成架构名 - 模型类的分发见 models/init.py文件名即模块名fairseq/models/__init__.py末尾会遍历models/目录下所有.py文件并importlib.import_module(fairseq.models. model_name)自动导入见 models/init.py。这就是为什么新模型文件必须放在fairseq/models/下、且装饰器注册后无需手动 import 即可被命令行识别同时它还会为每个已注册 model 生成一个带--archchoices 与模型专属参数的 argparse 组。最后定义一个命名架构之后即可通过--arch pytorch_tutorial_rnn引用from fairseq.models import register_model_architecture # register_model_architecture() 的第一个参数是上面注册的 # 模型名即 rnn_classifier。被装饰的函数接收单个 # 参数 *args* 并就地修改它以匹配目标架构。 register_model_architecture(rnn_classifier, pytorch_tutorial_rnn) def pytorch_tutorial_rnn(args): # 使用 getattr() 是为了优先采用命令行显式给出的参数 # 下面定义的默认值仅在未指定其他值时生效。 args.hidden_dim getattr(args, hidden_dim, 128)这一架构配置函数的设计使得同一 model 可以有多个预设例如不同隐藏维度、不同规模的架构且命令行显式参数始终覆盖默认值——getattr(args, hidden_dim, 128)正是这一优先级约定的实现方式。4. 注册新 Task加载词表与数据集Task 负责加载词表和构造数据集也可以控制批处理策略本教程复用 LanguagePairDataset 提供的批处理把分类建模为目标长度为 1 的 seq2seq。教程新建 simple_classification.pyimport os import torch from fairseq.data import Dictionary, LanguagePairDataset from fairseq.tasks import FairseqTask, register_task register_task(simple_classification) class SimpleClassificationTask(FairseqTask): staticmethod def add_args(parser): # 添加命令行参数数据位置与最大输入长度 parser.add_argument(data, metavarFILE, helpfile prefix for data) parser.add_argument(--max-positions, default1024, typeint, helpmax input length) classmethod def setup_task(cls, args, **kwargs): # 在这里可以做任务所需的任意设置包括加载词表、 # 初始化共享 Embedding 等。本例只加载两个词表。 input_vocab Dictionary.load(os.path.join(args.data, dict.input.txt)) label_vocab Dictionary.load(os.path.join(args.data, dict.label.txt)) print(| [input] dictionary: {} types.format(len(input_vocab))) print(| [label] dictionary: {} types.format(len(label_vocab))) return SimpleClassificationTask(args, input_vocab, label_vocab) def __init__(self, args, input_vocab, label_vocab): super().__init__(args) self.input_vocab input_vocab self.label_vocab label_vocab def load_dataset(self, split, **kwargs): Load a given dataset split (e.g., train, valid, test). prefix os.path.join(self.args.data, {}.input-label.format(split)) # 读取输入句。 sentences, lengths [], [] with open(prefix .input, encodingutf-8) as file: for line in file: sentence line.strip() # 按空格分词 tokens self.input_vocab.encode_line( sentence, add_if_not_existFalse, ) sentences.append(tokens) lengths.append(tokens.numel()) # 读取标签。 labels [] with open(prefix .label, encodingutf-8) as file: for line in file: label line.strip() labels.append( # 把标签转换为数值 ID。 torch.LongTensor([self.label_vocab.add_symbol(label)]) ) assert len(sentences) len(labels) print(| {} {} {} examples.format(self.args.data, split, len(sentences))) # 由于分类可以建模为目标序列长度为 1 的 seq2seq 任务 # 这里复用 LanguagePairDataset。 self.datasets[split] LanguagePairDataset( srcsentences, src_sizeslengths, src_dictself.input_vocab, tgtlabels, tgt_sizestorch.ones(len(labels)), # target 长度为 1 tgt_dictself.label_vocab, left_pad_sourceFalse, max_source_positionsself.args.max_positions, max_target_positions1, # 目标只有一个类别标签无需 teacher forcing。 # 若设为 TrueModel 的 forward() 会额外收到 # 一个名为 prev_output_tokens 的参数目标序列的 # 错位版本。 input_feedingFalse, ) def max_positions(self): Return the max input length allowed by the task. # source 应小于 args.max_positionstarget 最大长度为 1。 return (self.args.max_positions, 1) property def source_dictionary(self): Return the source :class:~fairseq.data.Dictionary. return self.input_vocab property def target_dictionary(self): Return the target :class:~fairseq.data.Dictionary. return self.label_vocab # 如果需要更精细地控制 batch 构造可以覆写此方法 # 但本教程无需这么做因为复用了 LanguagePairDataset 的批处理。 # # def get_batch_iterator( # self, dataset, max_tokensNone, max_sentencesNone, max_positionsNone, # ignore_invalid_inputsFalse, required_batch_size_multiple1, # seed1, num_shards1, shard_id0, # ): # (...)实现要点setup_task是类方法训练入口通过tasks.setup_task(args)按args.task从注册表分发到对应 Task 类见 tasks/init.py词表加载路径与预处理参数一一对应dict.input.txt/dict.label.txt来自第 2 节的--source-lang input --target-lang labelinput_feedingFalse决定了模型forward()只收到src_tokens与src_lengths即net_input键不会被塞入prev_output_tokens——这正是第 3 节模型签名与 Task 之间的契约max_positions返回(args.max_positions, 1)与 LanguagePairDataset 的max_source_positions/max_target_positions保持一致。从 tasks/init.py 的源码可以看到fairseq/tasks/目录下的.py文件同样被自动导入且每个成功注册的 Task 会生成一个专属 argparse 组task_name _parser把--task name和 Task 自身的add_args参数挂到命令行体系上——所以新 Task 只需放进fairseq/tasks/并正确register_task即可被fairseq-train感知无需修改任何入口脚本。5. 训练模型fairseq-train 命令行模型与任务注册完成后直接用现有fairseq-train工具训练只需指定新 Task--task simple_classification和新架构--arch pytorch_tutorial_rnn fairseq-train names-bin \ --task simple_classification \ --arch pytorch_tutorial_rnn \ --optimizer adam --lr 0.001 --lr-shrink 0.5 \ --max-tokens 1000 (...) | epoch 027 | loss 1.200 | ppl 2.30 | wps 15728 | ups 119.4 | wpb 116 | bsz 116 | num_updates 3726 | lr 1.5625e-05 | gnorm 1.290 | clip 0% | oom 0 | wall 32 | train_wall 21 | epoch 027 | valid on valid subset | valid_loss 1.41304 | valid_ppl 2.66 | num_updates 3726 | best 1.41208 | done training in 31.6 seconds训练输出的检查点文件会出现在checkpoints/目录中。补充两点通过--hidden-dim N可以覆盖命名架构中的默认隐藏维度默认 128因为pytorch_tutorial_rnn架构函数使用getattr保留了命令行优先的语义--max-tokens 1000控制按 token 数组批--lr-shrink 0.5触发 fairseq 内置的学习率衰减逻辑这些均属于 fairseq 通用训练参数由 options.py 中get_training_parser组装的数据集/优化/检查点参数组提供。6. 编写交互式评估脚本最后写一个短脚本 eval_classifier.py教程新建于工作目录非仓库文件直接导入 fairseq 对新输入做预测from fairseq import checkpoint_utils, data, options, tasks # 解析生成所需的命令行参数 parser options.get_generation_parser(default_tasksimple_classification) args options.parse_args_and_arch(parser) # Setup task task tasks.setup_task(args) # Load model print(| loading model from {}.format(args.path)) models, _model_args checkpoint_utils.load_model_ensemble([args.path], tasktask) model models[0] while True: sentence input(\nInput: ) # Tokenize into characters chars .join(list(sentence.strip())) tokens task.source_dictionary.encode_line( chars, add_if_not_existFalse, ) # Build mini-batch to feed to the model batch data.language_pair_dataset.collate( samples[{id: -1, source: tokens}], # bsz 1 pad_idxtask.source_dictionary.pad(), eos_idxtask.source_dictionary.eos(), left_pad_sourceFalse, input_feedingFalse, ) # Feed batch to the model and get predictions preds model(**batch[net_input]) # Print top 3 predictions and their log-probabilities top_scores, top_labels preds[0].topk(k3) for score, label_idx in zip(top_scores, top_labels): label_name task.target_dictionary.string([label_idx]) print(({:.2f})\t{}.format(score, label_name))脚本的关键环节与源码对应关系options.get_generation_parser(default_task...)与options.parse_args_and_arch(parser)正是 options.py 中定义的解析入口parse_args_and_arch内部做两次解析——先解析出--arch再把模型专属参数组如本例的--hidden-dim挂进 parser 后二次解析见 options.pycheckpoint_utils.load_model_ensemble([args.path], tasktask)从检查点恢复模型权重并构造模型实例见 checkpoint_utils.pydata.language_pair_dataset.collate(...)手工构造了与训练时 Task 相同的net_input批left_pad_sourceFalse、input_feedingFalse必须与 Task 中LanguagePairDataset的构造参数保持一致否则模型收到的键或填充方向就会错乱预测输出是 LogSoftmax 分布topk(k3)取出前三名及其 log-probability再用task.target_dictionary.string(...)把 ID 还原为标签名。运行方式注意必须传入原始数据目录names-bin/以便加载词表 python eval_classifier.py names-bin --path checkpoints/checkpoint_best.pt | [input] dictionary: 64 types | [label] dictionary: 24 types | loading model from checkpoints/checkpoint_best.pt Input: Satoshi (-0.61) Japanese (-1.20) Arabic (-2.86) Italian Input: Sinbad (-0.30) Arabic (-1.76) English (-4.08) Russian7. 总结本教程沉淀的 fairseq 扩展范式本教程虽然案例很小但完整覆盖了向 fairseq 添加新任务的四要素且每一项都有仓库源码可查证扩展步骤教程做法仓库中的支撑机制数据复用fairseq-preprocesslabel 作为长度 1 的 target词表dict.input.txt/dict.label.txt新 Modelregister_model(rnn_classifier)build_model(args, task)register_model_architecturefairseq/models/__init__.pymodels/目录自动导入新 Taskregister_task(simple_classification)setup_task加载词表、load_dataset复用LanguagePairDatasetfairseq/tasks/__init__.pytasks/目录自动导入训练与推理现有fairseq-train 自定义eval_classifier.pyfairseq/checkpoint_utils.py 的load_model_ensemble需要留意的限制教程代码中的 padding 掩码被明确留作读者练习见forward()中的 WARNING 注释因此在变长输入上严格结果会受 padding 影响评估脚本为交互式while True循环适合人工调试而非批量评测数据文件tutorial_names.tar.gz需另行下载仓库本身只包含教程文档与 fairseq 源码。把这套注册模型 - 注册任务 - 复用训练工具 - 脚本推理的范式替换为自己的 Encoder 与词表加载逻辑即可在 fairseq 之上快速落地任意单标签分类场景。【免费下载链接】unilmLarge-scale Self-supervised Pre-training Across Tasks, Languages, and Modalities项目地址: https://gitcode.com/GitHub_Trending/un/unilm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考