MLOnnx

pipeline pipeline

Exports a traditional machine learning model (i.e. scikit-learn) to ONNX.

Example

See the link below for a detailed example.

NotebookDescription
Export and run other machine learning modelsExport and run models from scikit-learn, PyTorch and moreOpen In Colab

Methods

Python documentation for the pipeline.

Exports a machine learning model to ONNX using ONNXMLTools.

Parameters:

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

TypeDescription

path to model output or model as bytes depending on output parameter

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

  1. def __call__(self, model, task="default", output=None, opset=12):
  2. """
  3. Exports a machine learning model to ONNX using ONNXMLTools.
  4. Args:
  5. model: model to export
  6. task: optional model task or category
  7. output: optional output model path, defaults to return byte array if None
  8. opset: onnx opset, defaults to 12
  9. Returns:
  10. path to model output or model as bytes depending on output parameter
  11. """
  12. # Convert scikit-learn model to ONNX
  13. model = convert_sklearn(model, task, initial_types=[("input_ids", StringTensorType([None, None]))], target_opset=opset)
  14. # Prune model graph down to only output probabilities
  15. model = select_model_inputs_outputs(model, outputs="probabilities")
  16. # pylint: disable=E1101
  17. # Rename output to logits for consistency with other models
  18. model.graph.output[0].name = "logits"
  19. # Find probabilities output node and rename to logits
  20. for node in model.graph.node:
  21. if node.output[0] == "probabilities":
  22. node.output[0] = "logits"
  23. # Save model to specified output path or return bytes
  24. model = save_onnx_model(model, output)
  25. return output if output else model