Transformers 自定义 Pipeline 开发指南：继承 Pipeline 基类、注册新任务并发布到 Hub-编程实验室

Transformers 自定义 Pipeline 开发指南：继承 Pipeline 基类、注册新任务并发布到 Hub

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

本文基于 Transformers 仓库官方文档 add_new_pipeline.md，完整讲解如何为 Transformers 添加一条全新的 pipeline：从输入/输出设计决策，到继承Pipeline基类并实现四个核心方法，再到通过PIPELINE_REGISTRY注册任务，最后给出发布到 Hub（push_to_hub+trust_remote_code）与直接合入 Transformers（补充测试用例）两条落地路径。读完本文，你可以独立完成一条自定义 pipeline 从实现、注册到社区共享的全过程。

一、设计决策：从输入和输出开始

官方文档给出的第一原则是：设计 pipeline 时，最少的工作量只需要为Pipeline提供一个适配任务的输入——这也是设计工作的起点。

具体而言，需要决策两件事：

确定Pipeline能接受哪些输入类型：可以是字符串、原始字节、字典等。文档建议尽可能保持输入为纯 Python 类型，因为兼容性更好；
确定Pipeline应返回的输出：同样保持为纯 Python 类型最简单，也最容易使用。

保持输入输出简单、理想情况下是 JSON 可序列化的，可以让用户无需学习新的对象类型即可运行你的Pipeline。此外，支持多种输入类型也是常见做法，以换取更好的易用性。例如，允许音频文件以文件名、URL 或原始字节三种方式提供，能给用户更大的灵活性。

二、创建 Pipeline：继承基类并实现 4 个方法

输入与输出确定后，就可以开始实现Pipeline了。你的 pipeline 应继承基类Pipeline，并包含 4 个方法：

from transformers import Pipeline class MyPipeline(Pipeline): def _sanitize_parameters(self, **kwargs): def preprocess(self, inputs, args=2): def _forward(self, model_inputs): def postprocess(self, model_outputs):

这 4 个方法在基类 base.py 中全部声明为@abstractmethod，因此子类必须实现，否则实例化会直接失败。下面逐一说明。

1. preprocess：把输入转成模型可用的张量格式

preprocess接收用户输入，将其转换为模型的合适输入格式：

def preprocess(self, inputs, maybe_arg=2): model_input = Tensor(inputs["input_ids"]) return {"model_input": model_input}

从基类 docstring（base.py）可以看到更精确的契约：preprocess接收特定 pipeline 的input_，返回一个包含_forward正常运行所需全部内容的字典——至少应包含一个张量，但可以附带任意其他字段。

2. _forward：模型相关逻辑的唯一归属地

文档强调：_forward不应被直接调用，forward才是推荐入口，因为它包含保障一切在预期设备上正确运行的防护逻辑。凡是与模型相关的事情都放在_forward，其他事情则分别归属preprocess或postprocess。

def _forward(self, model_inputs): outputs = self.model(**model_inputs) return outputs

这一设计在源码中有直接印证：基类的forward方法（base.py）在调用_forward之前会进入self.device_placement()上下文、通过get_inference_context()（PyTorch 下即torch.no_grad）关闭梯度计算，并用_ensure_tensor_on_device把输入张量搬到self.device、把输出搬回 CPU。也就是说，设备管理、推理模式等"热路径外围"逻辑由基类统一兜底，你的_forward只需专注"把输入喂给模型"这一件事，这也是preprocess/postprocess分离出来的根本原因——让热路径尽可能快。

3. postprocess：把模型输出整理成用户友好的结果

postprocess从_forward的模型输出中生成最终输出：

def postprocess(self, model_outputs, top_k=5): best_class = model_outputs["logits"].softmax(-1) return best_class

基类 docstring（base.py）指出：postprocess接收_forward的原始输出（通常是张量），并将其重新组织为更友好的形式，通常输出仅包含字符串和数字的列表或字典——这与第一节"输入输出保持简单、JSON 可序列化"的设计原则首尾呼应。

4. _sanitize_parameters：让用户能传入额外参数

_sanitize_parameters允许用户在初始化或调用Pipeline时传入额外参数。它返回3 个字典，分别作为关键字参数直接传给preprocess、_forward和postprocess。文档特别提醒：如果用户没有以额外参数调用 pipeline，就不要往字典里添加任何东西——这样函数签名中的默认值得以保留，行为更自然。基类对这一点的说明是：它不应被直接调用，而是由__init__和__call__自动调用并完成最终参数解析（base.py）。

例如，在postprocess中加一个top_k参数以返回最可能的 5 个类别，然后在_sanitize_parameters中检查用户是否传入了top_k，并加入postprocess_kwargs：

def _sanitize_parameters(self, **kwargs): preprocess_kwargs = {} if "maybe_arg" in kwargs: preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"] postprocess_kwargs = {} if "top_k" in kwargs: postprocess_kwargs["top_k"] = kwargs["top_k"] return preprocess_kwargs, {}, postprocess_kwargs

