Transformers pipeline 怎样指定推理设备

Transformers pipeline 的推理设备应在创建对象时通过 device 指定。强制 CPU 可写 device="cpu" 或 -1；单张 CUDA 卡可写 device="cuda:0"，需要另一张可用卡时使用对应编号。普通单设备模型由 pipeline 放到目标设备，文本等原始输入也由它完成预处理和张量移动。使用 device_map 分派的模型则不应同时设置 device。

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

先明确设备名称和当前进程可见范围

本文以 PyTorch 文本分类 pipeline 为例。环境需要安装相互兼容的 torch、transformers，并准备一个可用于文本分类的模型。机器存在 GPU，不等于当前 Python 环境就有可用 CUDA；容器、环境和驱动配置也会影响设备可见性。判断当前进程能否使用 CUDA，可以从 torch.cuda.is_available() 开始。

省略 device 的默认行为应以所用版本为准；需要固定部署位置时，显式写明目标比依赖默认值更清楚。device 是对象初始化参数，不能把它当作每次文本推理调用时临时切卡的通用开关。

创建单设备 pipeline 并执行一次推理

下面代码使用已经准备好的本地模型目录，避免把网络下载和设备问题混在一起。目录必须是真正的文本分类模型及对应分词器；如果文件缺失，local_files_only 会让加载失败，应先补齐本地资源。示例不在本机运行，也不预设分类标签或分数。

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline

# 换成已经准备好的文本分类模型目录，包含权重、配置和分词器文件。
model_dir = "./models/text-classifier"
tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
model = AutoModelForSequenceClassification.from_pretrained(
    model_dir, local_files_only=True
)

# 无 CUDA 时明确使用 CPU，不依赖不同版本的默认设备选择行为。
target = "cuda:0" if torch.cuda.is_available() else "cpu"
classifier = pipeline(
    task="text-classification",
    model=model,
    tokenizer=tokenizer,
    device=target,
)

# 普通文本输入交给 pipeline 预处理与设备移动。
try:
    print("pipeline device:", classifier.device)
    print("parameter device:", next(classifier.model.parameters()).device)
    print(classifier("The delivery arrived earlier than expected.", truncation=True))
finally:
    # 释放本例持有的模型引用；长驻服务应复用模型而非逐次销毁。
    del classifier
    del model

读者运行时，应先看 pipeline device 和 parameter device 是否与 target 对应，再看实际分类结果。这里的参数位置检查只适合普通单设备模型；分派模型的一项参数所在位置，不能代表整个模型。若选择 cuda:1，除了 CUDA 可用，还需确认有足够的可见设备。

同一服务通常在启动时创建 pipeline，并在请求间复用。代码中的 finally 只是演示对象引用的生命周期，不表示 GPU 缓存必定立即归还给操作系统，也不建议每处理一条文本就重新加载模型。

区分模型放置和输入张量移动

把字符串交给 classifier 时，分词器通常仍在主机侧生成输入张量，pipeline 再把需要的张量送到推理设备。设置 device="cuda:0" 不代表分词、文件读取和 Python 后处理全部在 GPU 上完成；看到 CPU 有活动不能直接判断设备参数失效。

如果绕过 pipeline，直接使用 classifier.model，就需要自己负责张量设备。下面片段应放在前例 finally 清理之前，使用相同的单设备模型；它不是接在整段清理代码之后直接运行的独立脚本。

# 与前例是两种调用方式；本段在前例清理之前使用同一个单设备模型。
encoded = tokenizer("The delivery arrived earlier than expected.", return_tensors="pt")
# 直接调用模型时，由调用方把各张量放到模型所在设备。
encoded = {name: value.to(target) for name, value in encoded.items()}
with torch.inference_mode():
    logits = classifier.model(**encoded).logits
print(logits.shape)

错误地只把模型放到 GPU，却把 tokenizer 产生的 CPU 张量直接传给模型，容易出现设备不匹配。反过来，普通字符串本身没有 .to("cuda") 这种张量移动语义。先确定调用的是 pipeline 还是原始模型，再选择谁来移动输入。

使用 device_map 时交给统一分派管理

device_map 面向模型加载和模块分派，可以按模型与资源条件把模块放在不同位置；device="cuda:0" 则表达普通单设备目标。不要既用 device_map="auto"，又给同一个 pipeline 指定 device。若传入模型已经带有 hf_device_map，也应保留分派机制管理，而不是再强行把整模型移动到某张卡。

需要 device_map="auto" 时，应按该模式准备 Accelerate 等相应依赖，并查看实际的 hf_device_map。它可能包含 GPU 与 CPU 的组合，不是“自动把所有层塞进一张显卡”的承诺。对于量化或带卸载钩子的模型，后续整体 .to() 也可能不受支持，不能套用普通模型迁移方法。

初次排查可以先让一个小的普通模型在单设备上跑通，再进入分派方案。遇到显存不足时，先看模型体积、输入长度和批大小；这类资源问题与设备名称写错是不同原因。这里不把精度、量化或分布式部署混入同一个最小示例。

常见问题

device=0 就能使用所有显卡吗？

不能。它选择一个设备，不会自动建立数据并行或模型分片。多设备模型放置属于另一套部署配置。

为什么 CPU 仍然占用明显？

分词、读取和后处理仍可能使用 CPU。应观察实际模型放置与计算过程，不能只看单一 CPU 利用率。

为什么传入现成模型后指定 device 报冲突？

先确认模型是否已由 Accelerate 等机制分派并带有设备映射。普通单设备移动与已分派模型不能用同一种强制迁移方式处理。

明确选一个设备、在初始化时固定位置、让输入遵守同一放置约定，是最小单设备推理闭环。模型改用 device_map 后，则应把放置责任统一交给分派机制。

参考：Transformers Pipelines 参数参考、Pipeline 设备管理实现。