Hugging Face pipelines 怎样为批量推理设置 batch size

如果你把一批文本交给 Hugging Face pipelines 做离线推理，batch_size 就是控制“每次送进模型多少条输入”的参数。最稳妥的做法不是直接填一个很大的数字，而是先确认设备和输入长度，再从较小批次开始记录吞吐、单请求延迟与显存风险。官方教程也提醒，批处理默认关闭，因为它不一定更快。

官方地址：https://huggingface.co/docs/transformers/main/pipeline_tutorial

先看 batch_size 影响的是哪一段

一次 pipeline 推理可以拆成输入读取、批次组织、模型推理、后处理和结果记录几个边界。batch_size 主要约束批次组织与模型输入的宽度，不会把不同任务自动合并成一个业务结果；输出仍然要按输入顺序整理。使用数据集时，Transformers 会通过数据加载器按批取数据，因此批次大小和输入长度会共同影响显存。

这里的关键判断是：你是在做离线吞吐，还是在做单请求响应。离线任务可以让多个输入共享一次模型调用；在线接口若每个请求只来一条数据，强行等待凑批反而会增加用户感知延迟。

用 KeyDataset 把批量参数放到数据集入口

对于较大的数据集，可以让 pipeline 逐项迭代，而不是一次把所有输入结果堆到内存中。下面是一个文本分类示例，代码只展示调用关系，未把输出数字当作本机实测结果：

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

# 用 GPU 承载模型；实际设备应按部署环境调整。
classifier = pipeline(
    task="text-classification",
    model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
    device=0,
)

dataset = load_dataset("stanfordnlp/imdb", split="unsupervised")

# 只取数据集中的 text 字段，并让 pipeline 每次处理 8 条。
for result in classifier(KeyDataset(dataset, "text"), batch_size=8, truncation="only_first"):
    # 这里可写入队列或文件；生产代码要处理下游失败和重试。
    print(result)

如果想从一开始固定调用默认值，也可以在创建 pipeline 时传入 batch_size；但对数据集推理，放在调用处更容易和具体任务的批处理策略一起记录。num_workers 负责数据加载并行度，不等于模型批次大小，两个参数不要混为一谈。

把 batch_size 当成风险控制参数

批次变大通常会增加一次模型调用的工作量，也可能提高吞吐；代价是显存峰值和单次处理延迟一起上升。尤其是文本长度不规整时，较长样本会让同批次的 padding 变多，所以“8 一定比 4 快”并不是可以直接套用的结论。

可以按下面的顺序做选择：

用小步递增和记录表收敛到可用值

调参时不要只记录“跑完了没有”。至少为每个批次保留设备、模型、输入长度范围、batch_size、总处理量、总耗时、吞吐、峰值显存和是否 OOM。先确认一组输入长度稳定的样本，再把 batch 从小到大逐档比较；出现 OOM 后回退到上一个安全值，并保留失败记录，方便日后更换模型或硬件时重新判断。

如果任务有严格的实时 SLA，批次大小还要受等待窗口限制：即使显存允许，排队凑批也可能让 p95 延迟失控。反过来，离线任务更关心单位时间完成多少样本，此时可以接受更高的单次延迟，但仍要为 OOM 和异常样本准备回退策略。

• 先确认 pipeline 实际使用的设备，再决定是否值得测试 batching。
• 让 sequence_length 或长度分布进入记录，避免只比较一个 batch 数字。
• 遇到 OOM 时减小批次或限制输入长度，不要用无限重试掩盖资源不足。

常见问题

batch_size 越大是不是一定越快？

不是。模型、硬件、输入长度和数据加载都会改变结果；官方建议用自己的模型、数据和设备测量。

CPU 上能不能设置 batch_size？

参数可以传入，但不代表会有收益。CPU 场景通常先以单条推理为基线，再用实际总耗时判断是否值得批处理。

为什么 batch_size 不大却仍然 OOM？

显存不仅由样本数量决定，还受每条输入的长度、padding、模型结构和精度影响。先检查长度分布，再回退批次或缩短输入。