DOCSEARCHbyAlgolia
Submit Docs

ONNXMLTools

repository·main·Indexed 22 days ago

https://github.com/onnx/onnxmltools

A toolkit for converting machine learning models from various frameworks into the ONNX format. It supports conversions from TensorFlow, scikit-learn, Core ML, XGBoost, LightGBM, CatBoost, H2O, and Spark ML. The framework utilizes an Intermediate Representation (IR) to decouple source formats from the target ONNX format through a pipeline of parsing, compiling, and converting.

Tokens
9.2K
Snippets
30
Records
54
Agent score
78%

What's inside onnxmltools

  1. Convert machine learning models to ONNX with onnxmltools

    main

    ONNXMLTools is a library used to convert models from various machine learning toolkits into the ONNX format. This allows models trained in different frameworks to be used with an ONNX-compatible backend like onnxruntime or onnxruntime-gpu for inference.

    Supported toolkits include:

    • Apple Core ML (up to version 3.1)
    • catboost
    • h2o (subset only)
    • Keras
    • LightGBM
    • libsvm
    • scikit-learn (subset of models)
    • SparkML
    • XGBoost

    The library leverages sklearn-onnx and tensorflow-onnx for specific conversions and provides its own converters for the other supported libraries.

  2. Understand the ONNXMLTools Conversion Framework

    main

    The conversion framework in onnxmltools is designed to translate computational graphs from one format (like Core ML, scikit-learn, or Keras) into ONNX. It relies on an Intermediate Representation (IR) to decouple the source format from the target ONNX format.

    Core Components:

    • Intermediate Representation (IR): A collection of data structures (Topology, Scope, Operator, Variable, Type) used to represent the graph.
    • Containers: Objects like RawModelContainer (and its derivatives CoremlModelContainer, SklearnModelContainer, KerasModelContainer) that store the raw model and its input/output names.
    • Parsers: Translate raw models into a Topology object.
    • Compiler: A collection of functions in the Topology class that performs graph optimization, shape inference, and basic checks via the compile() method.
    • Shape Calculators: Determine the type and shape of variables.
    • Converters: Functions that transform an Operator into ONNX components (nodes and initializers).
    • Registration: A mechanism to map Operator types to specific shape calculators and converters.

  3. How Shape Inference and Mapping works

    main

    Shape inference is a sequence of function calls where each call takes an Operator and calculates the type (including the shape field) of its output variables. A shape calculator is only invoked once all input variables for an operator have been initialized.

    Core ML to IR Shape Mapping Rules:

    • [C, H, W] $\rightarrow$ [N, C, H, W]
    • [C] $\rightarrow$ [N, C]
    • [S, C] $\rightarrow$ [N, C]
    • [S, C, H, W] $\rightarrow$ [N, C, H, W]
    • Scalar shapes (e.g., Int64Type) $\rightarrow$ [1, 1]

    Important Notes:

    • Batch Size (N): Core ML's batch size is ignored in the graph structure. By default, N=1 is used for traditional ML models and N='None' for neural networks.
    • Overriding Defaults: You can provide initial_types when calling convert(...) (e.g., in onnxmltools.convert.coreml.convert) to overwrite default types.
    • Scikit-learn: Typically expects [1, C] for feature vector inputs.

  4. How the Intermediate Representation (IR) works

    main

    The IR is the backbone of the conversion process, using a hierarchical structure to represent the computational graph:

    • Topology: The top-level structure. It contains Scope objects and provides graph-wide processing functions, such as topological_operator_iterator for traversing operators in execution order.
    • Scope: Acts as a container for operators and variables. It provides a naming mechanism to ensure all names are unique and allows for recursive parsing by isolating components.
    • Operator: The smallest unit of computation. Each Operator has an input list and an output list containing Variable objects. The actual logic of the operator is stored in its raw_operator field.
    • Variable: Represents data flowing through the graph. Each Variable has a type field containing shape information.

    Accessing Shapes: To access the shape of a variable x, use x.type.shape. This returns a list of integers and strings. The string 'None' is used to represent a variable-length coordinate.

  5. How to use RawModelContainers

    main

    A RawModelContainer (or its subclasses like CoremlModelContainer, SklearnModelContainer, or KerasModelContainer) is used to store the source model and define its entry and exit points.

    To ensure the entire graph is converted, you must correctly assign the input_names and output_names properties. These names define the roots and leaves of the computational graph. If an input name is omitted, the compiler may treat parts of the graph as unreachable and prune them.

  6. How Converters transform Operators to ONNX

    main

    A converter is a function that transforms a single Operator into one or more ONNX components (like NodeProto and initializers).

    Converter Signature: Every converter accepts three arguments:

    1. scope: A Scope object used to declare new operators and variables and to generate unique names.
    2. operator: The Operator object being converted. It contains the raw_operator and the input/output lists.
    3. container: An object used to store all created ONNX objects. These objects are eventually passed to an ONNX ModelProto.

    Implementation Detail: When a converter creates a sub-graph of ONNX nodes to simulate a raw_operator, it must use the scope naming functions to ensure that the input and output names of the new nodes are correctly connected.

  7. Prepare input dictionary for ONNX inference

    main

    Once the model is converted, you cannot pass a Spark DataFrame directly to an ONNX runtime. You must create a Python dict where keys are the input names and values are the corresponding TensorData (typically converted to NumPy arrays via Pandas).

    For simple cases, you can use the utility function buildInputDictSimple() and pass your testing DataFrame to it.

    Manual creation example:

    input_data = {}
