AIMET Keras Cross Layer Equalization APIs¶

User Guide Link¶

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

Examples Notebook Link¶

For an end-to-end notebook showing how to use Keras Cross Layer Equalization, please see here.

Introduction¶

AIMET functionality for Keras Cross Layer Equalization supports three techniques:

BatchNorm Folding
Cross Layer Scaling
High Bias Fold

Cross Layer Equalization API¶

Listed below is a comprehensive API to apply all available techniques under cross layer equalization. It performs ‘auto’ detection of candidate layers and applies the techniques. If there are no BatchNorm layers in a given model, BatchNorm fold and high bias fold shall be skipped.

API(s) for Cross Layer Equalization

aimet_tensorflow.keras.cross_layer_equalization.equalize_model(model)[source]¶: High-level API to perform Cross-Layer Equalization (CLE) on the given model :type model: Model :param model: tf.keras.Model :rtype: Model :return: CLE applied tf.keras.Model

Code Example¶

Required imports

import tensorflow as tf
from aimet_tensorflow.keras.cross_layer_equalization import equalize_model

Cross Layer Equalization in auto mode comprehensive

def cross_layer_equalization_auto():
    input_shape = (224, 224, 3)
    model = tf.keras.applications.ResNet50()

    cle_applied_model = equalize_model(model)

Primitive APIs¶

If the user would like to call the APIs individually, then the following APIs can be used:

Primitive APIs for Cross Layer Equalization