如何在Transformers库中使用CUDA进行加速推理？解决GPU微调后推理缓慢及设备不匹配问题

你的推测完全正确——报错的核心原因就是模型在GPU上，但输入张量还留在CPU上，导致两者设备不兼容。下面分两种场景给你具体的解决方案：

一、不使用pipeline的手动推理场景 #

针对你第一段BERT分类的代码，只需要把输入张量也同步移到GPU上即可，步骤如下：

1. 先定义设备（兼容CPU/GPU自动切换）
2. 把模型移到目标设备
3. 把编码后的所有输入张量都移到同一设备
4. （可选但推荐）推理时关闭梯度计算，提升速度并节省显存

修改后的完整代码：

import torch
import transformers

txt = "This was nice place"
model_path = "你的模型路径"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载模型并移到目标设备
model = transformers.BertForSequenceClassification.from_pretrained(model_path, num_labels=24)
model.to(device)
model.eval()  # 切换到推理模式，关闭训练相关的层

# 加载分词器并编码文本
tokenizer = transformers.BertTokenizer.from_pretrained('TurkuNLP/bert-base-finnish-cased-v1')
encoding = tokenizer.encode_plus(
    txt,
    add_special_tokens=True,
    truncation=True,
    padding="max_length",
    return_attention_mask=True,
    return_tensors="pt"
)

# 将所有输入张量移到GPU，和模型保持同一设备
encoding = {k: v.to(device) for k, v in encoding.items()}

# 推理阶段关闭梯度计算，节省显存并加速
with torch.no_grad():
    output = model(**encoding)
    output = output.logits.softmax(dim=-1).detach().cpu().flatten().numpy().tolist()

二、使用pipeline的场景 #

Hugging Face的pipeline可以直接通过参数指定使用的GPU，不需要手动处理输入张量——它会自动帮你完成设备同步。你只需要在创建pipeline时添加device参数：

方式1：直接指定GPU编号（比如cuda:0对应device=0） #

import transformers

txt = "This was nice place"
# 创建pipeline时指定使用第一个GPU（device=0对应cuda:0）
classifier = transformers.pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device=0  # 关键参数：指定使用GPU
)
result = classifier(txt)

方式2：通过device变量兼容自动切换 #

如果你想让代码自动适配CPU/GPU环境，可以先定义device变量，再传入pipeline：

import torch
import transformers

txt = "This was nice place"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# GPU环境下传入设备编号，CPU环境下传入-1
classifier = transformers.pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english",
    device=device.index if device.type == "cuda" else -1
)
result = classifier(txt)

额外优化建议 #

• 推理前务必调用model.eval()：关闭模型的训练模式（比如dropout、batch norm会切换到推理状态），提升推理稳定性和速度
• 使用torch.no_grad()上下文管理器：禁止梯度计算，大幅减少显存占用，加速推理
• 如果你的GPU显存较大，可以尝试批量处理文本，进一步提升推理效率

内容的提问来源于stack exchange，提问作者Mr. Engineer