zl程序教程

您现在的位置是:首页 >  其他

当前栏目

使用预训练语言模型预测阶段:GPU、CPU性能差别【Pegasus】

训练性能语言CPU 模型 预测 阶段 GPU
2023-09-27 14:20:38 时间

一、Pegasus

1、使用CPU(用时: 17.92682433128357 秒)

# https://github.com/huggingface/transformers/blob/master/src/transformers/models/pegasus/modeling_pegasus.py
import time

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

device = torch.device("cuda")

tokenizer = AutoTokenizer.from_pretrained(