onnxmltools - Facilitate Cross-Platform Machine Learning by Converting Models to ONNX

Introduction to ONNXMLTools

ONNXMLTools is a remarkable tool that simplifies the conversion of machine learning models from various toolkits into the Open Neural Network Exchange (ONNX) format. This capability is particularly useful for developers and researchers who wish to use their models across different platforms and environments without compatibility issues.

The project currently supports a wide range of machine learning frameworks, including but not limited to:

TensorFlow: Through a wrapper of the tf2onnx converter.
scikit-learn: Via a wrapper of the skl2onnx converter.
Apple Core ML: For implementing models initially built for Apple's ecosystem.
Spark ML: Although in an experimental phase, it shows potential for handling big data ML workflows.
LightGBM, libsvm, XGBoost, H2O, CatBoost: Popular models used in both academia and industry.

For PyTorch users, it's noteworthy that PyTorch has a built-in exporter for ONNX, which can be found through the PyTorch documentation.

Installation of ONNXMLTools

Installing ONNXMLTools is straightforward and can be done using Python's package manager, pip. Users can choose to install it from the PyPi repository or directly from the source.

To install the latest release from PyPi, use the following command:

pip install onnxmltools

To install from the source, execute:

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

Note: If installing from the source, remember to set the environment variable ONNX_ML=1 before installing the onnx package to ensure everything operates smoothly.

Dependencies

For a seamless experience, ONNXMLTools requires certain dependencies. Primarily, it depends on ONNX, NumPy, and ProtoBuf. Moreover, depending on the original toolkit used for the model, additional dependencies might be necessary, such as:

scikit-learn: For scikit-learn model conversions.
CoreMLTools: Needed for Core ML model conversions, specifically versions 3.1 or lower.
Keras and corresponding TensorFlow versions: Required for models built using Keras.
Framework-specific packages like LightGBM, SparkML, XGBoost, H2O, and CatBoost.

The tools and conversions are optimized for Python 3.7 and above.

Examples

To demonstrate the capabilities of ONNXMLTools, below are several examples of converting models from different frameworks to ONNX format:

Keras to ONNX Conversion

Here is an example of converting a Keras model:

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

# Define model
sub_input1 = Input(shape=(3,))
sub_mapped1 = Dense(3)(sub_input1)
sub_model1 = Model(inputs=sub_input1, outputs=sub_mapped1)

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

input1 = Input(shape=(3,))
input2 = Input(shape=(3,))
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)

# Convert model
onnx_model = onnxmltools.convert_keras(keras_model, target_opset=7)

CoreML to ONNX Conversion

Below is how you can convert a Core ML model:

import onnxmltools
import coremltools

coreml_model = coremltools.utils.load_spec('example.mlmodel')

onnx_model = onnxmltools.convert_coreml(coreml_model, 'Example Model')

onnxmltools.utils.save_model(onnx_model, 'example.onnx')

H2O to ONNX Conversion

For H2O model conversions, use the following:

import onnxmltools

onnx_model = onnxmltools.convert_h2o('/path/to/h2o/gbm_mojo.zip')

onnxmltools.utils.save_model(onnx_model, 'h2o_gbm.onnx')

Testing Model Converters

After conversion, it's crucial to verify the functionality of ONNX models. ONNXMLTools allows users to test these using the onnxruntime, which can handle the prediction tasks. Additionally, you can inspect model operators using the Netron tool or programmatically retrieve details such as the opset version.

Adding a New Converter

Expanding the tool's functionality by adding a new converter involves implementing the converter logic and creating unit tests to ensure it operates correctly. Tests help in maintaining reliability and quality across different use cases.

ONNXMLTools is open-source, licensed under the Apache License v2.0, facilitating contributions and collaborative improvements. By converting diverse model formats into ONNX, ONNXMLTools bridges the gap, allowing seamless deployment and interchangeability across a broad range of AI tools and platforms.