AIMET Keras Cross Layer Equalization APIs

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, input_shapes)

High-level API to perform Cross-Layer Equalization (CLE) on the given model

Parameters

  • model (Model) – tf.keras.Model

  • input_shapes (Union[None, Tuple, List[Tuple]]) – Input shape tuple or list of input tuple shape

Return type

Model

Returns

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, input_shape)

Primitive APIs

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