ML ONNX

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.

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

Exports a machine learning model to ONNX using ONNXMLTools.

Parameters:

Returns:

Source code in txtai/pipeline/train/mlonnx.py

31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
``` |  ```
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 = convertsklearn(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:
        for x,  in enumerate(node.output):
            if node.output[x] == “probabilities”:
                node.output[x] = “logits”

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

---|---