helios.onnx

Functions

export_to_onnx(→ None)

Export the given network to ONNX format.

Module Contents

helios.onnx.export_to_onnx(net: torch.nn.Module, net_args: torch.Tensor | tuple[torch.Tensor, Ellipsis], out_path: pathlib.Path, validate_output: bool = False, save_on_validation_fail: bool = True, rtol: float = 0.001, atol: float = 1e-05, **kwargs) → None[source]

Export the given network to ONNX format.

By default, the resulting onnx network will be validated through ONNX to ensure it’s valid. If you wish to validate the traced outputs to ensure they’re the same, set validate_output to true and change rtol/atol as needed.

Note

Starting in torch 2.9.1, torch.onnx.export defaults to dynamo=True. If dynamo is used (whether by default in torch >= 2.9.1 or by explicitly passing dynamo=True in kwargs), then the file path is omitted from the arguments to torch.onnx.export, using instead program.save() to write the file. If dynamo is not used, then it is passed on to torch.onnx.export.

Parameters:

  • net – the network to convert.

  • net_args – the input tensor(s) for tracing.

  • out_path – the path to save the exported network to.

  • validate_output – if true, validation is performed to ensure correctness. Defaults to false.

  • save_on_validation_fail – if true, the ONNX network is saved regardless of whether validation succeeds. If false, the ONNX network is deleted. Defaults to true.

  • rtol – relative tolerance threshold. Defaults to 1e-3.

  • atol – absolute tolerance threshold. Defaults to 1e-5.

  • kwargs – additional keyword arguments to torch.onnx.export.