MLOnnx

将传统机器学习模型（即 scikit-learn）导出为 ONNX。

示例

请参阅以下链接以获取详细示例。

Notebook	描述
导出并运行其他机器学习模型	从 scikit-learn、PyTorch 等导出并运行模型

方法

管道的 Python 文档。

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

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:
        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