AWS Trainium & Inferentia 文档

使用 AWS Neuron (Inf2/Trn1) 的推理管道

Hugging Face's logo
加入 Hugging Face 社区

并获取增强文档体验

开始使用

使用 AWS Neuron (Inf2/Trn1) 的推理管道

pipeline() 函数简化了使用来自 模型中心 的模型来执行加速推理以完成各种任务(如文本分类、问答和图像分类)的操作。

您也可以使用 Transformers 中的 pipeline() 函数并提供您的 NeurModel 模型类。

目前支持的任务是:

  • 特征提取
  • 填充掩码
  • 文本分类
  • 令牌分类
  • 问答
  • 零样本分类

Optimum pipeline 使用

虽然每个任务都有一个关联的 pipeline 类,但使用通用的 pipeline() 函数更简单,该函数将所有特定于任务的 pipeline 包含在一个对象中。pipeline() 函数会自动加载一个默认模型和分词器/特征提取器,它们能够为您执行任务进行推理。

  1. 首先通过指定一个推理任务来创建一个 pipeline
>>> from optimum.neuron.pipelines import pipeline

>>> classifier = pipeline(task="text-classification")
  1. 将您的输入文本/图像传递给 pipeline() 函数
>>> classifier("I like you. I love you.")
[{'label': 'POSITIVE', 'score': 0.9998838901519775}]

注意:pipeline() 函数中使用的默认模型没有针对推理进行优化或量化,因此与它们的 PyTorch 等效模型相比,不会有性能提升。

使用普通的 Transformers 模型并将其转换为 AWS Neuron

pipeline() 函数接受来自 Hugging Face Hub 的任何支持的模型。模型中心上有标签,可以让您筛选您想用于任务的模型。

要能够使用 Neuron 运行时加载模型,需要支持将所考虑的架构导出到 neuron。

您可以查看 这里 的支持的架构列表。

选择合适的模型后,您可以通过指定模型仓库来创建 pipeline()

>>> from optimum.neuron.pipelines import pipeline

# The model will be loaded to an NeuronModelForQuestionAnswering.
>>> neuron_qa = pipeline("question-answering", model="deepset/roberta-base-squad2", export=True)
>>> question = "What's my name?"
>>> context = "My name is Philipp and I live in Nuremberg."

>>> pred = neuron_qa(question=question, context=context)

也可以使用与 NeuronModelForXXX 类关联的 from_pretrained(model_name_or_path, export=True) 方法来加载它。

例如,以下是加载 ~neuron.NeuronModelForQuestionAnswering 类以进行问答的方法:

>>> from transformers import AutoTokenizer
>>> from optimum.neuron import NeuronModelForQuestionAnswering, pipeline

>>> tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")

>>> # Loading the PyTorch checkpoint and converting to the neuron format by providing export=True
>>> model = NeuronModelForQuestionAnswering.from_pretrained(
...     "deepset/roberta-base-squad2",
...     export=True
... )

>>> neuron_qa = pipeline("question-answering", model=model, tokenizer=tokenizer)
>>> question = "What's my name?"
>>> context = "My name is Philipp and I live in Nuremberg."

>>> pred = neuron_qa(question=question, context=context)

定义输入形状

NeuronModels 目前需要静态 input_shapes 来运行推理。 如果在提供 export=True 参数时未提供输入形状,将使用默认输入形状。 以下是如何为序列长度和批次大小指定输入形状的示例。

>>> from optimum.neuron.pipelines import pipeline

>>> input_shapes = {"batch_size": 1, "sequence_length": 64} 
>>> clt = pipeline("token-classification", model="dslim/bert-base-NER", export=True,input_shapes=input_shapes)
>>> context = "My name is Philipp and I live in Nuremberg."

>>> pred = clt(context)