独家 | 使用ONNX搭建NLP Transformers pipelines

2022-10-09 10:12:33 浏览数 (1)

代码语言：javascript复制
作者：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文件中。要安装它们，请运行以下命令。

代码语言：javascript复制
$ 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是我们要解决的任务。你可以通过执行以下代码看到可用的特征列表：

代码语言：javascript复制
from transformers.onnx.features import FeaturesManagerdistilbert_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格式的命令：

代码语言：javascript复制
$ python -m transformers.onnx     --model=dslim/bert-base-NER     --feature=token-classification     onnx/

通过Transformers pipeline来使用ONNX模型

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

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

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

代码语言：javascript复制
import torchfrom onnxruntime import (InferenceSession, SessionOptions, GraphOptimizationLevel)from transformers import (TokenClassificationPipeline, AutoTokenizer, AutoModelForTokenClassification)

创建一个ONNX模型会话：

代码语言：javascript复制
options = SessionOptions() # initialize session optionsoptions.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALLsession = InferenceSession("onnx/model.onnx", sess_options=options, providers=["CPUExecutionProvider"])# disable session.run() fallback mechanism, it prevents for a reset of the execution providersession.disable_fallback()

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

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

代码语言：javascript复制
from onnxruntime import get_all_providersget_all_providers()

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

使用ONNX模型创建pipeline

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

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

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

代码语言：javascript复制
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 usedby 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 createdbefore."""# This comes from the original implementation of the pipelinespecial_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 arraysoutputs_name = session.get_outputs()[0].name # get the name of the output tensorlogits = session.run(output_names=[outputs_name], input_feed=inputs)[0] # run the sessionlogits = torch.tensor(logits) # convert to torch tensor to be compatible with the original implementationreturn {"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 Falsemodel_inputs = self.tokenizer(sentence,return_attention_mask=True, # This is the only difference from the original implementationreturn_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_mappingmodel_inputs["sentence"] = sentencereturn model_inputs

运行pipeline

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

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

代码语言：javascript复制
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并检查输出。

代码语言：javascript复制
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模型，并用它创建一个管道：

代码语言：javascript复制
pytorch_pipeline = TokenClassificationPipeline(task="ner",model=model,tokenizer=tokenizer,framework="pt",aggregation_strategy="simple",)

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

代码语言：javascript复制
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次，以获得更准确的基准，并将所有内容放在一个表格中，以比较结果：

代码语言：javascript复制
import timeitfrom tabulate import tabulateresults = [["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模型的推理时间之比。

代码语言：javascript复制
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

编辑：黄继彦

译者简介

欧阳锦，一名在埃因霍温理工大学就读的硕士生。喜欢数据科学和人工智能相关方向。欢迎不同观点和想法的交流与碰撞，对未知充满好奇，对热爱充满坚持。

翻译组招募信息

工作内容：需要一颗细致的心，将选取好的外文文章翻译成流畅的中文。如果你是数据科学/统计学/计算机类的留学生，或在海外从事相关工作，或对自己外语水平有信心的朋友欢迎加入翻译小组。

你能得到：定期的翻译培训提高志愿者的翻译水平，提高对于数据科学前沿的认知，海外的朋友可以和国内技术应用发展保持联系，THU数据派产学研的背景为志愿者带来好的发展机遇。

其他福利：来自于名企的数据科学工作者，北大清华以及海外等名校学生他们都将成为你在翻译小组的伙伴。

点击文末“阅读原文”加入数据派团队~

转载须知

如需转载，请在开篇显著位置注明作者和出处（转自：数据派ID：DatapiTHU），并在文章结尾放置数据派醒目二维码。有原创标识文章，请发送【文章名称-待授权公众号名称及ID】至联系邮箱，申请白名单授权并按要求编辑。

发布后请将链接反馈至联系邮箱（见下方）。未经许可的转载以及改编者，我们将依法追究其法律责任。

点击“阅读原文”拥抱组织

pytorch 腾讯云测试服务 NLP服务 github git

0 人点赞