欢迎访问宙启技术站
智能推送

入门指南:在Python中使用ONNXRuntime进行模型部署

发布时间:2023-12-17 19:33:51

ONNXRuntime是一个用于高性能推理的开源软件,它支持用于机器学习模型的ONNX(开放神经网络交换)格式。使用ONNXRuntime,您可以在Python中轻松地部署和运行训练好的机器学习模型。

要开始使用ONNXRuntime,您需要首先安装它。可以使用Python的包管理工具pip来安装ONNXRuntime。在终端或命令提示符下运行以下命令:

pip install onnxruntime

安装完成后,您可以开始使用ONNXRuntime。

在下面的示例中,我们将演示如何使用ONNXRuntime来部署和运行一个预训练的图像分类模型。

首先,我们需要下载一个预训练的图像分类模型。ONNXRuntime官方网站提供了一些示例模型,可以在 https://github.com/onnx/models 上找到。在这个示例中,我们将使用ResNet50模型。

接下来,我们需要一个输入图像来进行分类。您可以选择任何一张图像,只需确保该图像与模型的输入大小一致。在这个示例中,我们将使用一个名为"input.jpg"的图像。

现在,让我们编写一个Python脚本来部署和运行模型。

import onnxruntime as rt
import numpy as np
from PIL import Image

# 加载模型
model_path = "resnet50.onnx"
sess = rt.InferenceSession(model_path)

# 加载输入图像
image_path = "input.jpg"
image = Image.open(image_path)
image = image.resize((224, 224))  # 将图像调整为模型期望的大小
image = np.array(image).astype(np.float32)  # 将图像转换为浮点数数组
image = np.transpose(image, (2, 0, 1))  # 转置图像的维度顺序

# 为输入创建一个新的numpy数组
input_name = sess.get_inputs()[0].name
inputs = {input_name: image}

# 运行模型
outputs = sess.run(None, inputs)

# 获取输出
output_name = sess.get_outputs()[0].name
output = outputs[0]

# 获取分类结果
labels = np.loadtxt("labels.txt", dtype=str)  # 加载类别标签
predicted_idx = np.argmax(output)  # 获取预测的索引
predicted_label = labels[predicted_idx]  # 获取预测的标签

# 打印结果
print("Predicted label:", predicted_label)

在这个脚本中,我们首先使用ONNXRuntime的InferenceSession类加载了我们的模型。然后,我们使用Python库PIL打开和调整输入图像的大小。接下来,我们将图像转换为Numpy数组,并将其转置以匹配模型的期望输入格式。

我们通过get_inputs()方法获取输入的名称,并创建一个包含输入图像的字典。然后,我们使用run()方法运行模型,并通过get_outputs()方法获取输出的名称。我们将分类结果输出到控制台。

要运行这个脚本,您需要确保已安装所有的依赖项,并将模型文件和输入图像放在与脚本相同的目录中。您还需要一个包含模型的类别标签的文本文件。您可以创建一个名为"labels.txt"的文本文件,每行一个类别标签。

现在,您可以使用ONNXRuntime来轻松地部署和运行机器学习模型。无论是用于图像分类、目标检测还是其他任务,ONNXRuntime都可以在Python中提供高性能的推理引擎。