10 Transformers - 任务容器类 Pipeline

管道(pipelines) 是使用模型进行推理的一种强大易用的方法;是从 Transformers库中抽取出大多数复杂代码构造而成的,用于提供完成特定功能的任务类集合,包括命名实体识别、掩码语言建模、情感分析、特征提取和问答等。

对于封装的具体任务对象,有两点需要注意:

  • Pipeline 不仅是实现特定功能任务类的父类;同时它又是继承于它的子类的容器。它已将各个任务对象封装进去,通过它可以调用这些完成具体功能的任务对象;
  • 具体的任务类可用于专有领域,如音频、计算机视觉、自然语言处理和多模态等。

Pipeline

Pipeline 类是 Transformers 库的核心接口,是管道中所有可用的任务类的封装器;是管道中可用的各个任务类的父类;是利用预训练模型进行推理的最简单方式。

Pipeline 支持 gpuApple Silicon和半精度权重,以加速推理并节省内存。

Transformers提供了两种方式执行任务,一种是通过通用的Pipeline,另外一种是使用具体的任务子类 如 TextGenerationPipelineVisualQuestionAnsweringPipeline等。 使用 Pipeline 则更简单,通过在 Pipelinetask 参数中设置任务标识符 来加载这些独立的任务类;通过设置特定的参数可以让Pipeline更适应对应的特定任务;如通过为自动语音识别设置 timestamps 参数以转录会议记录等。

每个任务类都配置有默认的预训练模型和预处理器。但如果需要,也可以通过配置 Pipelinemodel 参数来覆盖掉默认的模型。

调用流程

管道中每个任务都有对应的特定功能处理类,但使用通用的 Pipeline 调用则更为简单,该类集成了管道中所有实现特定功能的任务类。Pipeline 会根据指定的任务标识符自动加载默认模型和适用于所处理任务的预处理类,以实现推理功能。基本的流程如下:
call init run_single run_multi 是 否 否 传入输入和其他参数 解析参数(_sanitize_parameters) 输出 预处理(preprocess) 前馈处理(forward) 后处理(postprocess) 列表处理完? 预处理(preprocess) 单项输入? 前馈处理(forward) 后处理(postprocess) 通过传入的参数初始化配置 传入任务标识、模型标识/路径 或其他配置其他参数 解析参数(_sanitize_parameters) 开始 结束

从上面的流程可以看到,使用 Pipeline 很简单:第一步,实例化 Pipeline;第二步,调用实例函数获取结果 。如果没有的额外需要,两步就完成了任务的处理。下面举个文本生成的例子来了解如何使用 Pipeline

复制代码
from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b")
pipeline("the secret to baking a really good cake is ")

上例指定任务为文本生成(text-generation),使用模型为 google/gemma-2-2b,结果为:

