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

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 模型运行它。

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