Hugging Face Pipeline多进程/多线程实现零样本分类推理加速求助

• 你在主进程将模型加载到device=0（单GPU）后，开启多进程时spawn模式会为每个子进程复制一份完整的模型权重到同一块GPU，进程数越多显存占用越高，上万条数据时直接触发OOM
• 单GPU的推理瓶颈在GPU的计算吞吐量，用CPU多进程调度单GPU只会增加额外的IPC通信、数据拷贝开销，反而抵消并行收益，甚至比单进程更慢
• 默认pipeline是单条推理，没有启用batch处理，GPU的并行算力没有被充分利用

方案1：优先启用pipeline批处理（最简单，性价比最高） #

直接用pipeline自带的批处理能力，不需要改多进程，代码示例：

import os
import torch
from transformers import pipeline

os.environ["TOKENIZERS_PARALLELISM"] = "false"

classifier = pipeline(
    "zero-shot-classification", 
    model="typeform/distilbert-base-uncased-mnli", 
    device=0,
    torch_dtype=torch.float16 # 可选，开启FP16精度，提速降显存
)

if __name__ == '__main__':
    # 直接传入整个句子列表，batch_size根据显存大小调整，一般8-64都可
    results = classifier(d_lst, tag_values, batch_size=16)
    # 处理结果取top1标签和分数
    output = [(res['labels'][0], res['scores'][0]) for res in results]

• 这个方案不需要改复杂逻辑，只要加batch_size参数就能把GPU算力用满，提速效果远高于多进程调度单卡，显存占用也可控

方案2：多卡场景用多进程单卡单模型架构 #

如果你有多张GPU，可以每个进程绑定一张独立的GPU，避免多个进程抢同一块GPU的资源：

from pathos.multiprocessing import ProcessingPool as Pool
import multiprocess.context as ctx
import os
from functools import partial
ctx._force_start_method('spawn') 

os.environ["TOKENIZERS_PARALLELISM"] = "false"

def process_chunk(gpu_id, labels, chunk_sentences):
    # 模型必须在子进程内部初始化，不能在主进程初始化后传进去
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    from transformers import pipeline
    import torch
    classifier = pipeline(
        "zero-shot-classification", 
        model="typeform/distilbert-base-uncased-mnli", 
        device=0,
        torch_dtype=torch.float16
    )
    results = classifier(chunk_sentences, labels, batch_size=16)
    return [(res['labels'][0], res['scores'][0]) for res in results]

if __name__ == '__main__':
    # 按实际GPU数量调整
    num_gpus = 2
    chunks = [d_lst[i::num_gpus] for i in range(num_gpus)]
    pool = Pool(num_gpus)
    results = pool.map(partial(process_chunk, labels=tag_values), range(num_gpus), chunks)
    pool.close()
    pool.join()
    # 合并结果
    final_results = []
    for res in results:
        final_results.extend(res)

• 这个方案只有多卡场景下才有收益，单卡不要使用
• 模型必须在子进程内部初始化，避免不必要的显存拷贝

额外优化手段 #

• 如果对精度要求不高，可以开启int8量化，在pipeline初始化时加load_in_8bit=True参数，显存占用再降一半
• 数据量特别大的场景，可以把整个数据集拆成多个小批次依次处理，避免一次性加载所有数据到显存
• 不需要梯度计算的推理场景，全局加torch.no_grad()上下文管理器，进一步降低显存占用

内容的提问来源于stack exchange，提问作者Code_7