复制代码
[{'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的构造函数中, taskmodel参数至少有一个,用来确定任务的类型,其他的参数可选;另外,有些参数是从特定任务中提取的,而另一些是通用参数。 一般来说,你可以在任何地方指定参数,如:

复制代码
generator = pipeline(model="openai/whisper-large", my_parameter=1)
out = generator(...)  # This will use `my_parameter=1`.
out = generator(..., my_parameter=2)  # This will override and use `my_parameter=2`.
out = generator(...)  # This will go back to using `my_parameter=1`.

Pipleline的构造函数原型如下:

复制代码
Pipeline(self,
    model: Union["PreTrainedModel", "TFPreTrainedModel"],
    tokenizer: Optional[PreTrainedTokenizer] = None,
    feature_extractor: Optional[PreTrainedFeatureExtractor] = None,
    image_processor: Optional[BaseImageProcessor] = None,
    processor: Optional[ProcessorMixin] = None,
    modelcard: Optional[ModelCard] = None,
    framework: Optional[str] = None,
    task: str = "",
    args_parser: ArgumentHandler = None,
    device: Union[int, "torch.device"] = None,
    torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
    binary_output: bool = False,
    **kwargs
    )
task 参数

str 简单类型,可选 ,但与 model 参数二者必须至少保持一个,以便确定任务类型。

定义 Pipeline 所要执行的任务项;当前 Pipeline 所支持的任务项如下:

  • audio-classification: 音频分类标识符
  • automatic-speech-recognition: 自动语音标识符
  • depth-estimation: 深度估计标识符
  • document-question-answering: 文档问答标识符
  • feature-extraction: 特征提取标识符
  • fill-mask: 掩码填充标识符
  • image-classification: 图像分类标识符
  • image-feature-extraction: 图像特征提取标识符
  • image-segmentation: 图像分割标识符
  • image-text-to-text: 图像文本到文本标识符
  • image-to-image: 图像到图像标识符
  • image-to-text: 图像到文本标识符
  • mask-generation: 蒙版生成标识符
  • object-detection: 对象检测标识符
  • question-answering: 文本问答标识符
  • summarization: 文本摘要标识符
  • table-question-answering: 表格问答标识符
  • text2text-generation: 文生文标识符
  • text-classification (又名sentiment-analysis ): 文本分类标识符
  • text-generation: 文本生成标识符
  • text-to-audio (又名text-to-speech): 语音合成标识符
  • token-classification (又名 ner ): 命名实体识别标识符
  • translation: 翻译标识符
  • translation_xx_to_yy: 翻译标识符(指定源语种和目标语种)
  • video-classification: 视频分类标识符
  • visual-question-answering: 视频文档标识符
  • zero-shot-classification: 零样本文本分类标识符
  • zero-shot-image-classification: 零样本图像分类标识符
  • zero-shot-audio-classification: 零样本音频分类标识符
  • zero-shot-object-detection: 零样本对象检测标识符
model 参数

复合类型,支持str简单类型, 或 PreTrainedModel TFPreTrainedMode 等对象类型,可选 ,但与 task 参数两者必须有一个,确定任务的类型。

⚠️ 注意

如果 devicemodel 参数同时都设置了值,而且,当model中的 hf_device_map 参数也设置了值,则会提示异常。

定义 pipeline用来做预测的模型,可以是模型标识符,也可以是从 PreTrainedModel(适用于 PyTorch)或 TFPreTrainedModel(适用于 TensorFlow)继承的预训练模型的实际实例。

如果未指定,则会加载为任务加载默认的模型。

device 参数

复合类型,支持 strint简单类型,或是 torch.device 枚举类型。可选

pipeline支持很多硬件类型包括 GPU, CPU,Apple Silicon等。pipeline默认使用 CPU;即设定device=-1' 则表示使用GPU。如果设定device=0 如果你使用 device=n,则 pipeline 会自动将模型放置在指定的设备上。 无论你是使用 PyTorch 还是 Tensorflow,都可以使用这个参数。

⚠️ 注意

如果 devicemodel 参数同时都设置了值,而且,当model中的 hf_device_map 参数也设置了 值,则会提示异常。

你也可以让 Accelerate(一个用于分布式训练的库)自动选择如何在适当的设备上加载和存储模型权重。如果你有多个设备,这尤其有用。首先在最快的设备上加速加载并存储模型权重,然后根据需要将权重移动到其他设备(CPU、硬盘)。设置device_map="auto"Accelerate 选择设备。当然,如果需要 Accelerate,则首先需要安装它:

复制代码
!pip install -U accelerate

然后指定device="auto"

复制代码
from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device_map="auto")
pipeline("the secret to baking a really good cake is ")

如果pipeline 运行在 Apple silicon 上,需要设置为 device=mps

复制代码
from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device="mps")
pipeline("the secret to baking a really good cake is ")
torch_dtype 参数

复合类型,支持 str 简单类型或 torch.dtype 枚举类型。可选

直接作为 model_kwargs发送(只是一个更简单的快捷方式),以使用此模型的可用精度(torch.float16torch.bfloat16..."auto")。

tokenizer 参数

复合类型,支持 str 简单类型或 PreTrainedTokenizer对象类型。可选

pipeline 将使用该 tokenizer 对模型数据编码。这可以是标识符,也可以是继承自 PreTraindTokenizer 的实际预训练标记器实例。

如果没有提供,如果 processor 已经赋予,则通过 processor 中的 tokenizer属性提取;否则将加载给定模型的默认标记器(如果它是字符串);如果模型未指定或不是字符串,则加载配置的默认标记器(如果它是字符串)。

如果config 参数也没有给出或不是字符串,则将加载给定任务的默认标记器。

