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

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: