简介

ONNXMLTools 可以帮助您将来自不同机器学习工具包的模型转换为 ONNX 格式。目前支持以下工具包：

• TensorFlow（基于 tf2onnx 转换器）
• scikit-learn（基于 skl2onnx 转换器）
• Apple Core ML
• Spark ML（实验性支持）
• LightGBM
• libsvm
• XGBoost
• H2O
• CatBoost

PyTorch 自带 ONNX 导出功能，详情请参阅 这里。

安装

您可以从 PyPi 安装最新版本的 ONNXMLTools：

pip install onnxmltools

或者直接从源代码安装：

pip install git+https://github.com/onnx/onnxmltools

如果您选择从源代码安装 onnxmltools，则必须在安装 onnx 包之前设置环境变量 ONNX_ML=1。

依赖项

本包依赖于 ONNX、NumPy 和 ProtoBuf。如果您要转换 scikit-learn、Core ML、Keras、LightGBM、SparkML、XGBoost、H2O、CatBoost 或 LibSVM 的模型，则需要确保环境中已安装相应的库：

1. scikit-learn
2. CoreMLTools（版本 3.1 或更低）
3. Keras（版本 2.0.8 或更高）及对应的 TensorFlow 版本
4. LightGBM
5. SparkML
6. XGBoost
7. libsvm
8. H2O
9. CatBoost

ONNXMLTools 已在 Python 3.7+ 上进行测试。

示例

如果希望转换后的 ONNX 模型与特定的 ONNX 版本兼容，请在调用转换函数时指定 target_opset 参数。下面的 Keras 模型转换示例展示了这一点。您可以在 版本控制文档 中找到 ONNX 操作集（称为 opsets）与 ONNX 发布版本之间的对应关系。

Keras 到 ONNX 的转换

接下来，我们展示一个将 Keras 模型转换为 ONNX 模型的例子，其中 target_opset=7，对应 ONNX 发布版本 1.2。

import onnxmltools
from keras.layers import Input, Dense, Add
from keras.models import Model

# N：批次大小，C：子模型输入维度，D：最终模型输入维度
N, C, D = 2, 3, 3

# 定义一个子模型，它将成为我们最终模型的一部分
sub_input1 = Input(shape=(C,))
sub_mapped1 = Dense(D)(sub_input1)
sub_model1 = Model(inputs=sub_input1, outputs=sub_mapped1)

# 定义另一个子模型，它将成为我们最终模型的一部分
sub_input2 = Input(shape=(C,))
sub_mapped2 = Dense(D)(sub_input2)
sub_model2 = Model(inputs=sub_input2, outputs=sub_mapped2)

# 定义一个基于前两个子模型的模型
input1 = Input(shape=(D,))
input2 = Input(shape=(D,))
mapped1_2 = sub_model1(input1)
mapped2_2 = sub_model2(input2)
sub_sum = Add()([mapped1_2, mapped2_2])
keras_model = Model(inputs=[input1, input2], outputs=sub_sum)

# 开始转换！`target_opset` 参数是可选的。
onnx_model = onnxmltools.convert_keras(keras_model, target_opset=7)

CoreML 到 ONNX 的转换

以下是一个简单的代码片段，用于将 Core ML 模型转换为 ONNX 模型。

import onnxmltools
import coremltools

# 加载一个 Core ML 模型
coreml_model = coremltools.utils.load_spec('example.mlmodel')

# 将 Core ML 模型转换为 ONNX
onnx_model = onnxmltools.convert_coreml(coreml_model, 'Example Model')

# 保存为 protobuf 格式
onnxmltools.utils.save_model(onnx_model, 'example.onnx')

H2O 到 ONNX 的转换

以下是将 H2O MOJO 模型转换为 ONNX 模型的代码片段。唯一的先决条件是本地文件系统中已保存 MOJO 模型。

import onnxmltools

# 将 Core ML 模型转换为 ONNX
onnx_model = onnxmltools.convert_h2o('/path/to/h2o/gbm_mojo.zip')

# 保存为 protobuf 格式
onnxmltools.utils.save_model(onnx_model, 'h2o_gbm.onnx')

测试模型转换器

onnxmltools 可以将模型转换为 ONNX 格式，随后您便可以使用任意后端来执行预测。

检查转换后 ONNX 模型的操作集版本

您可以通过 Netron，一款神经网络模型查看器，来检查转换后 ONNX 模型的操作集版本。此外，您也可以通过以下代码行确定转换后模型的操作集版本：

opset_version = onnx_model.opset_import[0].version

