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:

Returns:

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