hyrax.model_exporters

Attributes

logger

Functions

export_to_onnx(model, sample, config, ctx)

Dispatching function to convert a ML framework model into an ONNX model.

_export_pytorch_to_onnx(model, sample, ...)

Specific implementation to convert PyTorch model to ONNX format. This

Module Contents

logger[source]
export_to_onnx(model, sample, config, ctx)[source]

Dispatching function to convert a ML framework model into an ONNX model.

Parameters:

  • model (ML framework model) – The model that was just trained using the ML framework. i.e. PyTorch

  • sample (Tensor) – A single sample from the training data loader. This is used to check the output of the ONNX model against the output of the PyTorch model.

  • config (dict) – The parsed config file as a nested dict

  • ctx (dict) – A context dictionary containing info needed for the conversion to ONNX.

_export_pytorch_to_onnx(model, sample, output_filepath, opset_version)[source]

Specific implementation to convert PyTorch model to ONNX format. This function will also: - Run sample through the model before converting the model to ONNX - Convert sample to a numpy array