unilab.training.onnx_export.verify_policy_onnx

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.