于是 pipeline 就支持了"按用户意愿返回 top-k 最可能标签"：

from transformers import pipeline pipeline = pipeline("my-task") # 返回 3 个最可能标签 pipeline("This is the best meal I've ever had", top_k=3) # 默认返回 5 个最可能标签 pipeline("This is the best meal I've ever had")

从源码结构看，这条参数链路是完整闭合的：Pipeline.__call__先调用self._sanitize_parameters(**kwargs)，再把返回的三组参数与初始化时保存的_preprocess_params/_forward_params/_postprocess_params合并（base.py）；单条输入场景下最终由run_single按preprocess→forward→postprocess的顺序串联执行（base.py）；批量输入则走get_iterator构建DataLoader数据流。一个真实实现可参考文本分类 pipeline 的_sanitize_parameters（text_classification.py），它接收function_to_apply、top_k并透传 tokenizer 相关 kwargs。

三、注册 pipeline：写入 PIPELINE_REGISTRY

pipeline 写好后，需要把你支持的新任务注册进PIPELINE_REGISTRY。注册项定义三样东西：

用pt_model指定支持的 PyTorch 模型类；
用default指定默认模型——必须来自一个模型表现符合预期的特定 revision（分支或 commit hash）；
用type指定期望的输入类型。

from transformers.pipelines import PIPELINE_REGISTRY from transformers import AutoModelForSequenceClassification PIPELINE_REGISTRY.register_pipeline( "new-task", pipeline_class=MyPipeline, pt_model=AutoModelForSequenceClassification, default={"pt": ("user/awesome-model", "branch-name")}, type="text", )

源码层面（base.py），PipelineRegistry.register_pipeline的签名为task, pipeline_class, pt_model=None, default=None, type=None，其内部行为可以补充文档未展开的细节：

若任务名已存在，会打印 "already registered. Overwriting..." 警告并覆盖旧注册项；
pt_model为None时置为空元组，非元组则自动包成单元素元组；
default若不是{"model": ...}形式，会被自动包一层{"model": default}，最终结构统一为{"model": (model_id, revision)}；
注册结果task_impl = {"impl": pipeline_class, "pt": pt_model, ...}会写入self.supported_tasks[task]，同时挂到pipeline_class._registered_impl上——后者正是后面push_to_hub能把注册信息写进config.json的关键（见下文第五节）。

对照仓库内置任务表的实现（pipelines/init.py），可以看到官方任务的注册结构与上述机制完全一致。例如text-classification条目：

"text-classification": { "impl": TextClassificationPipeline, "pt": (AutoModelForSequenceClassification,) if is_torch_available() else (), "default": {"model": ("distilbert/distilbert-base-uncased-finetuned-sst-2-english", "714eb0f")}, "type": "text", },

注意default中的第二个元素就是一个 commit 短 hash（714eb0f），印证了文档"默认模型应来自特定 revision"的要求。所有内置任务集中在SUPPORTED_TASKS字典中，并由PIPELINE_REGISTRY = PipelineRegistry(supported_tasks=SUPPORTED_TASKS, task_aliases=TASK_ALIASES)实例化（pipelines/init.py）。任务别名（如"sentiment-analysis" -> "text-classification"）在TASK_ALIASES中定义；PIPELINE_REGISTRY.check_task会先做别名归一化，再查任务表，未知任务抛出KeyError（base.py）。

四、分享你的 pipeline

官方给出两条分享路径：上传到 Hub（快，不需要 Transformers 团队评审），或直接把 pipeline 加进 Transformers（慢，需要评审并补充测试）。

路径一：上传到 Hub

把 pipeline 代码放入 Hub 上的一个 Python 文件。文档以"句子对分类（sentence pair classification）"为例，给出完整自定义 pipeline：

import numpy as np from transformers import Pipeline def softmax(outputs): maxes = np.max(outputs, axis=-1, keepdims=True) shifted_exp = np.exp(outputs - maxes) return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True) class PairClassificationPipeline(Pipeline): def _sanitize_parameters(self, **kwargs): preprocess_kwargs = {} if "second_text" in kwargs: preprocess_kwargs["second_text"] = kwargs["second_text"] return preprocess_kwargs, {}, {} def preprocess(self, text, second_text=None): return self.tokenizer(text, text_pair=second_text, return_tensors=self.framework) def _forward(self, model_inputs): return self.model(**model_inputs) def postprocess(self, model_outputs): logits = model_outputs.logits[0].numpy() probabilities = softmax(logits) best_class = np.argmax(probabilities) label = self.model.config.id2label[best_class] score = probabilities[best_class].item() logits = logits.tolist() return {"label": label, "score": score, "logits": logits}

