Pipeline

Pipeline 是一个简单而强大的推理 API，可直接用于各种机器学习任务，并支持 Hugging Face Hub 上的任意模型。

通过添加任务特定参数（例如为自动语音识别（ASR）流水线添加时间戳以转录会议记录）来定制 Pipeline，使其适配您的任务。Pipeline 支持 GPU、Apple Silicon 以及半精度权重，以加速推理并节省内存。

Transformers 包含两个 pipeline 类：通用的 Pipeline 以及许多针对特定任务的独立 pipeline，例如 TextGenerationPipeline。通过在 Pipeline 中设置 task 参数为相应的任务标识符，即可加载这些独立的 pipeline。您可以在各 pipeline 的 API 文档中找到其对应的任务标识符。

每个任务都配置了默认的预训练模型和预处理器，但如果您想使用其他模型，可以通过 model 参数进行覆盖。

例如，若要将 TextGenerationPipeline 与 Gemma 2 配合使用，请设置 task="text-generation" 和 model="google/gemma-2-2b"。

from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b")
pipeline("the secret to baking a really good cake is ")
[{'generated_text': 'the secret to baking a really good cake is 1. the right ingredients 2. the'}]

当存在多个输入时，请以列表形式传入。

from transformers import pipeline
from accelerate import Accelerator

device = Accelerator().device

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device=device)
pipeline(["the secret to baking a really good cake is ", "a baguette is "])
[[{'generated_text': 'the secret to baking a really good cake is 1. the right ingredients 2. the'}],
 [{'generated_text': 'a baguette is 100% bread.\n\na baguette is 100%'}]]

本指南将为您介绍 Pipeline，展示其各项功能，并说明如何配置其不同参数。

任务

Pipeline 兼容多种跨模态的机器学习任务。只需向 pipeline 传入合适的输入，其余工作将由它自动完成。

以下是一些使用 Pipeline 处理不同任务和模态的示例。

自动语音识别

图像分类

视觉问答

参数

至少，Pipeline 只需要一个任务标识符、一个模型以及相应的输入。不过，还有许多参数可用于配置流水线，从特定任务的参数到性能优化。

本节将向您介绍一些更为重要的参数。

设备

Pipeline 兼容多种硬件类型，包括 GPU、CPU、Apple Silicon 等。使用 device 参数来配置硬件类型。默认情况下，如果未设置 device，Pipeline 会自动将模型放置在第一个可用的加速器（CUDA GPU、Apple Silicon MPS、XPU 等）上，仅在没有可用加速器时才回退到 CPU。传入 device="cpu" 可强制其在 CPU 上运行。

GPU

Apple 芯片

批量推理

Pipeline 还可以通过 batch_size 参数处理输入批次。批量推理可能会提升速度，尤其是在 GPU 上，但并不保证一定有效。硬件、数据以及模型本身等其他变量也会影响批量推理是否能提升速度。因此，批量推理默认处于禁用状态。

在下面的示例中，当有 4 个输入且 batch_size 设置为 2 时，Pipeline 每次向模型传递 2 个输入的批次。

from transformers import pipeline
from accelerate import Accelerator

device = Accelerator().device

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device=device, batch_size=2)
pipeline(["the secret to baking a really good cake is", "a baguette is", "paris is the", "hotdogs are"])
[[{'generated_text': 'the secret to baking a really good cake is to use a good cake mix.\n\ni’'}],
 [{'generated_text': 'a baguette is'}],
 [{'generated_text': 'paris is the most beautiful city in the world.\n\ni’ve been to paris 3'}],
 [{'generated_text': 'hotdogs are a staple of the american diet. they are a great source of protein and can'}]]

批量推理的另一个良好用例是在 Pipeline 中流式处理数据。

from transformers import pipeline
from accelerate import Accelerator
from transformers.pipelines.pt_utils import KeyDataset
import datasets

device = Accelerator().device

