4.2.6.5. ONNX

horizon_plugin_pytorch.utils.onnx_helper.export_to_onnx(model, args, f, export_params=True, verbose=False, training=None, input_names=None, output_names=None, operator_export_type=None, opset_version=11, do_constant_folding=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None)

Export a (float or qat) model into ONNX format.

参数

  • model (torch.nn.Module or torch.jit.ScriptModule or ScriptFunction) – The model to be exported.

  • args (tuple or torch.Tensor) – Model inputs such that model(*args) is a valid invocation. Non-Tensor arguments are hard-coded into the exported model.

  • f (file-like or str) – A file-like object or a string containing a file name. A binary protocol buffer will be written to this file.

  • export_params (bool) – If True, all parameters will be exported.

  • verbose (bool) – If True, prints a description of the model being exported.

  • training (enum) – TrainingMode.EVAL, PRESERVE or TRAINING.

  • input_names (list of str) – Names to assign to the input nodes of the graph, in order.

  • output_names (list of str) – Names to assign to the output nodes of the graph, in order.

  • operator_export_type (enum) – ONNX / ONNX_FALLTHROUGH / ONNX_ATEN / ONNX_ATEN_FALLBACK.

  • opset_version (int) – ONNX opset version, default 11.

  • do_constant_folding (bool) – Apply the constant-folding optimization.

  • dynamic_axes (dict) – Axes of tensors that are dynamic (known only at run-time).

  • keep_initializers_as_inputs (bool) – If True, initializers are also added as graph inputs.

  • custom_opsets (dict) – Custom opset domain name to version.