这段代码完整体现了第二节的四个方法：preprocess借助self.tokenizer处理text_pair，postprocess输出纯 JSON 友好的字典。

把代码保存为pair_classification.py，然后导入并注册：

from pair_classification import PairClassificationPipeline from transformers.pipelines import PIPELINE_REGISTRY from transformers import AutoModelForSequenceClassification PIPELINE_REGISTRY.register_pipeline( "pair-classification", pipeline_class=PairClassificationPipeline, pt_model=AutoModelForSequenceClassification, )

文档称register_pipeline函数会把 pipeline 的详细信息（任务类型、pipeline 类、支持的后端）注册进模型的config.json。具体机制在源码中清晰可见：Pipeline.push_to_hub时若实例带有_registered_impl属性，就会把实现类转换成"模块名.类名"字符串、模型类转换为其类名，最终写为self.model.config.custom_pipelines（base.py）。对应的config.json片段形如：

"custom_pipelines": { "pair-classification": { "impl": "pair_classification.PairClassificationPipeline", "pt": [ "AutoModelForSequenceClassification" ], } },

接着调用Pipeline.push_to_hub推送 pipeline 到 Hub：包含代码的 Python 文件会被复制到 Hub，pipeline 的模型与 tokenizer 也会被保存并推送。此时你的 pipeline 就已可用在 Hub 上的命名空间下了：

from transformers import pipeline pipeline = pipeline(task="pair-classification", model="sgugger/finetuned-bert-mrpc") pipeline.push_to_hub("pair-classification-pipeline")

使用时，加载 pipeline 需加上trust_remote_code=True：

from transformers import pipeline pipeline = pipeline(task="pair-classification", trust_remote_code=True)

从源码结构看，这条远程加载链路的落点是工厂函数pipeline()：它会检查config.custom_pipelines（pipelines/init.py），经clean_custom_task把pt中的类名字符串还原为真正的模型类后，通过get_class_from_dynamic_module从仓库动态加载impl指向的类；若未设置trust_remote_code则会抛出明确提示"需先阅读代码以规避恶意使用"的ValueError（pipelines/init.py）。这解释了为何消费端必须显式传trust_remote_code=True。

路径二：直接加入 Transformers

把自定义 pipeline 加入 Transformers 本体，要求补充测试以确保一切按预期工作，并请求 Transformers 团队评审。流程如下：

添加代码模块：把 pipeline 代码作为新模块加入 pipelines 子包，并把它加入 pipelines/init.py 中定义的任务列表——在当前的仓库结构中，即导入你的 pipeline 类并新增一条SUPPORTED_TASKS条目（参考上文text-classification条目的impl/pt/default/type四字段写法）；
添加测试：在 tests/pipelines 下为新 pipeline 添加测试文件，可参考现有测试文件了解如何测试你的 pipeline。

测试写法上有三个文档明确强调、且在仓库测试中可对照实现的要点：

run_pipeline_test必须足够通用，并且要能运行在model_mapping定义的模型上（参考 test_pipelines_text_classification.py 中model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING的定义，以及 test_pipelines_text_classification.py 的run_pipeline_test实现）。这一点非常重要，因为它保证了未来与新模型的兼容性；
大量使用ANY占位符。由于测试模型权重是随机的，无法断言具体数值；使用ANY可以让测试匹配 pipeline 输出的结构类型而非具体值（ANY定义于 test_pipelines_common.py，例如self.assertEqual(nested_simplify(outputs), [{"label": ANY(str), "score": ANY(float)}])，见 test_pipelines_text_classification.py）；
额外实现 4 个测试（文档原文如此列出，其中两项为）：
1. test_small_model_pt：使用小模型（如hf-internal-testing/tiny-random-distilbert）验证 pipeline 返回正确的输出结构。结果不必有语义意义，但每个 pipeline 应返回同一结果（该测试同时验证了top_k等额外参数路径）；
2. test_large_model_pt：使用真实模型验证 pipeline 返回有意义的结果。这类测试运行缓慢，应标记为 slow（对应@slow装饰器）。

五、小结：自定义 pipeline 的完整生命周期

把本文流程串起来，一条自定义 pipeline 的落地路径是：

设计：确定输入/输出为尽量简单、JSON 可序列化的纯 Python 类型；
实现：继承Pipeline，实现_sanitize_parameters（参数路由，返回三组 kwargs 字典）、preprocess（输入→张量字典）、_forward（只放模型相关逻辑，由基类forward统一处理设备与推理上下文）、postprocess（模型输出→友好结果）；
注册：通过PIPELINE_REGISTRY.register_pipeline声明pipeline_class、pt_model、带 revision 的default与type；
发布：轻量路线用push_to_hub把代码与custom_pipelines配置一并推到 Hub，消费端以trust_remote_code=True加载；正式路线则进入 src/transformers/pipelines 与 tests/pipelines 参与仓库评审与 CI 测试。