Pytorch NLP 模型在进行推理时不使用 GPU

Question

我有一个在 Pytorch 上训练的 NLP 模型在 Jetson Xavier 上运行。我安装了 Jetson stats 来监控 CPU 和 GPU 的使用情况。当我运行 Python 脚本时，只有 CPU 个核心在负载上工作，GPU 条不增加。我在 Google 上搜索了关键字“How to check if pytorch is using the GPU?”并检查了 whosebug.com 等的结果。根据他们对面临类似问题的其他人的建议，cuda 可用并且我的 Jetson Xavier 中有 cuda 设备。但是，我不明白为什么GPU bar没有变化，CPU core bars走到尽头。

我不想用CPU，计算时间这么长。在我看来，它使用 CPU，而不是 GPU。我如何确定，如果它使用 CPU，我如何将其更改为 GPU？

注：模型取自huggingface transformers库。我试图在模型上使用 cuda() 方法。 (model.cuda()) 在这种情况下，使用了 GPU，但我无法从模型中获取输出并引发异常。

代码如下：

from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
import torch

BERT_DIR = "savasy/bert-base-turkish-squad"    

tokenizer = AutoTokenizer.from_pretrained(BERT_DIR)
model = AutoModelForQuestionAnswering.from_pretrained(BERT_DIR)
nlp=pipeline("question-answering", model=model, tokenizer=tokenizer)


def infer(question,corpus):
    try:
        ans = nlp(question=question, context=corpus)
        return ans["answer"], ans["score"]
    except:
        ans = None
        pass

    return None, 0

Answer 1

要使模型在 GPU 上运行，必须将数据和模型加载到 GPU：

您可以按如下方式进行：

from transformers import AutoTokenizer, AutoModelForQuestionAnswering, pipeline
import torch

BERT_DIR = "savasy/bert-base-turkish-squad"  
  
device = torch.device("cuda")

tokenizer = AutoTokenizer.from_pretrained(BERT_DIR)
model = AutoModelForQuestionAnswering.from_pretrained(BERT_DIR)
model.to(device) ## model to GPU

nlp=pipeline("question-answering", model=model, tokenizer=tokenizer)


def infer(question,corpus):
    try:
        ans = nlp(question=question.to(device), context=corpus.to(device)) ## data to GPU
        return ans["answer"], ans["score"]
    except:
        ans = None
        pass

    return None, 0

Answer 2

问题已通过加载包含设备参数的管道解决：

nlp = pipeline("question-answering", model=BERT_DIR, device=0)

Pytorch NLP 模型在进行推理时不使用 GPU

Pytorch NLP model doesn’t use GPU when making inference

python

pytorch

nvidia-jetson

huggingface-transformers