开始使用 Python 中的 ONNX Runtime

以下是快速指南，帮助您安装使用 ONNX 进行模型序列化和使用 ORT 进行推理的软件包。

目录

• 安装 ONNX Runtime
• 安装 ONNX 以导出模型
• PyTorch、TensorFlow 和 SciKit Learn 的快速入门示例
• Python API 参考文档
• 构建版本
• 了解更多

安装 ONNX Runtime

ONNX Runtime 有两个 Python 软件包。在任何一个环境中，一次只能安装其中一个软件包。GPU 软件包包含大部分 CPU 功能。

安装 ONNX Runtime CPU

如果您在基于 Arm® 的 CPU 和/或 macOS 上运行，请使用 CPU 软件包。

pip install onnxruntime

安装 ONNX Runtime GPU (CUDA 12.x)

ORT 的默认 CUDA 版本是 12.x。

pip install onnxruntime-gpu

安装 ONNX Runtime GPU (CUDA 11.8)

对于 Cuda 11.8，请使用以下说明从 ORT Azure Devops Feed 安装

pip install onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/

安装 ONNX 以导出模型

## ONNX is built into PyTorch
pip install torch

## tensorflow
pip install tf2onnx

## sklearn
pip install skl2onnx

PyTorch、TensorFlow 和 SciKit Learn 的快速入门示例

使用您最喜欢的框架训练模型，导出为 ONNX 格式，并在任何受支持的 ONNX Runtime 语言中进行推理！

PyTorch CV

在本示例中，我们将介绍如何将 PyTorch CV 模型导出为 ONNX 格式，然后使用 ORT 进行推理。创建模型的代码来自 Microsoft Learn 上的 PyTorch 基础知识学习路径。

• 使用 torch.onnx.export 导出模型

torch.onnx.export(model,                                # model being run
                  torch.randn(1, 28, 28).to(device),    # model input (or a tuple for multiple inputs)
                  "fashion_mnist_model.onnx",           # where to save the model (can be a file or file-like object)
                  input_names = ['input'],              # the model's input names
                  output_names = ['output'])            # the model's output names

• 使用 onnx.load 加载 onnx 模型

  import onnx
  onnx_model = onnx.load("fashion_mnist_model.onnx")
  onnx.checker.check_model(onnx_model)
  

• 使用 ort.InferenceSession 创建推理会话

import onnxruntime as ort
import numpy as np
x, y = test_data[0][0], test_data[0][1]
ort_sess = ort.InferenceSession('fashion_mnist_model.onnx')
outputs = ort_sess.run(None, {'input': x.numpy()})

# Print Result
predicted, actual = classes[outputs[0][0].argmax(0)], classes[y]
print(f'Predicted: "{predicted}", Actual: "{actual}"')

PyTorch NLP

在本示例中，我们将介绍如何将 PyTorch NLP 模型导出为 ONNX 格式，然后使用 ORT 进行推理。创建 AG News 模型的代码来自 本 PyTorch 教程。

• 处理文本并创建用于导出的示例数据输入和偏移量。

  import torch
  text = "Text from the news article"
  text = torch.tensor(text_pipeline(text))
  offsets = torch.tensor([0])
  

