MLOnnx

将传统的机器学习模型（例如 scikit-learn）导出为 ONNX 格式。

示例

参阅下方链接查看详细示例。

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

方法

流水线的 Python 文档。

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

使用 ONNXMLTools 将机器学习模型导出为 ONNX 格式。

参数

名称	类型	描述	默认值
model		要导出的模型	必需
task		可选的模型任务或类别	'default'
output		可选的输出模型路径，如果为 None 则默认为返回字节数组	None
opset		onnx opset，默认为 12	12

返回值

类型	描述
	模型输出路径或模型字节数组，取决于 output 参数

源代码位于 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