如果检查结果低于您在 onnxmltools.convert 函数中指定的 target_opset 值，这通常是预期行为。ONNXMLTools 转换器会逐个将操作符转换为 ONNX 格式，并为其找到最近更新的操作集版本。当所有操作符都转换完毕后，生成的 ONNX 模型将采用其包含的所有操作符中的最高操作集版本。

为了更具体地说明这一点，我们考虑一个包含 Abs 和 Add 两个操作符的模型。截至 2018 年 12 月，Abs 最近一次更新是在 opset 6，而 Add 最近一次更新是在 opset 7。因此，即使您请求 target_opset=8，转换后的 ONNX 模型的操作集版本也将始终是 7。这种转换行为是为了确保向后兼容性而设计的。

有关 ONNX 模型格式 的文档以及更多来自不同框架的模型转换示例，可在 ONNX 教程 仓库中找到。

测试所有现有转换器

所有转换器的单元测试都可以自动生成原始模型和转换后的模型，并使用 onnxruntime 或 onnxruntime-gpu 自动进行验证。这些单元测试案例均为标准的 Python 单元测试，您可以通过 pytest 命令行运行，例如：

python -m pytest --ignore .\tests\

大多数模型的测试需要 onnxruntime 和 numpy；与文本特征相关的转换则需要 pandas；稀疏特征则需要 scipy。还有一个测试需要 keras 来测试自定义操作符，这意味着也需要 sklearn 或其他机器学习库。

添加一个新的转换器

一旦转换器实现完毕，就会添加一个单元测试来确认其正常工作。在单元测试的最后，必须调用函数 dump_data_and_model 或任何等效函数，以导出预期的输出和转换后的模型。这些文件生成后，还需在 tests_backend 中添加相应的测试，以便使用运行时计算预测结果。

许可证

Apache许可证v2.0

ONNXMLTools 快速上手指南

ONNXMLTools 是一个用于将多种机器学习框架（如 scikit-learn、Core ML、LightGBM、XGBoost 等）训练的模型转换为 ONNX 格式的工具库，便于模型在不同推理引擎间部署。

环境准备

系统要求

• Python 版本：3.7 及以上

前置依赖

本工具依赖 ONNX、NumPy 和 ProtoBuf。根据你要转换的源模型框架，还需安装对应的库：

源框架	需安装的额外依赖
scikit-learn	scikit-learn
Core ML	coremltools (版本 ≤ 3.1)
Keras/TensorFlow	keras (≥ 2.0.8) 及对应版本的 tensorflow
LightGBM	lightgbm
Spark ML	pyspark
XGBoost	xgboost
LibSVM	libsvm
H2O	h2o
CatBoost	catboost

注意：如果你计划从源码安装，请在安装 onnx 包之前设置环境变量：export ONNX_ML=1。

安装步骤

方式一：通过 PyPI 安装（推荐）

使用国内镜像源加速安装：

pip install onnxmltools -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从源码安装

pip install git+https://github.com/onnx/onnxmltools

基本使用

以下示例展示如何将一个 Keras 模型转换为 ONNX 格式。其他框架（如 Core ML、H2O 等）的使用方式类似，只需调用对应的 convert_<framework> 函数。

示例：Keras 模型转 ONNX

import onnxmltools
from keras.layers import Input, Dense, Add
from keras.models import Model

# 1. 构建一个简单的 Keras 模型
N, C, D = 2, 3, 3

sub_input1 = Input(shape=(C,))
sub_mapped1 = Dense(D)(sub_input1)
sub_model1 = Model(inputs=sub_input1, outputs=sub_mapped1)

sub_input2 = Input(shape=(C,))
sub_mapped2 = Dense(D)(sub_input2)
sub_model2 = Model(inputs=sub_input2, outputs=sub_mapped2)

input1 = Input(shape=(D,))
input2 = Input(shape=(D,))
mapped1_2 = sub_model1(input1)
mapped2_2 = sub_model2(input2)
sub_sum = Add()([mapped1_2, mapped2_2])

keras_model = Model(inputs=[input1, input2], outputs=sub_sum)

# 2. 转换为 ONNX 模型
# target_opset 是可选参数，用于指定目标 ONNX 算子集版本
onnx_model = onnxmltools.convert_keras(keras_model, target_opset=7)

# 3. 保存模型
onnxmltools.utils.save_model(onnx_model, 'keras_model.onnx')

其他框架转换简述

• Core ML: 使用 onnxmltools.convert_coreml(coreml_model, 'Model Name')
• H2O: 使用 onnxmltools.convert_h2o('/path/to/model.zip')

转换完成后，你可以使用 Netron 查看模型结构，或使用 onnxruntime 进行推理验证。