18---使用 ONNX 导出和运行模型

浅浅的花香味﹌ 2023-09-27 19:22 15阅读 0赞

ONNX运行时为机器学习模型提供了一种通用的序列化格式。ONNX 支持多种不同的平台/语言，并具有有助于减少推理时间的内置功能。

PyTorch 对将 Torch 模型导出到 ONNX 具有强大的支持。这使得可以将 Hugging Face Transformer 和/或其他下游模型直接导出到 ONNX。

ONNX 开辟了一条使用多种语言和平台进行直接推理的途径。例如，模型可以直接在 Android 上运行以限制发送到第三方服务的数据。ONNX 是一个令人兴奋的发展，充满希望。微软还发布了Hummingbird，它可以将传统模型（sklearn、决策树、逻辑回归……）导出到 ONNX。

本文将介绍如何使用 txtai 将模型导出到 ONNX。然后，这些模型将直接在 Python、JavaScript、Java 和 Rust 中运行。目前，txtai 通过它的 API 支持所有这些语言，这仍然是推荐的方法。

安装依赖  
安装txtai和所有依赖项。由于本文使用ONNX量化，我们需要安装pipeline extras包。  
pip install txtai\[pipeline\] datasets  
使用 ONNX 运行模型  
让我们开始吧！以下示例将情绪分析模型导出到 ONNX 并运行推理会话。  
import numpy as np

from onnxruntime import InferenceSession, SessionOptions  
from transformers import AutoTokenizer  
from txtai.pipeline import HFOnnx

## Normalize logits using sigmoid function ##

sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))

## Export to ONNX ##

onnx = HFOnnx()  
model = onnx(“distilbert-base-uncased-finetuned-sst-2-english”, “text-classification”)

## Start inference session ##

options = SessionOptions()  
session = InferenceSession(model, options)

## Tokenize ##

tokenizer = AutoTokenizer.from\_pretrained(“distilbert-base-uncased-finetuned-sst-2-english”)  
tokens = tokenizer(\[“I am happy”, “I am mad”\], return\_tensors=“np”)

## Print results ##

outputs = session.run(None, dict(tokens))  
print(sigmoid(outputs\[0\]))  
\[\[0.01295124 0.9909526 \]  
\[0.9874723 0.0297817 \]\]  
就这样，有结果！文本分类模型使用两个标签来判断情绪，0 表示负面，1 表示正面。上面的结果显示了每个文本片段的每个标签的概率。

ONNX 管道加载模型，将图形转换为 ONNX 并返回。请注意，没有提供输出文件，在这种情况下，ONNX 模型作为字节数组返回。如果提供了输出文件，则此方法返回输出路径。

训练和导出文本分类模型  
接下来，我们将 ONNX 管道与 Trainer 管道相结合，以创建“训练并导出到 ONNX”工作流程。  
from datasets import load\_dataset  
from txtai.pipeline import HFTrainer

trainer = HFTrainer()

## Hugging Face dataset ##

ds = load\_dataset(“glue”, “sst2”)  
data = ds\[“train”\].select(range(10000)).flatten\_indices()

## Train new model using 10,000 SST2 records (in-memory) ##

model, tokenizer = trainer(“google/electra-base-discriminator”, data, columns=(“sentence”, “label”))

## Export model trained in-memory to ONNX (still in-memory) ##

output = onnx((model, tokenizer), “text-classification”, quantize=True)

## Start inference session ##

options = SessionOptions()  
session = InferenceSession(output, options)

## Tokenize ##

tokens = tokenizer(\[“I am happy”, “I am mad”\], return\_tensors=“np”)

## Print results ##

outputs = session.run(None, dict(tokens))  
print(sigmoid(outputs\[0\]))  
\[\[0.01525715 0.975399 \]  
\[0.97395283 0.04432926\]\]  
结果与上一步类似，尽管此模型仅在 sst2 数据集的一小部分上进行训练。让我们保存此模型以备后用。  
onnx = HFOnnx()  
text = onnx((model, tokenizer), “text-classification”, “text-classify.onnx”, quantize=True)  
导出句子嵌入模型  
ONNX 管道还支持导出使用句子转换器包训练的句子嵌入模型。  
embeddings = onnx(“sentence-transformers/paraphrase-MiniLM-L6-v2”, “pooling”, “embeddings.onnx”, quantize=True)  
现在让我们使用 ONNX 运行模型。  
from sklearn.metrics.pairwise import cosine\_similarity

options = SessionOptions()  
session = InferenceSession(embeddings, options)

tokens = tokenizer(\[“I am happy”, “I am glad”\], return\_tensors=“np”)

outputs = session.run(None, dict(tokens))\[0\]

print(cosine\_similarity(outputs))  
\[\[0.99999994 0.8474637 \]  
\[0.8474637 0.9999997 \]\]  
上面的代码标记了两个单独的文本片段（“我很高兴”和“我很高兴”）并通过 ONNX 模型运行它。

这将输出两个嵌入数组，并使用余弦相似度对这些数组进行比较。正如我们所见，这两个文本片段具有接近的语