作者：Thomas Chaigneau
翻译：欧阳锦校对：和中华
本文约3000字，建议阅读7分钟本文介绍了如何使用ONNX构建真实世界的NLP应用。

如何用ONNX构建真实世界的NLP应用，而不仅仅是为了张量做基准测试。

图片源自网络

ONNX是一种用于神经网络的机器学习格式。它是可移植的，开源的，并且在不牺牲准确性的情况下提高推理速度，真的很厉害。

我发现了很多关于ONNX基准的文章，但没有一篇文章介绍将其用于真实世界NLP任务的简便方法。我还在Hugging Face的discord server上回答了很多关于ONNX以及将其用于NLP的最佳方式的问题。

这就是我决定写这篇博文的原因。我想帮助你使用ONNX与超强的Transformers pipelines获得尽可能好的结果。

本教程将告诉你如何将Hugging Face的NLP Transformers模型导出到ONNX，以及如何将导出的模型与适当的Transformers pipeline一起使用。我使用命名实体识别(NER)模型作为例子，但它并不限于NER。

所有的代码片段都可以在相关的GitHubrepo中找到。所以不用担心复制的问题，只需克隆仓库并在阅读这篇博文时运行notebook。

运行环境

首先，你需要安装所需的依赖项。建议使用一个隔离的环境以避免冲突。

该项目需要Python 3.8或更高版本。你可以使用任何你想要的软件包管理器。我推荐在本教程中使用conda。所有需要的依赖项都列在requirements.txt文件中。要安装它们，请运行以下命令。

$ conda create -y -n hf-onnx python=3.8$ conda activate hf-onnx
$ git clone https://github.com/ChainYo/transformers-pipeline-onnx.git$ cd transformers-pipeline-onnx
$ pip install -r requirements.txt

导出ONNX模型

对于这个例子，我们可以使用Hugging Face库中的任何TokenClassification模型，因为我们要解决的任务是NER。

我选择了dslim/bert-base-NER模型，因为它是一个基础模型，意味着在CPU上的计算时间适中。另外，BERT架构是NER的一个不错的选择。

Huggging Faces的Transformers库提供了一个方便的方法来将模型导出为ONNX格式。你可以参考官方文档了解更多细节。

我们使用上面提到的bert-base-NER模型和token-classification作为特征。token-classification是我们要解决的任务。你可以通过执行以下代码看到可用的特征列表：

from transformers.onnx.features import FeaturesManager
distilbert_features = list(FeaturesManager.get_supported_features_for_model_type("bert").keys())
print(distilbert_features)
>>> ['default', 'masked-lm', 'causal-lm', 'sequence-classification', 'token-classification', 'question-answering']

通过调用转换脚本，你必须指定模型名称，从本地目录或直接从Hugging Face的枢纽中指定。你还需要指定如上所示的特征。输出文件将被保存在output目录中。

我们把onnx/作为输出目录。这就是ONNX模型将被保存的地方。

我们让opset参数为默认值，该参数在模型的ONNX配置中定义。

最后，我们还将atol参数设为默认值，即1e-05。这是原始PyTorch模型和ONNX模型之间数值精度的公差。

下面是将模型导出为ONNX格式的命令：

$ python -m transformers.onnx \    --model=dslim/bert-base-NER \    --feature=token-classification \    onnx/

通过Transformers pipeline来使用ONNX模型

现在我们已经将模型导出为ONNX格式，我们可以用Transformers pipeline来使用它，这个过程很简单。

• 用ONNX模型创建一个会话，允许你将模型加载到管道中并进行推理。
• 覆盖管道的_forward和preprocess方法以使用ONNX模型。
• 运行管道。

首先，让我们导入所需的包。

import torch
from onnxruntime import (
InferenceSession, SessionOptions, GraphOptimizationLevel
)
from transformers import (
TokenClassificationPipeline, AutoTokenizer, AutoModelForTokenClassification
)

创建一个ONNX模型会话：

options = SessionOptions() # initialize session options
options.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
session = InferenceSession(
"onnx/model.onnx", sess_options=options, providers=["CPUExecutionProvider"]
)
# disable session.run() fallback mechanism, it prevents for a reset of the execution provider
session.disable_fallback()

