ONNX Optimizer

0.4.2 · active · verified Thu Apr 16

ONNX Optimizer is a Python library that provides a C++ library for performing arbitrary optimizations on ONNX models, offering a growing list of prepackaged optimization passes. It aims to share optimization work between various ONNX backend implementations, making models more efficient for inference. The project is actively maintained with frequent minor releases, often every few months.

Common errors

```
subprocess.CalledProcessError: Command '['/usr/bin/cmake', '--build', '.', '--', '-j', '4']' returned non-zero exit status 2.
```
cause This error typically occurs during installation from source (when `pip` tries to build the wheel) if the `cmake` build system or other necessary C++ build tools (like `protobuf` development headers) are not present on the system.

fix Ensure `cmake` is installed and properly configured, and that `protobuf` development files are available. For example, on Ubuntu: `sudo apt-get update && sudo apt-get install cmake build-essential libprotobuf-dev protobuf-compiler`.
```
AttributeError: module 'onnx' has no attribute 'optimizer' or ImportError: cannot import name 'optimizer' from 'onnx'
```
cause Attempting to import the optimizer directly from the `onnx` package, which is no longer the correct path for `onnx` versions 1.9.0 and newer. The optimizer functionality was moved to the `onnxoptimizer` package.

fix Install the `onnxoptimizer` package separately (`pip install onnxoptimizer`) and change your import statements from `from onnx import optimizer` to `import onnxoptimizer`.
```
[ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Invalid rank for input: float_input Got: 1 Expected: 2 Please fix either the inputs/outputs or the model.
```
cause While not directly an `onnxoptimizer` error, this is a common issue when using optimized ONNX models with `onnxruntime` if the input shape provided to the runtime session does not match the expectations of the optimized model. This could indicate a misunderstanding of the model's new input shape after optimization, or an issue with how the model was optimized.

fix Verify the expected input shapes of the optimized ONNX model using `onnx.load("optimized_model.onnx").graph.input` and ensure your input data conforms to these shapes. Use tools like Netron to visually inspect the optimized model graph.

Warnings

breaking The ONNX optimizer functionality was moved from the `onnx` package to a separate `onnxoptimizer` package. Direct imports like `from onnx import optimizer` will fail for `onnx` versions 1.9.0 and later.
Fix: Install `onnxoptimizer` as a separate package (`pip install onnxoptimizer`) and update imports to `import onnxoptimizer` or `from onnxoptimizer import optimize`.
gotcha Building `onnxoptimizer` from source (e.g., when pre-built wheels are unavailable or for specific development needs) requires `protobuf` to be installed on your system.
Fix: Ensure `protobuf` is installed before attempting to build from source (e.g., `pip install protobuf` or `sudo apt-get install libprotobuf-dev protobuf-compiler`).
gotcha Some advanced or custom graph structures might not be optimized as expected. The optimizer relies on exact subgraph matching for certain transformations.
Fix: Carefully inspect the optimized graph using tools like Netron. If expected optimizations are not applied, consider simplifying the model or implementing custom passes if necessary.

Install

pip install onnxoptimizer Install from PyPI

Imports

optimize
wrong:
```
from onnx import optimizer
```
correct:
```
import onnxoptimizer
optimized_model = onnxoptimizer.optimize(model)
```
The ONNX optimizer was moved to its own package, 'onnxoptimizer', after ONNX version 1.9. Direct import from 'onnx' will fail.

Quickstart

This quickstart demonstrates how to load an ONNX model (here, a dummy one is created), apply optimizations using `onnxoptimizer.optimize()`, and then save the resulting optimized model. It also includes an optional step to verify the optimized model using ONNX Runtime.

import onnx
import onnxoptimizer

# Create a dummy ONNX model for demonstration
# In a real scenario, you would load an existing .onnx file
from onnx import helper
from onnx import TensorProto

graph_def = helper.make_graph(
    [
        helper.make_node("Add", ["input1", "input2"], ["output"]), # Example operation
    ],
    "simple_graph",
    [
        helper.make_tensor_value_info("input1", TensorProto.FLOAT, [1, 2]),
        helper.make_tensor_value_info("input2", TensorProto.FLOAT, [1, 2]),
    ],
    [
        helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 2]),
    ],
)

model = helper.make_model(graph_def, producer_name='test-model')

# Save the dummy model (optional, for verification)
onnx.save(model, "original_model.onnx")
print("Original model saved to original_model.onnx")

# Optimize the model
# You can specify passes, e.g., ['fuse_bn_into_conv']
# By default, a set of common fusion and elimination passes are used.
optimized_model = onnxoptimizer.optimize(model)

# Save the optimized model
onnx.save(optimized_model, "optimized_model.onnx")
print("Optimized model saved to optimized_model.onnx")

# Optional: Verify the optimized model (requires onnxruntime)
try:
    import onnxruntime as ort
    sess_options = ort.SessionOptions()
    _ = ort.InferenceSession("optimized_model.onnx", sess_options)
    print("Optimized model successfully loaded by ONNX Runtime.")
except ImportError:
    print("onnxruntime not installed. Cannot verify optimized model.")
except Exception as e:
    print(f"Error loading optimized model with ONNX Runtime: {e}")

view raw JSON →