如何在Python中为Transformer模型指定使用CPU？

Question

**如何在Python中为Transformer模型指定使用CPU进行推理或训练？** 在使用如Hugging Face的Transformers库加载预训练模型（如BERT、GPT等）时，默认情况下模型会根据设备情况自动选择运行在GPU或CPU上。然而，在某些场景下，我们希望显式地将模型指定到CPU上运行，例如在没有GPU支持的环境中或进行轻量级测试时。 常见问题是：**如何确保Transformer模型在CPU上运行？** 解决方法通常包括：使用`.to("cpu")`或`.cpu()`方法将模型移动到CPU；设置环境变量如`CUDA_VISIBLE_DEVICES=-1`禁用GPU；或在推理时指定`device="cpu"`参数。此外，还需注意数据张量也应位于CPU上，以避免设备不匹配错误。正确掌握这些方法有助于在CPU环境下高效运行Transformer模型。

Answer 1

一、引言：为何要将Transformer模型运行在CPU上？

在深度学习模型部署和测试阶段，尤其是在资源受限或没有GPU支持的环境中，我们常常需要将模型运行在CPU上。对于Hugging Face Transformers库中的预训练模型（如BERT、GPT等），虽然默认情况下会根据设备情况自动选择计算设备，但显式指定CPU运行可以避免运行时错误、提升调试效率，并确保代码的可移植性。

二、基础知识：设备管理在PyTorch中的工作原理

在PyTorch中，张量和模型都有一个设备属性（device），可以是"cpu"或"cuda"。模型和数据必须位于同一设备上才能进行计算。例如，若模型在GPU上而输入数据在CPU上，会抛出RuntimeError。

因此，控制模型和张量的设备是关键。

常见设备相关方法包括：

• model.to("cpu")：将模型移动到CPU。
• model.cpu()：等价于model.to("cpu")。
• tensor.to("cpu")：将张量移动到CPU。

三、具体操作：如何在Hugging Face Transformers中指定CPU

1. 加载模型并显式移动到CPU

from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("model_path").to("cpu")
  

2. 使用环境变量禁用GPU

设置环境变量CUDA_VISIBLE_DEVICES=-1可全局禁用所有GPU设备，强制所有计算在CPU上进行。

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("model_path")
  

3. 在推理时指定设备

某些模型或库支持在调用时指定设备参数：

from transformers import pipeline

nlp = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english", device="cpu")
result = nlp("I love using Transformers!")
  

四、进阶技巧与注意事项

1. 数据张量也需在CPU上

即使模型在CPU上，如果输入数据仍在GPU上，仍会引发错误。应确保输入数据也使用.to("cpu")或直接在CPU上创建。

inputs = tokenizer("Hello world!", return_tensors="pt").to("cpu")
  

2. 检查设备状态

可通过以下方式检查模型和张量所在的设备：

print(model.device)  # 输出: device(type='cpu')
print(inputs["input_ids"].device)  # 输出: device(type='cpu')
  

3. 多设备推理时的注意事项

如果模型使用了DataParallel或多GPU并行，需先移除并行设置再移动到CPU。

if isinstance(model, torch.nn.DataParallel):
    model = model.module
model = model.to("cpu")
  

五、流程图：CPU运行模型的典型流程

graph TD
    A[加载模型] --> B{是否需要在CPU运行?}
    B -->|是| C[设置CUDA_VISIBLE_DEVICES=-1 或 .to("cpu")]
    B -->|否| D[使用GPU]
    C --> E[加载Tokenizer并编码输入]
    E --> F[将输入张量移至CPU]
    F --> G[执行推理或训练]

六、总结

在使用Hugging Face Transformers库时，通过显式地将模型和数据移动到CPU，或通过环境变量控制设备选择，可以有效确保模型在CPU环境下运行。理解设备管理机制、注意张量与模型的同步设备状态，是实现CPU推理和训练的关键。