ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

Transformers Pipelines 完全指南:用 `pipeline()` 一行代码完成推理、批量加速与自定义扩展

Transformers Pipelines 完全指南:用 `pipeline()` 一行代码完成推理、批量加速与自定义扩展 Transformers Pipelines 完全指南用pipeline()一行代码完成推理、批量加速与自定义扩展【免费下载链接】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 官方文档 Pipelines 章节日文版与 英文版 同源的深度解读与源码级扩充。Pipelines 是 Transformers 面向推理场景提供的最易用 API它把「加载模型 → 预处理分词/图像处理/特征提取→ 模型前向 → 后处理」这一整条复杂链路抽象成一个可调用对象覆盖命名实体识别NER、掩码语言建模、情感分析、特征提取、问答等任务。读完本文你将掌握pipeline()工厂函数的完整用法单条/批量/流式输入、数据集迭代、自定义模型、批量推理batching与 chunk batching 的底层原理与调优经验、FP16 推理加速以及如何通过子类化定制属于自己的 pipeline。一、Pipelines 的两层抽象结构Pipelines 体系分为两层见 pipelines.mdpipeline工厂函数也是最强大的入口对象。它根据你传入的任务名task自动选择并封装所有其他 pipeline负责加载模型、分词器、图像处理器等组件任务专属 pipeline 类按模态分为 音频Audio、计算机视觉Computer Vision、自然语言处理NLP、多模态Multimodal 四类每个类只专注一个具体任务。从源码看pipeline()的完整签名src/transformers/pipelines/init.py提供了丰富的参数核心包括参数类型作用taskstr决定返回哪种 pipeline如text-classification、ner、automatic-speech-recognitionmodelstr/PreTrainedModel模型标识符或模型实例不传则加载该任务默认模型configstr/PreTrainedConfig配置标识符或配置实例决定模型架构tokenizer/feature_extractor/image_processor/processor各预处理组件均可省略pipeline会自动按「model → config → task 默认值」的优先级推断加载revisionstr默认main指定 Hub 上的分支、tag 或 commit iduse_fastbool默认True是否优先使用 Fast tokenizerdeviceint/str/torch.device分配设备cpu、cuda:1、mps等不传时自动放到首个可用加速器device_mapstr/dict配合accelerate使用device_mapauto自动规划多卡/卸载策略不可与device同时使用dtypestr/torch.dtype默认auto模型加载精度可显式传torch.float16trust_remote_codebool默认False是否信任 Hub 仓库中的自定义代码pipeline_classAny传入自定义 pipeline 类以覆盖默认实现值得注意的约束源码中会直接抛出RuntimeErrortask与model至少指定一个只给tokenizer/feature_extractor而不给model会报错因为无法保证该组件与默认模型兼容。二、pipeline 抽象从单条输入到流式迭代2.1 单条输入与任务名最简单的用法是只传任务名pipeline会自动加载该任务的默认模型 pipe pipeline(text-classification) pipe(This restaurant is awesome) [{label: POSITIVE, score: 0.9998743534088135}]2.2 用指定模型时任务名可以省略如果模型在 Hub 上已经声明了自己的任务可以只传model而忽略task前提是模型仓库的 config 中定义了pipeline_tag/ 任务映射 pipe pipeline(modelFacebookAI/roberta-large-mnli) pipe(This restaurant is awesome) [{label: NEUTRAL, score: 0.7313136458396912}]2.3 列表输入一次处理多个样本 pipe pipeline(text-classification) pipe([This restaurant is awesome, This restaurant is awful]) [{label: POSITIVE, score: 0.9998743534088135}, {label: NEGATIVE, score: 0.9996669292449951}]2.4 直接迭代 Dataset推荐的数据集推理方式对于完整数据集官方推荐直接把datasets.Dataset传给 pipeline既不需要把整个数据集一次性载入内存也无需自己写 batch 循环在 GPU 上的速度与手写循环相当。配合KeyDataset可以只取字典中的某个字段import datasets from transformers import pipeline from transformers.pipelines.pt_utils import KeyDataset from tqdm.auto import tqdm pipe pipeline(automatic-speech-recognition, modelfacebook/wav2vec2-base-960h, device0) dataset datasets.load_dataset(superb, nameasr, splittest) # KeyDataset 只返回数据集条目 dict 中的指定 key这里是音频文件路径 # 句子对场景请使用 KeyPairDataset for out in tqdm(pipe(KeyDataset(dataset, file))): print(out) # {text: NUMBER TEN FRESH NELLY IS WAITING ON YOU GOOD NIGHT HUSBAND}KeyDataset的实现非常轻量src/transformers/pipelines/pt_utils.py它只是把dataset[i][key]包装成一个torch.utils.data.DatasetKeyPairDataset则同时取两个 key 并组装成{text: ..., text_pair: ...}供句子对任务使用pt_utils.py。2.5 生成器输入适合数据流场景pipeline 也接受生成器适合从数据库、消息队列或 HTTP 请求流中边取边推from transformers import pipeline pipe pipeline(text-classification) def data(): while True: # 数据可来自数据集、数据库、队列或服务器上的 HTTP 请求 # Caveat: 因为这是迭代式输入不能使用 num_workers 1 做多线程预处理 # 但依然可以开 1 个线程做预处理主线程跑推理 yield This is a test for out in pipe(data()): print(out)2.6 底层调用链preprocess → forward → postprocesspipeline的调用之所以灵活是因为每个 pipeline 都由三个阶段构成。Pipeline基类src/transformers/pipelines/base.py把整个工作流定义为Input - Tokenization (预处理) - Model Inference (模型前向) - Post-Processing (任务相关后处理) - Output基类中还内置了设备选择逻辑自动识别 CUDA / MPS / XPU / HPU / NPU / MLU / MUSA 等后端base.py同时如果模型已用accelerate加载带有hf_device_map会禁止再传device避免冲突base.py。三、Pipeline batching批量推理的提速与陷阱所有 pipeline 都支持 batching——只要输入走的是「流式」路径传入 list、Dataset或generator即可启用from transformers import pipeline from transformers.pipelines.pt_utils import KeyDataset import datasets dataset datasets.load_dataset(stanfordnlp/imdb, nameplain_text, splitunsupervised) pipe pipeline(text-classification, device0) for out in pipe(KeyDataset(dataset, text), batch_size8, truncationonly_first): print(out) # [{label: POSITIVE, score: 0.9998743534088135}] # 输出与之前完全一致只是内容以 batch 形式送入模型⚠️ 官方文档特别警告batching 并不自动等于性能提升取决于硬件、数据和具体模型可能是 10 倍加速也可能是 5 倍变慢。3.1 提速案例短文本 GPU文档给出 GTX 970 上的实测数据5000 条等长短文本无 batch 与不同 batch_size 对比配置耗时吞吐无 batching26s187.52 it/sbatch_size84s1205.95 it/sbatch_size642s2478.24 it/sbatch_size2561s2554.43 it/s收益递减GPU 已饱和测试脚本核心是自定义torch.utils.data.Dataset并在不同batch_size下循环调用from transformers import pipeline from torch.utils.data import Dataset from tqdm.auto import tqdm pipe pipeline(text-classification, device0) class MyDataset(Dataset): def __len__(self): return 5000 def __getitem__(self, i): return This is a test dataset MyDataset() for batch_size in [1, 8, 64, 256]: print(- * 30) print(fStreaming batch_size{batch_size}) for out in tqdm(pipe(dataset, batch_sizebatch_size), totallen(dataset)): pass3.2 变慢案例长度不齐导致 padding 放大如果数据集中偶发超长样本例如每 64 条里出现 1 条 100 倍长度的句子整批都会被 pad 到最长长度[64, 4]会退化成[64, 400]速度大幅下降batch 更大时甚至直接 OOM 崩溃class MyDataset(Dataset): def __len__(self): return 5000 def __getitem__(self, i): if i % 64 0: n 100 else: n 1 return This is a test * n文档给出的实测对比无 batch 5s / 183.69 it/sbatch_size83s / 265.74 it/sbatch_size6426s / 37.80 it/sbatch_size256直接RuntimeError: CUDA out of memory。3.3 官方经验法则用你的硬件、你的负载去测反复测真实数字是唯一准则延迟敏感线上实时推理→ 不要 batch使用 CPU → 不要 batchGPU 上追求吞吐批量跑静态数据时对「自然」数据序列长度不可控→ 默认不要 batch先测再试探性加上并加 OOM 兜底恢复逻辑序列长度非常规整 → batch 大概率收益很大测到 OOM 为止GPU 越大batching 越可能划算一旦启用 batching务必确保能优雅处理 OOM。四、Pipeline chunk batching零样本分类与问答的特殊处理zero-shot-classification和question-answering很特殊一条输入可能触发模型的多次前向如零样本分类要为每个候选标签分别打分。普通 pipeline 的三段式流程preprocessed pipe.preprocess(inputs) model_outputs pipe.forward(preprocessed) outputs pipe.postprocess(model_outputs)在ChunkPipeline继承自Pipeline见 src/transformers/pipelines/base.py中变成对每个 chunk 循环all_model_outputs [] for preprocessed in pipe.preprocess(inputs): model_outputs pipe.forward(preprocessed) all_model_outputs.append(model_outputs) outputs pipe.postprocess(all_model_outputs)从源码可以看到ZeroShotClassificationPipeline正是ChunkPipeline的子类。这种变化对调用方完全透明——使用方式与普通 pipeline 一致且 pipeline 会自动帮你处理 batch因此你无需关心一条输入实际触发了几次前向可以独立于输入内容去优化batch_size上一节的所有注意事项依然适用。五、Pipeline FP16 推理一行代码开启半精度加速模型可以以 FP16 精度运行在 GPU 上通常显著提速并节省显存且大多数模型不会出现明显精度损失——模型越大越不容易受影响。启用方式非常简单只需给pipeline()构造函数传dtypetorch.float16或dtypefloat16from transformers import pipeline pipe pipeline(text-classification, modelFacebookAI/roberta-large-mnli, dtypetorch.float16)注意两点仅对PyTorch 后端模型有效你的输入会在内部自动转换为 FP16。对应地pipeline()的dtype参数默认是auto按模型保存时的精度加载可以显式覆盖为torch.float16、torch.bfloat16、torch.float32等见 src/transformers/pipelines/init.py。六、Pipeline custom code子类化定制你的 pipeline如果默认 pipeline 不满足需求官方建议先到社区提交 issue——pipeline 的设计目标就是易用并覆盖大多数场景。想快速实验时直接子类化目标 pipeline 并重写方法即可class MyPipeline(TextClassificationPipeline): def postprocess(self, model_outputs, **kwargs): # 你的定制逻辑 scores scores * 100 # 更多定制逻辑 return results my_pipeline MyPipeline(modelmodel, tokenizertokenizer, ...) # 或走 pipeline 工厂传入 pipeline_class my_pipeline pipeline(modelxxxx, pipeline_classMyPipeline)这种方式可以完全接管 preprocess / forward / postprocess 中的任意环节。关于从零实现一个全新 pipeline 的完整流程参见文档 add_new_pipeline日文版对应docs/source/en/add_new_pipeline.md。所有 pipeline 类的基类Pipelinesrc/transformers/pipelines/base.py还支持binary_output构造参数——对feature-extraction这类输出超大张量的任务设为True时输出以 pickle 二进制格式保存避免转成嵌套列表文本。七、音频Audio任务AudioClassificationPipeline音频分类输入音频路径/字节流/数组输出类别标签AutomaticSpeechRecognitionPipeline自动语音识别ASR即语音转文本支持带外部语言模型的 CTC 解码源码中会尝试通过pyctcdecodekenlm自动加载WithLM解码器见 src/transformers/pipelines/init.pyTextToAudioPipeline文本转音频TTS即文生音/语音合成ZeroShotAudioClassificationPipeline零样本音频分类无需微调即可按候选标签分类音频。八、计算机视觉Computer Vision任务DepthEstimationPipeline深度估计输出每个像素的深度图ImageClassificationPipeline图像分类ImageSegmentationPipeline图像分割语义/实例/全景分割ObjectDetectionPipeline目标检测输出边界框、类别与置信度VideoClassificationPipeline视频分类ZeroShotImageClassificationPipeline零样本图像分类ZeroShotObjectDetectionPipeline零样本目标检测。英文版文档中还包含KeypointMatchingPipeline关键点匹配、MaskGenerationPipeline掩码生成等更新任务日文版正文未列出但源码中均已实现见 src/transformers/pipelines 目录下的keypoint_matching.py、mask_generation.py。九、自然语言处理NLP任务FillMaskPipeline掩码语言建模填充句子中的[MASK]位置NerPipeline命名实体识别ner任务别名是TokenClassificationPipeline的别名/特化源码中二者同源见 src/transformers/pipelines/token_classification.py详细 API 参考TokenClassificationPipelineTableQuestionAnsweringPipeline表格问答对结构化表格数据回答自然语言问题TextClassificationPipeline文本分类text-classification别名sentiment-analysis即情感分析等TextGenerationPipeline文本生成自回归生成底层调用generateTokenClassificationPipelinetoken 级分类词性标注、NER 等ZeroShotClassificationPipeline零样本文本分类属于ChunkPipeline。十、多模态Multimodal任务DocumentQuestionAnsweringPipeline文档问答基于文档图像 问题输出答案FeatureExtractionPipeline特征提取feature-extraction把文本转为稠密向量常用于 embeddingImageFeatureExtractionPipeline图像特征提取ImageTextToTextPipeline图像 文本到文本视觉语言模型、图像字幕等。多模态 pipeline 的预处理依赖processorProcessorMixin同时处理文本与图像pipeline()加载 processor 时会校验其类型不是ProcessorMixin实例会直接抛TypeError见 src/transformers/pipelines/init.py。十一、常用任务别名速查从pipeline()的 docstringsrc/transformers/pipelines/init.py可以整理出以下常用任务标识符及其别名任务标识符返回的 pipeline 类可用别名text-classificationTextClassificationPipelinesentiment-analysistoken-classificationTokenClassificationPipelinenertext-to-audioTextToAudioPipelinetext-to-speechautomatic-speech-recognitionAutomaticSpeechRecognitionPipeline—fill-maskFillMaskPipeline—question-answeringQuestionAnsweringPipeline—zero-shot-classificationZeroShotClassificationPipeline—feature-extractionFeatureExtractionPipeline—image-classificationImageClassificationPipeline—image-segmentationImageSegmentationPipeline—object-detectionObjectDetectionPipeline—document-question-answeringDocumentQuestionAnsweringPipeline—image-text-to-textImageTextToTextPipeline—十二、官方文档示例模型与任务组合使用pipeline()的 docstring 中还给出了「指定模型 指定 tokenizer」的经典组合写法src/transformers/pipelines/init.pyfrom transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer # 情感分析 pipeline默认模型 analyzer pipeline(sentiment-analysis) # NER pipeline显式传入模型与 tokenizer model AutoModelForTokenClassification.from_pretrained(dbmdz/bert-large-cased-finetuned-conll03-english) tokenizer AutoTokenizer.from_pretrained(google-bert/bert-base-cased) recognizer pipeline(ner, modelmodel, tokenizertokenizer)结语Pipelines 把「预处理 → 推理 → 后处理」整条链路收敛为一行代码同时通过流式输入、batching、chunk batching、FP16 与自定义子类化保留了工程上所需的灵活性。核心要点回顾入口统一所有任务都从pipeline()工厂函数出发任务、模型、预处理组件均可自动推断吞吐优先GPU 上跑大批量静态数据时用 batching但必须先测基准、警惕长度不齐导致的 padding 放大与 OOM特殊任务特殊处理零样本分类与问答走ChunkPipeline单输入多前向由框架自动管理加速手段FP16dtypetorch.float16是 GPU 上零成本提速的最简方案扩展方式子类化Pipeline/具体任务类重写preprocess/forward/postprocess或通过pipeline_class参数接入自定义实现。如需了解 pipeline 中涉及的 tokenizer、image processor、feature extractor 与 processor 的详细 API可继续阅读仓库中的 tokenizer、image_processor、feature_extractor 与 processors 文档实现全新 pipeline 的完整指南见 add_new_pipeline。【免费下载链接】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创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进