ONNXMLTools简介
ONNXMLTools是一个开源的Python库,旨在将各种机器学习框架训练的模型转换为开放神经网络交换格式(ONNX)。ONNX是一种用于表示深度学习模型的开放格式,可以在不同的框架和工具之间进行交换。ONNXMLTools的出现使得模型的互操作性和可移植性得到了极大的提升。
主要特性
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开发中扮演更加重要的角色。
