MLOnnx

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

Methods

Python documentation for the pipeline.

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"

    # Find probabilities output node and rename to logits
    for node in model.graph.node:
        if node.output[0] == "probabilities":
            node.output[0] = "logits"

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