ONNXMLTools: 跨平台机器学习模型转换工具

ONNXMLTools简介

ONNXMLTools是一个开源的Python库,旨在将各种机器学习框架训练的模型转换为开放神经网络交换格式(ONNX)。ONNX是一种用于表示深度学习模型的开放格式,可以在不同的框架和工具之间进行交换。ONNXMLTools的出现使得模型的互操作性和可移植性得到了极大的提升。

ONNXMLTools logo

主要特性

ONNXMLTools支持将以下框架和工具的模型转换为ONNX格式:

TensorFlow (通过封装tf2onnx转换器)
scikit-learn (通过封装skl2onnx转换器)
Apple Core ML
Spark ML (实验性支持)
LightGBM
libsvm
XGBoost
H2O
CatBoost

值得注意的是,PyTorch已经内置了ONNX导出功能,因此不需要通过ONNXMLTools进行转换。

安装和依赖

你可以通过pip轻松安装ONNXMLTools的最新版本:

pip install onnxmltools

或者从源代码安装:

pip install git+https://github.com/microsoft/onnxconverter-common
pip install git+https://github.com/onnx/onnxmltools

ONNXMLTools依赖于ONNX、NumPy和ProtoBuf库。此外,根据你要转换的模型类型,你可能还需要安装相应的机器学习库,如scikit-learn、CoreMLTools、Keras、LightGBM等。

ONNXMLTools支持Python 3.7及以上版本。

使用示例

Keras模型转换

下面是一个将Keras模型转换为ONNX格式的简单示例:

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

# 定义Keras模型
input1 = Input(shape=(3,))
input2 = Input(shape=(3,))
dense1 = Dense(3)(input1)
dense2 = Dense(3)(input2)
sum = Add()([dense1, dense2])
keras_model = Model(inputs=[input1, input2], outputs=sum)

# 转换为ONNX模型
onnx_model = onnxmltools.convert_keras(keras_model, target_opset=7)

Core ML模型转换

以下代码展示了如何将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')

# 保存ONNX模型
onnxmltools.utils.save_model(onnx_model, 'example.onnx')

高级特性

指定目标操作集版本

在转换模型时,你可以通过target_opset参数指定目标ONNX操作集版本,以确保生成的ONNX模型与特定版本兼容:

onnx_model = onnxmltools.convert_keras(keras_model, target_opset=7)

检查ONNX模型的操作集版本

你可以通过以下代码查看转换后的ONNX模型使用的操作集版本:

opset_version = onnx_model.opset_import[0].version

值得注意的是,最终ONNX模型的操作集版本可能会低于指定的target_opset。这是因为ONNXMLTools会为每个操作符选择最近更新的版本,然后取所有操作符版本的最大值作为最终模型的版本。

贡献和测试

ONNXMLTools欢迎社区贡献。如果你想添加新的转换器,可以实现转换逻辑并添加相应的单元测试。所有的转换器单元测试都可以生成原始模型和转换后的模型,并通过onnxruntime或onnxruntime-gpu自动进行检查。

你可以使用pytest运行单元测试:

python -m pytest --ignore .\tests\

总结

ONNXMLTools为机器学习从业者提供了一个强大的工具,能够将各种流行框架的模型转换为统一的ONNX格式。这不仅提高了模型的可移植性,还为模型部署和推理提供了更多的灵活性。无论你是使用TensorFlow、scikit-learn还是其他支持的框架,ONNXMLTools都能帮助你轻松地将模型转换为ONNX格式,为跨平台和跨框架的模型应用铺平道路。

随着机器学习技术的不断发展,模型互操作性的重要性日益凸显。ONNXMLTools作为连接不同机器学习生态系统的桥梁,必将在未来的AI开发中扮演更加重要的角色。