feature_extractor 参数

复合类型,支持 str 简单类型或 PreTrainedFeatureExtractor对象类型。可选

pipeline 将使用的特征提取器对模型的数据进行编码。

特征提取器用于非NLP模型,如语音或视觉模型以及多模态模型。多模态模型还需要传递标记器。

如果没有指定,将加载给定模型的默认特征提取器(如果是字符串)。

如果未指定模型或不是字符串,如果 processor 已经赋予,则通过 processor 中的 feature_extractor属性提取; 否则则加载配置的默认特征提取器(如果是字符串);如果config参数也没有给出或不是字符串,则将加载给定任务的默认特征提取器。

image_processor 参数

复合类型,支持 str 简单类型或 BaseImageProcessor对象类型。可选

指定 pipeline 对模型的图像进行预处理的处理器。该参数值可以是字符串标识名,也可以是继承自 BaseImageProcessor的对象实例。

图像处理器用于需要图像输入的视觉模型多模态模型 。多模态模型还需要传递 tokenizer

如果没有指定,如果 processor 已经赋予,则通过 processor 中的 image_processor属性提取;否则将加载给定 model 的默认图像处理器(如果是字符串)。如果 model 未指定或不是字符串,则加载 config 的默认图像处理器(如果是字符串)。

processor 参数

复合类型,支持 str 简单类型或 ProcessorMixin对象类型。可选

指定 pipeline 预处理模型的数据的处理器。可以是字符串标识符,也可以是从 ProcessorMixin 继承的处理器对象实例。指定的处理器用于需要多模态输入的多模态模型,例如,需要文本和图像输入的模型。

如果没有提供,将加载给定 model 的默认处理器(如果是字符串)。如果 model 未指定或不是字符串,则加载 config 参数配置的默认处理器(如果是字符串)。

framework 参数

str 简单类型。可选

指定使用的框架, pt 表示 PyTorch 框架;tf 表示 TensorFlow框架。注意:指定的框架必须确定已经安装

如果没有指定框架,将默认为当前安装的框架。如果没有指定框架并且两个框架都安装了,则使用 model 对应的默认框架,如果没有提供模型,则默认使用的是 PyTorch 框架。

binary_output 参数

布尔型,可选

指定是以序列化格式输出(即 pickle)或原始输出数据(如文本)。

modelcard 参数

复合类型,支持 str 简单类型或 ModelCard 对象类型。 可选

设置输入此 Pipeline实例所需模型的模型卡。

args_parser 参数

ArgumentHandler 对象类型,可选

引用负责对提供给此 Pipeline 实例的参数对象解析参数。

kwargs 参数

Dict[str, Any] 类型。可选

为继承于 Pipeline 类的特定任务类传递其所需的其他关键字参数。各特定任务类所需的具体特定参数在特定任务类上进行说明,在此不做说明。

assistant_model 参数

复合类型,支持 str简单类型或 PreTrainedModel对象类型。可选

为任务类配置的辅助模型。

assistant_tokenizer 参数

复合类型,支持 str简单类型或 `` 对象类型。可选

为任务类配置的辅助标记器。

batch_size 参数

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)

在确定批量推理是否有助于提高性能时,请记住以下一般经验规则:

  1. 唯一确定的方法是测量模型、数据和硬件的性能。
  2. 如果您受限于延迟(例如实时推理产品),请不要批量推理。
  3. 如果您使用 CPU,请不要批量推理。
  4. 如果你不知道数据的 sequence_length,请不要批量推断。测量性能,迭代地添加到 sequence_length中,并包括内存不足(OOM)检查以从故障中恢复。
  5. 如果你的 sequence_length是规则的,可采用批量推理,并一直推送它,直到遇到 OOM 错误。GPU 越大,对批量推理的帮助就越大。
  6. 如果你决定进行批量推理,一定要确保你可以处理 OOM 错误。
num_workers 参数

int类型,可选

是否多进程读取数据;默认0,为0时使用主进程来处理数据。大于0则多进程处理数据,加快数据处理速度。

torch_dtype 属性

复合类型,支持 str 简单类型或 torch.dtype 枚举类型。

获取 torch 的类型。

__call__ 内置方法

