Optimizing a Model using the Optimizer

Introduction

The ONNX Script Optimizer tool provides the user with the functionality to optimize an ONNX model by performing optimizations and clean-ups such as constant folding, dead code elimination, etc.

Usage

In order to utilize the optimizer tool,

import onnxscript

onnxscript.optimizer.optimize(model)

optimize API

The onnxscript.optimizer.optimize call takes in several optional parameters that allows the caller to further fine-tune the process of optimization.

onnxscript.optimizer.optimize(model: ModelProto, num_iterations: int = 2, *, onnx_shape_inference: bool = True, stop_if_no_change: bool = True, external_data_folder: str = '', **kwargs: Any) ModelProto[source]

Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.

Parameters:
  • model (onnx.ModelProto) – The model to optimize.

  • num_iterations (int, optional) – Number of iterations to perform.

  • onnx_shape_inference (bool, optional) – Whether to perform onnx shape inference on the model. Set this to False to turn off onnx shape inference, and rely on model carried shapes and types. This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries the symbolic shapes recorded from dynamo tracing.

  • stop_if_no_change (bool, optional) – Whether to stop if no change is detected.

  • external_data_folder (str, optional) – The folder to store external data.

  • **kwargs – Additional keyword arguments. For BC purposes.

Description of optimizations applied by onnxscript.optimizer.optimize

Optimization ‘onnxscript.optimizer.` + …

Description

Constant folding
constant_folding.fold_constants

Applies constant folding optimization to the model.

Constant propagation
constant_folding.fold_constants

Applies constant propagation optimization to the model. Applied as part of the constant folding optimization.

Sequence simplification
constant_folding.fold_constants

Simplifies Sequence based ops (SequenceConstruct, ConcatFromSequence) present in the model. Applied as part of the constant folding optimization.

Remove unused nodes
remove_unused.remove_unused_nodes

Removes unused nodes from the model.

Remove unused functions
remove_unused_function.remove_unused_functions

Removes unused function protos from the model.

Inline functions with unused outputs
simple_function_folding.inline_functions_with_unused_outputs

Inlines function nodes that have unused outputs.

Inline simple functions
simple_function_folding.inline_simple_functions

Inlines simple functions based on a node count threshold.

List of pattern rewrite rules applied by onnxscript.optimizer.optimize

onnxscript.rewriter.broadcast_to_matmul

onnxscript.rewriter.cast_constant_of_shape

onnxscript.rewriter.gemm_to_matmul_add

onnxscript.rewriter.no_op