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: _ModelProtoOrIr, num_iterations: int = 2, *, onnx_shape_inference: bool = True, stop_if_no_change: bool = True, input_size_limit: int = 512, output_size_limit: int = 262144, inline: bool = True) → _ModelProtoOrIr[source]

Optimizes a model.

Parameters:

  • model – The model to be optimized.

  • num_iterations – Number of times the optimization loop is repeated.

  • onnx_shape_inference – Applies node-level shape-inference as part of optimization

  • input_size_limit – Will not apply constant folding to ops with any input of size greater than this. Does not apply to special ops like Shape() and Size().

  • output_size_limit – Will not rewrite any foldable-op into a Constant op if the size of the output tensor is greater than this.

  • stop_if_no_change – Stop the optimization loop if no change is detected in an iteration.

  • inline – If True, inlines all functions in the model.

Returns:

The optimized model. If the input was a ModelProto, the output will also be a ModelProto. If the input was an ir.Model, the output will also be an ir.Model.

Description of optimizations applied by onnxscript.optimizer.optimize

Optimization

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.