__call__ 是 Python 的特殊方法(magic method),用于使对象变成可调用(callable)。当一个对象实现了 该方法时,就可以像调用函数一样调用这个对象。

该内置函数是 Pipeline 的核心内置函数,完成对输入做预处理注入标记等,推理、后处理等一系列的过程。 函数原型如下:

复制代码
__call__(self, inputs, *args, num_workers=None, batch_size=None, **kwargs)

其中:

  • inputs:需要推理的输入,复合类型,可以是 str字符串类型,也可以 list 类型,也可以是 Dict类型;
  • args:命令行参数,扩展,目前暂未用到;
  • num_works:工作进程数,如果为0,表示主进程处理,大于1表示使用的工作进程数;
  • batch_size:指定批次大小,如为2,则表示一个批次可以同时传入两个输入;
  • kwargs:特定任务所需的特定参数均通过 kwargs 传入,具体参数解析由 _sanitize_parameters 函数处理,参数具体意义分别在具体的特定任务类中说明,在此不做累述。

_sanitize_parameters 内置方法

_sanitize_parameters 内置方法是一个抽象方法,需要子类自主实现。该函数在 __init____call__两个内置函数中分别调用,用来为预处理、前驱处理和后处理从给定的参数字典中提取可用的参数。如果参数未传入,则使用默认的参数。函数原型如下:

复制代码
_sanitize_parameters(self, **pipeline_parameters)

其中:

  • pipeline_parameters:传入的参数字典,实际就是在 __init____call__两个内置函数调用时传入的 kwargs

_forward 内置方法

_forward是实现前馈处理算法的抽象内置函数,需要子类实现。

该函数在预处理结束后接收准备好的字典,并在模型上运行它。此方法可能涉及GPU或CPU,但应该与之无关。设计此函数的目的是因为预处理和后处理都会执行前馈处理,因此,这个方法如有可能要保证它高效运行。

该函数不应直接调用,在 forward 方法中调用它。他们实现方式基本是相同的,只是 forward 包含围其他的额外代码,以确保张量和模型在同一设备上,在代码中屏蔽训练部分以使它推理更快。

原型如下:

复制代码
_forward(self, input_tensors: Dict[str, GenericTensor], **forward_parameters: Dict) -> ModelOutput

其中:

  • input_tensors: 预处理好的字典张量;
  • forward_parameters: 前馈处理所需的参数字典。

device_placement 方法

定义上下文管理器(Context Manager),允许以框架无关的方式在用户指定的设备上分配张量。

返回上下文管理器。

get_inference_context 方法

返回一个上下文管理器,用于在推理过程中关闭梯度计算,提高推理效率。函数原型:

复制代码
get_inference_context(self)

输入参数:无。

输出参数:返回上下文管理器 torch.no_grad

transform 方法

transform是专为 Scikit 库和 Keras 库提供的公开的接口方法,是一个抽象函数,需要子类实现。为实现数据转换用的,在前馈函数(forward)处理完成后,推理函数(__call__)前执行。原型为:

复制代码
transform(self, X)

predict 方法

predict是专为 Scikit 库和 Keras 库提供的公开的接口方法,是一个抽象函数,需要子类实现。为实现预测用的,在前馈函数(forward)处理完成后,推理函数(__call__)前执行。原型为:

复制代码
predict(self, X)

preprocess 方法

preprocess函数实现预处理功能,它是一个抽象函数,需要子类实现;设计思想为:对传入的原始输入预处理后为前馈处理提供的所有必须的预处理结果字典,并返回一个字典,其中至少应该包括一个张量,当然也可以包含其他项。

函数原型:

复制代码
preprocess(self, input_: Any, **preprocess_parameters: Dict) -> Dict[str, GenericTensor]

输入参数:

  • input_:原始输入;
  • preprocess_parameters:预处理所需的配置参数字典。

输出参数:张量字典

postprocess 方法

Postprocess 函数实现后处理功能,是一个抽象函数,需要子类实现具体功能。设计思想为:将接收 _forward 前馈方法的原始输出张量,并将它们重新格式化成更友好的结果。一般来说,它将输出一个列表或命令或结果(包含简单的字符串和数字)。

函数原型:

复制代码
postprocess(self, model_outputs: ModelOutput, **postprocess_parameters: Dict) -> Any