# KeyDataset is a utility that returns the item in the dict returned by the dataset
dataset = load_dataset("stanfordnlp/imdb", name="plain_text", split="unsupervised")
pipeline = pipeline(task="text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device=device)
for out in pipeline(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    print(out)

在判断批量推理是否有助于提升性能时，请牢记以下一般性经验法则。

唯一能确定结果的方法是在你的模型、数据和硬件上测量性能。
如果受延迟限制（例如实时推理产品），请勿进行批量推理。
如果使用 CPU，请勿进行批量推理。
如果不确定数据的 sequence_length，请勿进行批量推理。请测量性能，逐步增加 sequence_length，并加入内存溢出（OOM）检查，以便在失败时恢复。
如果sequence_length是固定的，请使用批量推理，并不断增大批量大小，直到出现 OOM 错误。GPU 越大，批量推理带来的收益越明显。
如果您决定执行批量推理，请务必确保能够处理 OOM 错误。

特定任务的参数

Pipeline 接受各个任务流水线所支持的任意参数。请务必查阅各个任务流水线，以了解可用的参数类型。如果您找不到适用于自身使用场景的参数，欢迎在 GitHub 上提交一个 issue 来提出需求！

以下示例展示了部分可用的任务特定参数。

自动语音识别

文本生成

块批处理

在某些情况下，你需要分块处理数据。

对于某些数据类型，单个输入（例如一个非常长的音频文件）可能需要先被分块为多个部分，然后才能进行处理
对于某些任务，例如零样本分类或问答，单个输入可能需要多次前向传播，这可能会导致 batch_size 参数出现问题

ChunkPipeline 类专为处理这些用例而设计。两个 pipeline 类的使用方式相同，但由于 ChunkPipeline 可以自动处理批处理，因此您无需担心输入会触发多少次前向传播。相反，您可以独立于输入来优化 batch_size。

下面的示例展示了它与 Pipeline 的区别。

# ChunkPipeline
all_model_outputs = []
for preprocessed in pipeline.preprocess(inputs):
    model_outputs = pipeline.model_forward(preprocessed)
    all_model_outputs.append(model_outputs)
outputs =pipeline.postprocess(all_model_outputs)

# Pipeline
preprocessed = pipeline.preprocess(inputs)
model_outputs = pipeline.forward(preprocessed)
outputs = pipeline.postprocess(model_outputs)

大型数据集

对于大型数据集的推理，您可以直接遍历数据集本身。这样可以避免立即为整个数据集分配内存，也无需自行创建批次。请尝试使用 batch_size 参数进行批量推理，以查看是否能提升性能。

from transformers.pipelines.pt_utils import KeyDataset
from transformers import pipeline
from accelerate import Accelerator
from datasets import load_dataset

device = Accelerator().device

dataset = load_dataset("stanfordnlp/imdb", name="plain_text", split="unsupervised")
pipeline = pipeline(task="text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device=device)
for out in pipeline(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    print(out)

使用 Pipeline 在大型数据集上运行推理的其他方法包括使用迭代器或生成器。

def data():
    for i in range(1000):
        yield f"My example {i}"

pipeline = pipeline(model="openai-community/gpt2", device=0)
generated_characters = 0
for out in pipeline(data()):
    generated_characters += len(out[0]["generated_text"])

大型模型

Accelerate 为使用 Pipeline 运行大型模型提供了几项优化。请确保已先安装 Accelerate。

!pip install -U accelerate

device_map="auto" 设置非常有用，它可以自动将模型优先分配到速度最快的设备（GPU）上，如果还有剩余，再分发到其他较慢的设备（如 CPU、硬盘）。

Pipeline 支持半精度权重（torch.float16），这可以显著提升速度并节省内存。对于大多数模型，尤其是较大的模型，性能损失可以忽略不计。如果您的硬件支持，可以改用 torch.bfloat16 以获得更大的数值范围。

输入会在内部转换为 torch.float16，并且仅适用于使用 PyTorch 后端的模型。

最后，Pipeline 还支持量化模型，以进一步降低内存占用。请确保已安装 bitsandbytes 库，然后在 pipeline 的 model_kwargs 中添加 quantization_config。

import torch
from transformers import pipeline, BitsAndBytesConfig

pipeline = pipeline(model="google/gemma-7b", dtype=torch.bfloat16, device_map="auto", model_kwargs={"quantization_config": BitsAndBytesConfig(load_in_8bit=True)})
pipeline("the secret to baking a good cake is ")
[{'generated_text': 'the secret to baking a good cake is 1. the right ingredients 2. the right'}]

Update on GitHub