hyrax.model_exporters

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 uses

_export_pytorch_to_onnx_v2(model, sample, ...)

Currently unused.

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) – This sample is the result of running a batch of data through the data loader and the model’s prepare_inputs function. It is used to compare 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 uses the older (torch<2.9) export capabilities. And only supports up the opset version 20.

Parameters:

  • model (torch.nn.Module) – The PyTorch model to be converted to ONNX format.

  • sample (np.ndarray or list of np.ndarray) – A sample of input data to the model. This is used to trace the model during the export process.

  • output_filepath (pathlib.Path) – The file path where the ONNX model will be saved.

  • opset_version (int) – The ONNX opset version to use for the export.

_export_pytorch_to_onnx_v2(model, sample, output_filepath, opset_version)[source]#

Currently unused. Specific implementation to convert PyTorch model to ONNX format using torch Dynamo export capabilities.

Parameters:

  • model (torch.nn.Module) – The PyTorch model to be converted to ONNX format.

  • sample (np.ndarray or list of np.ndarray) – A sample of input data to the model. This is used to trace the model during the export process.

  • output_filepath (pathlib.Path) – The file path where the ONNX model will be saved.

  • opset_version (int) – The ONNX opset version to use for the export.