输入参数:

  • model_outputs:由 _forward 输出的原始张量
  • postprocess_parameters:后处理所需的配置参数字典

输出:由具体实现确定

forward 方法

foward 方法是对外提供的接口方法,实际处理和内置函数_forward 类似,只是比在 forward_parameters增加了些额外处理代码。原型如下:

复制代码
forward(self, model_inputs, **forward_params)

其中:

  • model_inputs:预处理好的字典张量;
  • forward_params:前馈处理所需的参数字典。

iterate 方法

迭代给定的输入数据集,对每个输入数据进行模型处理,并通过生成器返回结果。函数原型:

复制代码
iterate(self, inputs, preprocess_params, forward_params, postprocess_params)

输入参数:

  • inputs:待处理的原始输入;
  • preprocess_params:预处理所需的参数字典;
  • forward_params:前馈处理所需的参数字典;
  • postprocess_params:后处理所需的参数字典。

输出参数:返回推理完成后的结果集的迭代器。

ensure_tensor_on_device 方法

检查并确保PyTorch 张量在本实例的设备上。实质是调用内部函数 _ensure_tensor_on_device 实现的。函数原型:

复制代码
ensure_tensor_on_device(self, **inputs)
  • 输入参数:
    • inputs:输入张量,关键字参数应该是 torch.Tensor,其他部分忽略;
  • 输出参数:
    • 返回字典张量,和输入基本相同,只是多了设备信息。

check_model_type 方法

检查实例自身配置的模型是否在此实例的类所支持的支持模型列表上。函数原型:

复制代码
check_model_type(self, supported_models: Union[List[str], dict])

输入参数:

  • supported_modelsPipelien具体类所支持的模型名列表或者是带有模型类值得字典。

输出参数:无

get_iterator 方法

获取数据迭代器,用于模型推理过程中的数据加载和预处理。函数原型:

复制代码
get_iterator(
        self, inputs, num_workers: int, batch_size: int, preprocess_params, forward_params, postprocess_params
    )

输入参数:

  • inputs
  • num_workers:工作进程数;
  • batch_size:每个批处理大小,如为2,则表示为每个批次允许两个输入;
  • preprocess_params:预处理所需的参数字典;
  • forward_params:前馈处理所需的参数字典;
  • postprocess_params:后处理所需的参数字典。

输出参数:获取的结果是用于推理过程用的数据和处理器包括特征提取器、数据矫正器、数据集等的迭代器。

run_multi 方法

run_multi函数实现对原始多输入序列进行预处理、前馈推理和后处理后,返回处理后的结果集合。原型为:

复制代码
run_multi(self, inputs, preprocess_params, forward_params, postprocess_params)

输入参数:

  • inputs:待处理的原始输入;
  • preprocess_params:预处理所需的参数字典;
  • forward_params:前馈处理所需的参数字典;
  • postprocess_params:后处理所需的参数字典。

输出参数:返回推理完成后的结果集集合

run_single 方法

run_single函数实现对原始单输入进行预处理、前馈推理和后处理后,返回处理后的结果。原型为:

复制代码
run_single(self, inputs, preprocess_params, forward_params, postprocess_params)

输入参数:

  • inputs:待处理的原始输入;
  • preprocess_params:预处理所需的参数字典;
  • forward_params:前馈处理所需的参数字典;
  • postprocess_params:后处理所需的参数字典。

输出参数:返回推理完成后的结果集

save_pretrained 方法

将修改好的配置保存到指定的路径,原型如下:

复制代码
save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        safe_serialization: bool = True,
        **kwargs,
    )

输入参数:

  • save_directory:字符串路径或 os.PathLike 实例;
  • safe_serialization:保存为safetensors或为传统的 PyTorchTensorflow 的格式。True为使用 safetensors 方式;
  • kwargs:为特定任务所需的附加参数字典。

输出参数:无

加速

使用流加速

管道中所有任务类都支持批处理!无论何时,管道中的任务类接受的输入如果是列表(list)、数据集(dataset)或生成器(generator),都可以使用任务类的 "流" 能力。

复制代码
from transformers import pipeline
from transformers.pipelines.pt_utils import KeyDataset
import datasets

