unilab.training.onnx_export

ONNX export and onnxruntime numeric verification helpers for play entrypoints.

Functions

export_policy_onnx(export_module, onnx_path, ...)

Export export_module to onnx_path and print the artifact path.

verify_policy_onnx(export_module, onnx_path, ...)

Compare PyTorch and ONNX Runtime outputs on identical inputs.
unilab.training.onnx_export.export_policy_onnx(export_module, onnx_path, export_inputs, *, input_names, output_names=None, opset_version=17)[source]

Export export_module to onnx_path and print the artifact path.

Parameters:

  • export_module (Module) – Module traced by torch.onnx.export.

  • onnx_path (str) – Destination file path for the exported graph.

  • export_inputs (tuple[Tensor, ...]) – Positional example inputs matching input_names.

  • input_names (list[str]) – ONNX input names, aligned positionally with export_inputs.

  • output_names (list[str] | None) – ONNX output names; defaults to ["action"].

  • opset_version (int) – ONNX opset version; defaults to 17.

Return type:

None

unilab.training.onnx_export.verify_policy_onnx(export_module, onnx_path, verify_inputs, *, input_names, max_diff_tol=0.0001)[source]

Compare PyTorch and ONNX Runtime outputs on identical inputs.

Runs export_module and the exported graph at onnx_path on verify_inputs, prints the max/mean absolute difference, and prints a warning when the max difference exceeds max_diff_tol.

Parameters:

  • export_module (Module) – Module that was exported to onnx_path. Tuple outputs are compared on their first element.

  • onnx_path (str) – Exported ONNX graph to verify.

  • verify_inputs (tuple[Tensor, ...]) – Positional inputs matching input_names.

  • input_names (list[str]) – ONNX input names, aligned positionally with verify_inputs.

  • max_diff_tol (float) – Absolute max-difference tolerance before warning.

Return type:

tuple[float, float]

Returns:

(max_diff, mean_diff) between PyTorch and ONNX Runtime outputs.