AWS Trainium & Inferentia 文档
使用 AWS Neuron (Inf2/Trn1) 进行推理管线
并获得增强的文档体验
开始使用
使用 AWS Neuron (Inf2/Trn1) 进行推理管线
pipeline()
函数使得使用 模型中心 的模型在各种任务(如文本分类、问答和图像分类)上进行加速推理变得简单。
您还可以使用 Transformers 的 pipeline() 函数,并提供您的 NeuronModel 模型类。
目前支持的任务有
特征提取
掩码填充
文本分类
令牌分类
问答
零样本分类
Optimum pipeline 用法
虽然每个任务都有一个关联的 pipeline 类,但使用通用的 pipeline()
函数更简单,它将所有特定于任务的 pipeline 包装在一个对象中。pipeline()
函数会自动加载一个默认模型和 tokenizer/feature-extractor,它们能够为您的任务执行推理。
- 首先通过指定推理任务来创建一个 pipeline
>>> from optimum.neuron.pipelines import pipeline
>>> classifier = pipeline(task="text-classification")
- 将您的输入文本/图像传递给
pipeline()
函数
>>> classifier("I like you. I love you.")
[{'label': 'POSITIVE', 'score': 0.9998838901519775}]
注意: pipeline()
函数中使用的默认模型未针对推理进行优化或量化,因此与 PyTorch 对应模型相比,性能不会有提升。
使用原始 Transformers 模型并转换为 AWS Neuron
pipeline()
函数接受来自 Hugging Face Hub 的任何受支持模型。模型中心上有标签,允许您筛选出您想用于任务的模型。
为了能够使用 Neuron Runtime 加载模型,所考虑的架构需要支持导出到 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)