experimental_experiment.torch_bench.export_model_helper

class experimental_experiment.torch_bench.export_model_helper.WrapExecutorchForTorch(model: Any, forward_method: Callable)

Wraps a executorch model.

class experimental_experiment.torch_bench.export_model_helper.WrapForTorch(torch_model: Any)

Wraps a torch model.

class experimental_experiment.torch_bench.export_model_helper.WrapInferenceSessionForTorch(sess: Any, nvtx: bool = False)

Wraps an onnxruntime.InferenceSession to overload method run to support torch.Tensor.

experimental_experiment.torch_bench.export_model_helper.common_export(model: Any, inputs: List[Any], exporter: str = 'custom', target_opset: int = 18, folder: str = '', filename: str = 'model.onnx', dynamic_shapes: Any | None = None, verbose: int = 0, optimize_oxs: str = '', ort_optimize: bool = False, large_model: bool = False, order: bool = False, enable_pattern: str | List[str] | None = None, disable_pattern: str | List[str] | None = None, stats: Dict[str, Any] | None = None)

Exports a model into a folder.

Parameters:

• model – model

• exporter – torchscript, onnx_dynamo, dynamo_export, custom, …

• folder – folder to export into

• filename – onnx filename

• inputs – inputs

• dynamic_shapes – dynamic shapes

• target_opset – target opset

• optimize_oxs – run optimization with onnxscript

• enable_pattern – patterns to apply

• disable_pattern – patterns not to apply

• verbose – verbosity

• stats – if not None, populates this dictionary with statistics about time

• optimize_oxs – optimize

• ort_optimize – produces a file showing onnxruntime optimizations

• large_model – save weights as external

• order – optimize order

Returns:

onnx proto

experimental_experiment.torch_bench.export_model_helper.compute_weight_size(model: Any) → int

Returns the model size for a torch model or an onnx model. That includes the weights.

experimental_experiment.torch_bench.export_model_helper.run_inference(model: Any, example_inputs: List[Any], warmup: int = 5, repeat: int = 5, verbose: int = 0) → dict[str, Any]

Runs multiple times the same inference.

Args:

model: torch model to run example_inputs: dummy inputs warmup: number of iterations to warmup repeat: number of iterations to repeat verbose: verbosity

Returns:

statistcs

experimental_experiment.torch_bench.export_model_helper.run_onnx_inference(model: ModelProto, example_inputs: List[Any], warmup: int = 5, repeat: int = 5, verbose: int = 0, ort_optimize: bool = True, torch_model: Any | None = None) → Dict[str, Any] | None

Runs multiple times the same inference with onnxruntime.

Args:

model: torch model to run example_inputs: dummy inputs warmup: number of iterations to warmup repeat: number of iterations to repeat verbose: verbosity ort_optimize: enable, disable onnxruntime optimizations torch_model: if not empty, measure the discrepancies

Returns:

statistcs