独家 | 使用ONNX搭建NLP Transformers pipelines

作者：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配置中定义。