input_data['label'] = test_df.select('label').toPandas().values
# ... (repeat for all desired inputs)

  8. Convert Spark ML models to ONNX

    main

    To convert a Spark ML pipeline to an ONNX model, you must provide the API with the specific Tensor types for the inputs. The process involves defining initial_types, calling convert_sparkml(), and preparing input data as a dictionary for inference.

    Key Requirements:

    • The input names in initial_types must match the column names in your DataFrame and the inputCol(s) values used when creating your Spark Pipeline.
    • For numerical data, it is recommended to use FloatTensorType for all numbers to avoid compatibility issues.
    • For certain components like Word2Vec, the conversion may only support a batch size of 1.
    # 1. Define initial types
initial_types = [
    ("label", StringTensorType([1, 1])),
    # (repeat for required inputs)
]

# 2. Convert the pipeline model
pipeline_model = pipeline.fit(training_data)
onnx_model = convert_sparkml(pipeline_model, 'My Sparkml Pipeline', initial_types)

# 3. Save the model
with open("model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

# 4. Prepare input data for inference
input_data = {}
input_data['label'] = test_df.select('label').toPandas().values

# 5. Run inference using onnxruntime
import onnxruntime
sess = onnxruntime.InferenceSession(onnx_model, providers=["CPUExecutionProvider"])
output = sess.run(None, input_data)

  9. Prepare input types for Spark ML conversion

    main

    When calling convert_sparkml(), you must provide a list of tuples representing the input names and their corresponding Tensor types.

    For simple conversion tasks, you can use the utility function buildInitialTypesSimple() from convert.sparkml.utils by passing your test DataFrame.

    For manual definition, use the following format:

    initial_types = [
    ("column_name", TensorType([shape])),
]

    Note: Input names must match your Spark DataFrame column names and the inputCol(s) specified in your Spark Pipeline.

    initial_types = [
    ("label", StringTensorType([1, 1])),
    # (repeat for the required inputs)
]

  10. Typical Model Conversion Procedure

    main

    Converting a model follows a three-stage pipeline:

    1. Parsing: Translate the raw model into an Intermediate Representation (IR) Topology using a parser (e.g., onnxmltools.convert.coreml._parse).
    2. Compiling: Call compile() on the Topology object. This optimizes the graph, performs shape inference, and applies post-processing rules.
    3. Converting: Invoke the converters for all operators in topological order. This is handled by calling convert_topology(...) in _topology.py.

    High-level API Entrypoints:

    • Core ML: onnxmltools.convert.coreml.convert
    • Scikit-learn: onnxmltools.convert.sklearn.convert

  11. Run ONNX model predictions with ONNX Runtime

    main

    After converting a model to ONNX, use onnxruntime to perform inference. Create an InferenceSession by providing the model path and specifying an execution provider (e.g., CPUExecutionProvider). Use sess.run() to execute the model, passing the desired output names and a dictionary mapping input names to the input data (ensuring data types like numpy.float32 match the model requirements).

    import onnxruntime as rt
import numpy as np

# Load the session
sess = rt.InferenceSession("logreg_iris.onnx", providers=["CPUExecutionProvider"])

# Get the input name from the session
input_name = sess.get_inputs()[0].name

# Run inference
# Note: X_test must be cast to the correct type (e.g., float32)
pred_onx = sess.run([label_name], {input_name: X_test.astype(np.float32)})[0]

  12. Install ONNXMLTools

    main

    You can install the latest release of ONNXMLTools from PyPI using pip.

    If you choose to install from source, you must set the environment variable ONNX_ML=1 before installing the onnx package.

    # Install from PyPI
pip install onnxmltools

# Install from source
pip install git+https://github.com/onnx/onnxmltools