这里我们将只使用CPUExecutionProvider，它是ONNX模型的默认执行提供者（execution provider）。你可以为会话提供一个或多个执行提供者。例如，你可以使用CUDAExecutionProvider来在GPU上运行模型。默认情况下，会话将从列表中的第一个开始，使用机器上可用的那个。

Onnxruntime提供了一个函数来查看所有可用的执行提供者。

from onnxruntime import get_all_providers
get_all_providers()

正如你所看到的，有很多可用于各种用例和配置的provider。

使用ONNX模型创建pipeline

现在我们有一个带有ONNX模型的会话可以使用，我们可以继承原来的TokenClassificationPipeline类以使用ONNX模型。

为了充分理解原理，你可以参考TokenClassificationPipeline python类的源代码。

我们将只覆盖_forward和preprocess方法，因为其他方法不依赖于模型格式。

class OnnxTokenClassificationPipeline(TokenClassificationPipeline):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def _forward(self, model_inputs):
"""
Forward pass through the model. This method is not to be called by the user directly and is only used
by the pipeline to perform the actual predictions.
This is where we will define the actual process to do inference with the ONNX model and the session created
before.
"""
# This comes from the original implementation of the pipeline
special_tokens_mask = model_inputs.pop("special_tokens_mask")
offset_mapping = model_inputs.pop("offset_mapping", None)
sentence = model_inputs.pop("sentence")
inputs = {k: v.cpu().detach().numpy() for k, v in model_inputs.items()} # dict of numpy arrays
outputs_name = session.get_outputs()[0].name # get the name of the output tensor
logits = session.run(output_names=[outputs_name], input_feed=inputs)[0] # run the session
logits = torch.tensor(logits) # convert to torch tensor to be compatible with the original implementation
return {
"logits": logits,
"special_tokens_mask": special_tokens_mask,
"offset_mapping": offset_mapping,
"sentence": sentence,
**model_inputs,
}
# We need to override the preprocess method because the onnx model is waiting for the attention masks as inputs
# along with the embeddings.
def preprocess(self, sentence, offset_mapping=None):
truncation = True if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0 else False
model_inputs = self.tokenizer(
sentence,
return_attention_mask=True, # This is the only difference from the original implementation
return_tensors=self.framework,
truncation=truncation,
return_special_tokens_mask=True,
return_offsets_mapping=self.tokenizer.is_fast,
)
if offset_mapping:
model_inputs["offset_mapping"] = offset_mapping
model_inputs["sentence"] = sentence
return model_inputs

运行pipeline

已经设置好了一切，现在我们可以运行管道。

像往常一样，管道需要一个分词器、一个模型和一个任务。我们将使用ner任务。

model_name_from_hub = "dslim/bert-base-NER"
tokenizer = AutoTokenizer.from_pretrained(model_name_from_hub)
model = AutoModelForTokenClassification.from_pretrained(model_name_from_hub)
onnx_pipeline = OnnxTokenClassificationPipeline(
task="ner",
model=model,
tokenizer=tokenizer,
framework="pt",
aggregation_strategy="simple",
)

看看我们是否可以运行pipeline并检查输出。

sequence = "Apple was founded in 1976 by Steve Jobs, Steve Wozniak and Ronald Wayne to develop and sell Wozniak's Apple I personal computer"
onnx_pipeline(sequence)

ONNX pipeline输出

在这里，管道在ONNX模型下运行良好! 现在，我们有了一条完全可以使用ONNX的NER管道。

可以看看下面可选读的基准测试章节，看看与原始的PyTorch模型相比，该模型的表现如何，或者直接跳到最末尾结论部分，对该过程进行快速总结。

对完整pipeline进行基准测试（可选读）

我们将对ONNX模型和PyTorch模型的推理时间做基准测试。

我们首先需要加载PyTorch模型，并用它创建一个管道：

pytorch_pipeline = TokenClassificationPipeline(
task="ner",
model=model,
tokenizer=tokenizer,
framework="pt",
aggregation_strategy="simple",
)

