AIMET ONNX Cross Layer Equalization APIs

User Guide Link

To learn more about this technique, please see Cross-Layer Equalization

Introduction

AIMET functionality for Cross Layer Equalization has 3 features-

BatchNorm Folding
Cross Layer Scaling
High Bias Fold

Cross Layer Equalization API

The following API performs BatchNorm fold followed by Cross Layer Scaling followed by High Bias Fold.

Note: High Bias fold will not happen when the below API is used, if the model does not have BatchNorm layers

API for Cross Layer Equalization

aimet_onnx.cross_layer_equalization.equalize_model(model)[source]

High-level API to perform Cross-Layer Equalization (CLE) on the given model. The model is equalized in place.

Parameters:: model (ModelProto) – Model to equalize

Note: It is recommended to use onnx-simplifier before applying cross layer equalization.

Code Example

Required imports

from onnxsim import simplify
from aimet_onnx.cross_layer_equalization import equalize_model

Cross Layer Equalization in auto mode

def cross_layer_equalization():
    onnx_model = Model()
    # Simplify the model
    onnx_model, _ = simplify(onnx_model)
    equalize_model(onnx_model)