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, out_path: pathlib.Path, validate_output: bool = False, 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.

Parameters:

  • net – the network to convert.

  • net_args – the input tensor 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.

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

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

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