onnx_diagnostic.export.api

onnx_diagnostic.export.api.get_main_dispatcher(use_control_flow_dispatcher: bool = False, onnx_plugs: List[EagerDirectReplacementWithOnnx] | None = None) → Any[source][source]

Creates a custom dispatcher for the custom exporter.

onnx_diagnostic.export.api.method_to_onnx(mod: Module, method_name: str = 'forward', input_names: Sequence[str] | None = None, target_opset: int | Dict[str, int] | None = None, verbose: int = 0, filename: str | None = None, output_names: List[str] | None = None, output_dynamic_shapes: Dict[str, Any] | Tuple[Any] | None = None, exporter: str = 'onnx-dynamo', exporter_kwargs: Dict[str, Any] | None = None, save_ep: str | None = None, optimize: bool = True, optimizer_for_ort: bool = True, use_control_flow_dispatcher: bool = False, onnx_plugs: List[EagerDirectReplacementWithOnnx] | None = None, inline: bool = True, convert_after_n_calls: int = 2, patch_kwargs: Dict[str, Any] | None = None, skip_kwargs_names: Set[str] | None = None, dynamic_shapes: Dict[str, Any] | Tuple[Any] | None = None) → Callable[source][source]

Exports one method into ONNX for a module into ONNX. It returns a new method which must be called by the user at least twice with different values for the dynamic dimension between triggering the conversion into ONNX.

Parameters:

  • mod_meth – function to export into ONNX

  • input_names – input names for the onnx model (optional)

  • target_opset – opset to target, if not specified, each converter keeps its default value

  • verbose – verbosity level

  • filename – output filename, mandatory, the onnx model is saved on disk

  • output_names – to change the output of the onnx model

  • output_dynamic_shapes – to overwrite the dynamic shapes names

  • exporter – exporter to use (onnx-dynamo, modelbuilder, custom)

  • exporter_kwargs – additional parameters sent to the exporter

  • save_ep – saves the exported program

  • optimize – optimizes the model

  • optimizer_for_ort – optimizes the model for onnxruntime

  • use_control_flow_dispatcher – use the dispatcher created to supported custom loops (see onnx_diagnostic.export.control_flow_onnx.loop_for_onnx())

  • onnx_plugs – the code was modified to replace some parts with onnx translation

  • inline – inline local functions

  • convert_after_n_calls – converts the model after this number of calls.

  • patch_kwargs – patch arguments

  • skip_kwargs_names – use default values for these parameters part of the signature of the method to export

  • dynamic_shapes – dynamic shapes to use if the guessed ones are not right

Returns:

the output of the selected exporter, usually a structure including an onnx model

See Export a model through method generate (with Tiny-LLM) for an example.

onnx_diagnostic.export.api.to_onnx(mod: Module | GraphModule, args: Sequence[Tensor] | None = None, kwargs: Dict[str, Tensor] | None = None, input_names: Sequence[str] | None = None, target_opset: int | Dict[str, int] | None = None, verbose: int = 0, dynamic_shapes: Dict[str, Any] | Tuple[Any] | None = None, filename: str | None = None, output_names: List[str] | None = None, output_dynamic_shapes: Dict[str, Any] | Tuple[Any] | None = None, exporter: str = 'onnx-dynamo', exporter_kwargs: Dict[str, Any] | None = None, save_ep: str | None = None, optimize: bool = True, optimizer_for_ort: bool = True, use_control_flow_dispatcher: bool = False, onnx_plugs: List[EagerDirectReplacementWithOnnx] | None = None, inline: bool = True) → Any[source][source]

Exports one model into ONNX. Common API for exporters. By default, the models are optimized to use the most efficient kernels implemented in onnxruntime.

Parameters:

  • mod – torch model

  • args – unnamed arguments

  • kwargs – named arguments

  • input_names – input names for the onnx model (optional)

  • target_opset – opset to target, if not specified, each converter keeps its default value

  • verbose – verbosity level

  • dynamic_shapes – dynamic shapes, usually a nested structure included a dictionary for each tensor

  • filename – output filename

  • output_names – to change the output of the onnx model

  • output_dynamic_shapes – to overwrite the dynamic shapes names

  • exporter – exporter to use (onnx-dynamo, modelbuilder, custom)

  • exporter_kwargs – additional parameters sent to the exporter

  • save_ep – saves the exported program

  • optimize – optimizes the model

  • optimizer_for_ort – optimizes the model for onnxruntime

  • use_control_flow_dispatcher – use the dispatcher created to supported custom loops (see onnx_diagnostic.export.control_flow_onnx.loop_for_onnx())

  • onnx_plugs – the code was modified to replace some parts with onnx translation

  • inline – inline local functions

Returns:

the output of the selected exporter, usually a structure including an onnx model

A simple example:

to_onnx(
    model,
    kwargs=inputs,
    dynamic_shapes=ds,
    exporter=exporter,
    filename=filename,
)

Some examples using control flows are available in onnx_diagnostic.export.control_flow_onnx.loop_for_onnx() or onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx.