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

2024-03-29 10:56:41 发布

您现在位置:Python中文网/ 问答频道 /正文

我有一个在Pytorch上训练的NLP模型,将在Jetson Xavier上运行。我安装了Jetson stats来监控CPU和GPU的使用情况。当我运行Python脚本时,只有CPU内核可以加载,GPU条不会增加。我在谷歌上搜索过“如何检查pytorch是否在使用GPU?”等关键词,并在stackoverflow.com等网站上查看了结果。根据他们对其他面临类似问题的人的建议,cuda是可用的,我的Jetson Xavier中有cuda设备。然而,我不明白为什么GPU条没有改变,CPU核心条走到了尽头

我不想使用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

Tags: from模型importmodelgpudircpucuda
2条回答

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

相关问题 更多 >