AIMET TensorFlow 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