Pipeline

Pipeline 是一个简单但功能强大的推理 API，可用于使用 Hugging Face Hub 中的任何模型执行各种机器学习任务。

使用任务特定的参数定制 Pipeline 以适应您的任务，例如为自动语音识别 (ASR) Pipeline 添加时间戳以转录会议记录。Pipeline 支持 GPU、Apple Silicon 和半精度权重，以加速推理并节省内存。

Transformers 有两个 Pipeline 类，一个通用的 Pipeline 和许多独立的任务特定 Pipeline，例如 TextGenerationPipeline 或 VisualQuestionAnsweringPipeline。通过在 Pipeline 的 `task` 参数中设置任务标识符来加载这些独立的 Pipeline。您可以在其 API 文档中找到每个 Pipeline 的任务标识符。

每个任务都配置为使用默认的预训练模型和预处理器，但如果需要使用不同的模型，可以使用 `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

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device="cuda")
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，从任务特定参数到性能优化。

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

设备

Pipeline 兼容多种硬件类型，包括 GPU、CPU、Apple Silicon 等。使用 `device` 参数配置硬件类型。默认情况下，Pipeline 在 CPU 上运行，由 `device=-1` 表示。

GPU

Apple silicon

批量推理

Pipeline 也可以使用 `batch_size` 参数处理批量输入。批量推理可能会提高速度，尤其是在 GPU 上，但这并非保证。其他变量，如硬件、数据和模型本身，都会影响批量推理是否能提高速度。因此，默认情况下禁用批量推理。

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

from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device="cuda", 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 transformers.pipelines.pt_utils import KeyDataset
import datasets

# KeyDataset is a utility that returns the item in the dict returned by the dataset
dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipeline = pipeline(task="text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device="cuda")
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 接受每个独立任务 Pipeline 支持的任何参数。请务必查看每个独立的任务 Pipeline，以了解可用的参数类型。如果您找不到适用于您用例的参数，请随时在 GitHub 上提交问题请求！

以下示例演示了一些可用的任务特定参数。

自动语音识别

文本生成

分块批量处理

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

对于某些数据类型，单个输入（例如，一个非常长的音频文件）可能需要分成多个部分才能处理
对于某些任务，例如零样本分类或问答，单个输入可能需要多次前向传递，这可能会导致 `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 datasets import load_dataset

dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipeline = pipeline(task="text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device="cuda")
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` 中添加 `load_in_8bit=True`。

import torch
from transformers import pipeline, BitsAndBytesConfig

pipeline = pipeline(model="google/gemma-7b", torch_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'}]

< > 在 GitHub 上更新