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

问题描述

我有一个经过Pytorch训练的NLP模型，可以在Jetson Xavier中运行。我安装了Jetson统计信息来监视cpu和GPU的使用情况。当我运行Python脚本时，只有cpu内核可以在负载下工作，GPU条不会增加。我在Google上搜索了关键字“如何检查pytorch是否正在使用GPU？”的关键字。并在stackoverflow.com等网站上检查结果。根据他们对面临类似问题的其他人的建议，cuda可用，并且我的Jetson Xavier中有cuda设备。但是，我不明白为什么GPU栏不会改变，cpu核心栏会消失。

我不想使用cpu，计算时间很长。我认为它使用cpu，而不是GPU。如何确定它是否使用cpu，如何将其更改为GPU？

注意：该模型取自“ havingface”变形金刚库。我试图在模型上使用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

解决方法

要使模型在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

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

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

huggingface-transformers nvidia-jetson python pytorch