Skip to content

MLOnnx

pipeline pipeline

Exports a traditional machine learning model (i.e. scikit-learn) to ONNX.

Example

See the link below for a detailed example.

Notebook Description
Export and run other machine learning models Export and run models from scikit-learn, PyTorch and more Open In Colab

Methods

Python documentation for the pipeline.

__call__(self, model, task='default', output=None, opset=12) special

Exports a machine learning model to ONNX using ONNXMLTools.

Parameters:

Name Type Description Default
model

model to export

required
task

optional model task or category

'default'
output

optional output model path, defaults to return byte array if None

None
opset

onnx opset, defaults to 12

12

Returns:

Type Description

path to model output or model as bytes depending on output parameter

Source code in txtai/pipeline/train/mlonnx.py
def __call__(self, model, task="default", output=None, opset=12):
    """
    Exports a machine learning model to ONNX using ONNXMLTools.

    Args:
        model: model to export
        task: optional model task or category
        output: optional output model path, defaults to return byte array if None
        opset: onnx opset, defaults to 12

    Returns:
        path to model output or model as bytes depending on output parameter
    """

    # Convert scikit-learn model to ONNX
    model = convert_sklearn(model, task, initial_types=[("input_ids", StringTensorType([None, None]))], target_opset=opset)

    # Prune model graph down to only output probabilities
    model = select_model_inputs_outputs(model, outputs="probabilities")

    # pylint: disable=E1101
    # Rename output to logits for consistency with other models
    model.graph.output[0].name = "logits"
    model.graph.node[0].output[0] = "logits"

    # Save model to specified output path or return bytes
    model = save_onnx_model(model, output)
    return output if output else model