在 Python 中使用便携式 ONNX AI 模型

Keith Pijanowski

5.00/5 (1投票)

2020年9月14日

CPOL

4分钟阅读

10357

在本文中，我简要概述了 ONNX Runtime 和 ONNX 格式。

下载源代码 - 547.1 KB

在本系列文章中，我们将介绍如何在 2020 年使用可移植的神经网络，您将学习如何在 x64 架构上安装 ONNX 并在 Python 中使用它。

微软与 Facebook 和 AWS 共同开发了 ONNX。 ONNX 格式和 ONNX Runtime 都获得了行业支持，以确保所有重要的框架都能够将其图导出到 ONNX，并且这些模型可以在任何硬件配置上运行。

ONNX Runtime 是一个用于运行已转换为 ONNX 格式的机器学习模型的引擎。传统的机器学习模型和深度学习模型（神经网络）都可以导出为 ONNX 格式。该运行时可以在 Linux、Windows 和 Mac 上运行，并且可以在各种芯片架构上运行。它还可以利用硬件加速器，例如 GPU 和 TPU。但是，并非每种操作系统、芯片架构和加速器的组合都有一个安装包，因此如果您没有使用常见的组合之一，则可能需要从源代码构建运行时。请查看 ONNX Runtime 网站，获取所需组合的安装说明。本文将展示如何在具有默认 CPU 的 x64 架构上以及具有 GPU 的 x64 架构上安装 ONNX Runtime。

除了能够在许多硬件配置上运行之外，还可以从大多数流行的编程语言中调用运行时。本文的目的是展示如何在 Python 中使用 ONNX Runtime。我将展示如何安装 onnxruntime 包。安装 ONNX Runtime 后，我将把之前导出的 MNIST 模型加载到 ONNX Runtime 中，并使用它来进行预测。

安装和导入 ONNX Runtime

在使用 ONNX Runtime 之前，您需要安装 onnxruntime 包。以下命令将在具有默认 CPU 的 x64 架构上安装运行时

pip install onnxruntime

要在具有 GPU 的 x64 架构上安装运行时，请使用以下命令

pip install onnxruntime-gpu

安装后，可以使用以下 import 语句将其导入到您的模块中

import onnxruntime

加载 ONNX 模型

将 ONNX 模型加载到 ONNX Runtime 中与转换一样简单，实际上只需一行代码。下面的函数演示了如何将 ONNX 模型加载到 ONNX Runtime 中。 ONNX 模型的实际加载只是一个命令。

def load_onnx_model(onnx_model_file):
   try:
       session = onnxruntime.InferenceSession(onnx_model_file)

   except (InvalidGraph, TypeError, RuntimeError) as e:
       # It is possible for there to be a mismatch between the onnxruntime and the
       # version of the onnx model format.
       print(e)
       raise e
   return session

使用 ONNX Runtime 进行预测

下面的函数演示了如何使用加载 ONNX 模型时创建的 ONNX 会话。这里有一些值得注意的地方。首先，您需要查询会话以获取其输入。这是使用会话的 get_inputs() 方法完成的。此输入的名称用于创建传递给会话运行命令的输入字典。我们的 MNIST 模型只有一个输入参数，它是一个图像列表。（在 image_samples 参数中。）如果您的模型有多个输入参数，那么 get_inputs() 将为每个参数都有一个条目。

def onnx_infer(onnx_session, image_samples):
   input_name = onnx_session.get_inputs()[0].name
   result = onnx_session.run(None, {input_name: image_samples})
   probabilities = np.array(result[0])
   print(type(probabilities))
   print(probabilities)

   # Generate arg maxes for predictions
   predictions = np.argmax(probabilities, axis=1)
   return predictions

大多数神经网络不直接返回预测。它们返回每个输出类的概率列表。对于我们的 MNIST 模型，每个图像的返回值将是 10 个概率的列表。概率最高的条目就是预测。您可以进行一个有趣的测试，将 ONNX 模型返回的概率与原始模型在创建模型的框架中运行时返回的概率进行比较。理想情况下，模型格式和运行时的更改不应更改产生的任何概率。这将创建一个良好的单元测试，每次模型发生更改时都会运行该测试。

总结和后续步骤

在本文中，我简要概述了 ONNX Runtime 和 ONNX 格式。然后，我展示了如何在 ONNX Runtime 中使用 Python 加载和运行 ONNX 模型。

本文的代码示例包含一个可用的控制台应用程序，演示了此处显示的所有技术。此代码示例是 Github 存储库的一部分，该存储库探索了使用神经网络预测 MNIST 数据集中找到的数字。具体来说，有一些示例展示了如何在 Keras、PyTorch、TensorFlow 1.0 和 TensorFlow 2.0 中创建神经网络。

如果您想了解更多关于导出为 ONNX 格式和使用 ONNX Runtime 的信息，请查看本系列中的其他文章。

在 Python 中使用便携式 ONNX AI 模型

安装和导入 ONNX Runtime

加载 ONNX 模型

使用 ONNX Runtime 进行预测

总结和后续步骤

参考文献