dataset = datasets.load_dataset("imdb", name="plain_text", split="unsupervised")
pipe = pipeline("text-classification", device=0)
for out in pipe(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    print(out)

然而,这并不意味着完全解决了性能问题,它有可能是10倍的加速,也有可能是5倍的减速,这取决于硬件、数据和实际使用的模型。下面是一个加速的例子:

复制代码
from transformers import pipeline
from torch.utils.data import Dataset
from tqdm.auto import tqdm

pipe = pipeline("text-classification", device=0)

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(f"Streaming batch_size={batch_size}")
    for out in tqdm(pipe(dataset, batch_size=batch_size), total=len(dataset)):
        pass

下面是一个减速的例子:

复制代码
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

与其他句子相比,这是一个偶尔很长的句子。在这种情况下,整个批处理将需要 400token,因此整个批处理将是 [64,400] 而不是 [64,4],从而导致严重减速。更糟糕的是,如果批量更大,程序就会崩溃。对于这个问题没有好的通用解决方案,只能根据您的用例来定。经验法则是:

  • 利用现有的硬件,不断的测量性能

  • 如果是延迟受限的(实时产品做推理),不要批处理。

  • 如果使用的是CPU,不要批处理。

  • 如果使用的吞吐量(在一堆静态数据上运行你的模型),在GPU上,然后:

    • 如果你不知道 sequence_length 的大小,默认情况下不要批处理,测量并尝试添加它,添加 OOM 检查以在失败时恢复(如果不控制 sequence_length,将会在某个时候恢复)。
    • 如果你的 sequence_length 是超规则的,那么批处理就更会出问题;不断测量并推动它直到你得到 OOM
    • GPU越大,批处理就有效
  • 一旦你启用批处理,确保你可以很好地处理 OOM

使用分块

zero-shot-classificationquestion-answering 在某种意义上是特别定制的任务类,即单个输入可能会产生多个前置处理。在正常情况下,这将因为 batch_size参数大小而导致问题。

为了规避这个问题,这两个管道都有点特殊,它们是继承于 ChunkPipeline 而不是直接继承于常规的 Pipeline

复制代码
preprocessed = pipe.preprocess(inputs)
model_outputs = pipe.forward(preprocessed)
outputs = pipe.postprocess(model_outputs)

将上面改造下变成:

复制代码
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)

对以上两段代码应做比较,使用的是相同的方式。这是一个简化的视图: Pipeline 可以自动处理批处理!这意味着您不必关心输入实际会触发多少前置处理,您可以独立于输入优化 batch_size

使用半精度浮点数

模型可以在 FP16 上运行,这可以在节省内存的同时显着提高GPU速度。大多数模型不会因此而遭受明显的性能损失。模型越大,出现这种情况的可能性就越小。

要启用 FP16 推理,通过设置 torch_dtype="torch"torch_dtype="Float16" 并传入。注意,这只适用于 PyTorch 后端。在传入输入前需将输入转换为半精度浮点数值。

相关推荐
拓端研究室4 小时前
2025人工智能AI研究报告:算力、应用、风险与就业|附1000+份报告PDF、数据、可视化模板汇总下载
人工智能
西***63475 小时前
核心技术解码:LED 视频处理器如何筑牢高清显示基石?
人工智能
王中阳Go5 小时前
三年前,我帮万人转Go;今天,聊聊Go/Java程序员如何抢占AI高地
人工智能·后端·go
STLearner5 小时前
AAAI 2026 | 时空数据(Spatial-temporal)论文总结[上](时空预测,轨迹挖掘,自动驾驶等)
大数据·人工智能·python·深度学习·机器学习·数据挖掘·自动驾驶
Coremail15 小时前
AI赋能,安全领航!Coremail受邀出席CERNET第三十一届学术年会
人工智能·安全
快乐非自愿5 小时前
AI低代码:敏捷开发与跨界协作的应用场景探索
人工智能·低代码·敏捷流程
Codebee5 小时前
于钉钉投诉的申诉与思考:行业观察的价值,在于敢说真实的问题
人工智能·产品经理·阿里巴巴
fie88895 小时前
基于MATLAB的多幅图像拼接
人工智能·计算机视觉·matlab
whaosoft-1435 小时前
51c视觉~3D~合集10
人工智能