项目简介
onnxmltools 是一个强大的工具库,旨在帮助用户将来自不同机器学习工具集的模型转换成 ONNX 格式。目前,它支持的工具集包括 Tensorflow、scikit-learn、Apple Core ML、Spark ML、LightGBM、libsvm、XGBoost、H2O 以及 CatBoost。对于 Pytorch,其自带 ONNX 导出功能。
安装指南
用户可以通过两种方式安装 onnxmltools。第一种是在 PyPi 上下载最新发布的版本,执行以下命令:
pip install onnxmltools
第二种方法是从源代码安装,需要执行以下命令:
pip install git+https://github.com/microsoft/onnxconverter-common
pip install git+https://github.com/onnx/onnxmltools
如果选择从源代码安装,在安装 onnx
包之前,需要设置环境变量 ONNX_ML=1
。
依赖
onnxmltools 依赖于 ONNX、NumPy 和 ProtoBuf。如果转换的是 scikit-learn、Core ML、Keras、LightGBM、SparkML、XGBoost、H2O、CatBoost 或者 LibSVM 的模型,还需要安装相应的软件包。该工具已在 Python 3.7 及以上版本中测试。
转换示例
Keras 转 ONNX 示例
以下示例展示了如何将一个 Keras 模型转换为 ONNX 模型,并指定 target_opset=7
,这对应于 ONNX 版本 1.2:
import onnxmltools
from keras.layers import Input, Dense, Add
from keras.models import Model
# 定义子模型
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)
# 转换为 ONNX 模型
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')
# 转换为 ONNX 模型
onnx_model = onnxmltools.convert_coreml(coreml_model, 'Example Model')
# 保存为 protobuf 格式
onnxmltools.utils.save_model(onnx_model, 'example.onnx')
H2O 转 ONNX 示例
以下代码片段展示了如何将 H2O MOJO 模型转换为 ONNX 模型:
import onnxmltools
# 转换为 ONNX 模型
onnx_model = onnxmltools.convert_h2o('/path/to/h2o/gbm_mojo.zip')
# 保存为 protobuf 格式
onnxmltools.utils.save_model(onnx_model, 'h2o_gbm.onnx')
测试模型转换器
onnxmltools 可以将模型转换为 ONNX 格式,然后使用任意后端进行预测计算。用户可以使用 Netron 或代码查看转换后的 ONNX 模型的操作符集版本。这是确保模型向后兼容的重要步骤。
添加新转换器
一旦实现了新的转换器,需要添加单元测试以确保其工作正常。在单元测试结束时,需调用相关函数以输出预期结果和转换后的模型,并在 tests_backend
中添加测试以使用运行时计算预测。
onnxmltools 项目遵循 Apache License v2.0 授权协议。