Transformers：如何使用 CUDA 进行推理？

Question

我已经使用 GPU 微调了我的模型，但推理过程非常慢，我认为这是因为推理默认使用 CPU。这是我的推理代码：

txt = "This was nice place"
model = transformers.BertForSequenceClassification.from_pretrained(model_path, num_labels=24)
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")
output = model(**encoding)
output = output.logits.softmax(dim=-1).detach().cpu().flatten().numpy().tolist()

这是我的第二个推理代码，它使用管道（针对不同的模型）：

classifier = transformers.pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
result = classifier(txt)

如何强制 Transformer 库在 GPU 上进行更快的推理？我尝试添加

model.to(torch.device("cuda"))

但这会引发错误：

Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu

我认为问题与数据未发送到 GPU 有关。这里有一个类似的问题：pytorch 摘要因 Huggingface 模型 II 失败：预期所有张量都在同一设备上，但发现至少两个设备，cuda：0 和 cpu

使用和不使用管道时如何将数据发送到 GPU？非常感谢任何建议。

Answer 1

在执行推理之前您也应该将输入传输到 CUDA：

device = torch.device('cuda')

# transfer model
model.to(device)

# define input and transfer to device
encoding = tokenizer.encode_plus(txt, 
     add_special_tokens=True, 
     truncation=True, 
     padding="max_length", 
     return_attention_mask=True, 
     return_tensors="pt")

encoding = encoding.to(device)

# inference
output = model(**encoding)

请注意，

nn.Module.to

已就位，而torch.Tensor.to则不在（它会复制！）。

Answer 2

对于管道代码问题

问题在于

transformers.pipeline

 使用 CPU 的默认行为。但是，例如，从 

here 您可以添加 device=0

 参数以使用第一个 GPU。

• device=0

   利用 GPU cuda:0

• device=1

   利用 GPU cuda：1

pipeline = pipeline(TASK, model=MODEL_PATH, device=0) 您的代码变为：

classifier = transformers.pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english", device=0)
result = classifier(txt)