我们将用相同的数据和3个不同的序列长度来测试两个管道：

sequences = {
"short_sequence": "Hello my name is Thomas and I love HuggingFace.",
"medium_sequence": "Winston Churchill was born in 1874 in Stoke-on-Trent, England, to a German father, William and Elizabeth Churchill.",
"long_sequence": """The first person to reach the summit of Everest was the South Nepalese Everest Gurun, who was a member of the Royal Nepal Expedition, led by the Nepalese Mountaineer, Sir Edmund Hillary. Hilary lived in the Himalayas for a time. He sadly died in 1953 at the age of 88."""
}

让我们比较一下每个管道在3个不同序列长度下的推理时间。我们将对每个序列的长度重复300次，以获得更准确的基准，并将所有内容放在一个表格中，以比较结果：

import timeit
from tabulate import tabulate
results = [["Sequence Length", "PyTorch", "ONNX"]]
for k, v in sequences.items():
results.append( [k, timeit.timeit(lambda: pytorch_pipeline(v), number=300), timeit.timeit(lambda: onnx_pipeline(v), number=300)])
print(tabulate(results, headers="firstrow"))

基准测试结果

看起来很不错! 看来，对于每个序列长度，ONNX模型都比原来的PyTorch模型快得多。让我们计算一下ONNX模型和PyTorch模型的推理时间之比。

print(f"For a short sequence: ONNX is {results[1][1]/results[1][2]:.2f}x faster than PyTorch")
print(f"For a medium sequence: ONNX is {results[2][1]/results[2][2]:.2f}x faster than PyTorch")
print(f"For a long sequence: ONNX is {results[3][1]/results[3][2]:.2f}x faster than PyTorch")

基准测试比例

在长序列上，我们几乎实现了3倍的速度提升！我们甚至没有根据模型结构和模型运行的硬件做任何优化，而这是可以用ONNX做的。

优化可能非常有用，但这是一个很深的话题，在这篇文章中无法涵盖。但知道你能做到这一点是很好的，我们可以在未来的文章中探讨它。

另外，我们的测试是在CPU上进行的，但我看到的所有GPU上的基准测试都比CPU上的基准测试更令人印象深刻。查看这篇（https://towardsdatascience.com/nlp-transformers-pipelines-with-onnx-9b890d015723）很不错的伟大的文章，了解更多关于不同架构和推理配置的基准测试。

结论

综上所述，我们已经用ONNX建立了一个完全正常的NER管道。我们将PyTorch模型转换为ONNX模型，并对原有的pipeline类进行了继承以适应ONNX模型的新要求。最后，我们将ONNX模型与原来的PyTorch模型进行了基准测试，并比较了结果。

不幸的是，PyTorch模型必须与ONNX模型一起加载。这是因为Transformers pipeline需要加载PyTorch模型，特别是对模型的配置。

我正在寻找一种方法来避免这种PyTorch模型的加载，因为它可能在某些系统上产生RAM问题。

我们用来使其工作的过程对于Hugging Face的Transformers库中的每个模型和任务都是一样的。

你唯一需要关注的是，模型架构是否有为ONNX实现的配置。你可以在文档中看到完整的架构列表。

如果你要找的架构还没有实现，你仍然可以创建它，并向Transformers库进行pull request以添加它。这正是我几个月前为CamemBERT架构所做的事情。你可以在Transformers的GitHub repo中查看完整的PR。

GitHub repo链接：https://github.com/huggingface/transformers/pull/14059

我希望你觉得这篇文章有用，有趣。如果你有任何问题或面临任何问题，请告诉我。我很想增加更多的例子和对其他NLP任务的支持，所以如果你有任何想法或要求，请告诉我!

如有疑问或问题，请在GitHub上或在下面的评论中打开一个问题。

P.S. 我还计划增加另一个基准部分，以测试ONNX模型是否能达到与原始模型相同的结果。

感谢Katherine Prairie和Ben Huberman。

原文标题：

NLP Transformers pipelines with ONNX

原文链接：

https://towardsdatascience.com/nlp-transformers-pipelines-with-onnx-9b890d015723