简介
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/microsoft/onnxconverter-common
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的模型,您需要一个安装了以下相应包的环境:
- scikit-learn
- CoreMLTools (3.1版本或更低)
- Keras (2.0.8版本或更高)及对应的Tensorflow版本
- LightGBM
- SparkML
- XGBoost
- libsvm
- H2O
- 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版本。
opset_version = onnx_model.opset_import[0].version
如果检查ONNX模型的opset结果小于您在onnxmltools.convert函数中指定的target_opset
数字,请放心,这可能是预期的行为。ONNXMLTools转换器的工作原理是将每个算子单独转换为ONNX格式,并找到它最近更新的对应opset版本。一旦所有算子都转换完成,结果ONNX模型将具有所有算子的最大opset版本。
为具体说明这一点,让我们考虑一个包含两个算子Abs和Add的模型。截至2018年12月,Abs最近在opset 6中更新,而Add最近在opset 7中更新。因此,转换后的ONNX模型的opset将始终为7,即使您要求target_opset=8
。转换器的行为被定义为这样以确保向后兼容性。
有关ONNX模型格式的文档以及从不同框架转换模型的更多示例可以在ONNX教程存储库中找到。
测试所有现有转换器
所有转换器单元测试都可以生成原始模型和转换后的模型,以便使用onnxruntime或onnxruntime-gpu自动检查。单元测试用例都是普通的Python单元测试用例,您可以使用pytest命令行运行它们,例如:
python -m pytest --ignore .\tests\
大多数模型需要onnxruntime、numpy,与文本特征相关的转换需要pandas,稀疏特征需要scipy。一个测试还需要keras来测试自定义算子。这意味着需要sklearn或任何机器学习库。
添加新的转换器
一旦实现了转换器,就需要添加一个单元测试来确认它能正常工作。在单元测试的最后,必须调用dump_data_and_model函数或任何等效函数来转储预期输出和转换后的模型。一旦生成这些文件,必须在tests_backend中添加相应的测试,以使用运行时计算预测。