• 导出模型

  # Export the model
  torch.onnx.export(model,                     # model being run
                  (text, offsets),           # model input (or a tuple for multiple inputs)
                  "ag_news_model.onnx",      # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=10,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input', 'offsets'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})
  

• 使用 onnx.load 加载模型

  import onnx
  onnx_model = onnx.load("ag_news_model.onnx")
  onnx.checker.check_model(onnx_model)
  

• 使用 ort.InferenceSession 创建推理会话

  import onnxruntime as ort
  import numpy as np
  ort_sess = ort.InferenceSession('ag_news_model.onnx')
  outputs = ort_sess.run(None, {'input': text.numpy(),
                              'offsets':  torch.tensor([0]).numpy()})
  # Print Result
  result = outputs[0].argmax(axis=1)+1
  print("This is a %s news" %ag_news_label[result[0]])
  

TensorFlow CV

在本示例中，我们将介绍如何将 TensorFlow CV 模型导出为 ONNX 格式，然后使用 ORT 进行推理。使用的模型来自 Keras resnet50 的 GitHub Notebook。

• 获取预训练模型

import os
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50
import onnxruntime

model = ResNet50(weights='imagenet')

preds = model.predict(x)
print('Keras Predicted:', decode_predictions(preds, top=3)[0])
model.save(os.path.join("/tmp", model.name))

• 将模型转换为 onnx 并导出

import tf2onnx
import onnxruntime as rt

spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = model.name + ".onnx"

model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=13, output_path=output_path)
output_names = [n.name for n in model_proto.graph.output]

• 使用 rt.InferenceSession 创建推理会话

providers = ['CPUExecutionProvider']
m = rt.InferenceSession(output_path, providers=providers)
onnx_pred = m.run(output_names, {"input": x})

print('ONNX Predicted:', decode_predictions(onnx_pred[0], top=3)[0])

SciKit Learn CV

在本示例中，我们将介绍如何将 SciKit Learn CV 模型导出为 ONNX 格式，然后使用 ORT 进行推理。我们将使用著名的 iris 数据集。

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)

from sklearn.linear_model import LogisticRegression
clr = LogisticRegression()
clr.fit(X_train, y_train)
print(clr)

LogisticRegression()

• 将模型转换或导出为 ONNX 格式

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

initial_type = [('float_input', FloatTensorType([None, 4]))]
onx = convert_sklearn(clr, initial_types=initial_type)
with open("logreg_iris.onnx", "wb") as f:
    f.write(onx.SerializeToString())

• 使用 ONNX Runtime 加载和运行模型。我们将使用 ONNX Runtime 来计算此机器学习模型的预测。

import numpy
import onnxruntime as rt

sess = rt.InferenceSession("logreg_iris.onnx")
input_name = sess.get_inputs()[0].name
pred_onx = sess.run(None, {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

OUTPUT:
 [0 1 0 0 1 2 2 0 0 2 1 0 2 2 1 1 2 2 2 0 2 2 1 2 1 1 1 0 2 1 1 1 1 0 1 0 0
  1]

• 获取预测类别

可以更改代码以通过将特定输出的名称指定到列表中来获取该输出。

import numpy
import onnxruntime as rt

sess = rt.InferenceSession("logreg_iris.onnx")
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run(
    [label_name], {input_name: X_test.astype(numpy.float32)})[0]
print(pred_onx)

Python API 参考文档

转到 ORT Python API 文档

构建版本

如果使用 pip，请在下载之前运行 pip install --upgrade pip。

工件	描述	支持的平台
onnxruntime	CPU（发布版）	Windows (x64)、Linux (x64, ARM64)、Mac (X64),
nightly	CPU（开发版）	与上述相同
onnxruntime-gpu	GPU（发布版）	Windows (x64)、Linux (x64, ARM64)
用于 CUDA 11.* 的 onnxruntime-gpu	GPU（开发版）	Windows (x64)、Linux (x64, ARM64)
用于 CUDA 12.* 的 onnxruntime-gpu	GPU（开发版）	Windows (x64)、Linux (x64, ARM64)

安装用于 CUDA 11.* 的 onnxruntime-gpu 的示例

python -m pip install onnxruntime-gpu --extra-index-url=https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-11-nightly/pypi/simple/

安装用于 CUDA 12.* 的 onnxruntime-gpu 的示例

python -m pip install onnxruntime-gpu --pre --extra-index-url=https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/

有关 Python 编译器版本说明，请参阅此页面

了解更多

• Python 教程
• TensorFlow 与 ONNX Runtime
• PyTorch 与 ONNX Runtime
• scikit-learn 与 ONNX Runtime