AIMET TensorFlow Cross Layer Equalization Primitive API

Introduction

If a user wants to modify the order of Cross Layer Equalization, not use some features, or manually tweak the list of layers that need to be equalized, the following APIs can be used.

Higher level API can be used for using one or more features one after the other. It automatically finds the layers to be folded or scaled.

Lower level APIs can be used to manually tweak the list of layers to be folded. The user has to pass the list of layers in the correct order that they appear in the model.

Note: Before using High Bias fold, Cross Layer Scaling (CLS) needs to be applied and scaling factors obtained from CLS need to be plugged in to High Bias Fold. And, if there are batchnorm layers, they need to be folded and the info saved to be plugged into high bias fold API.

Higher Level APIs for Cross Layer Equalization

API for Batch Norm Folding

API for Cross Layer Scaling

API for High Bias Folding

Code Examples for Higher Level APIs

Required imports

import tensorflow as tf
from aimet_tensorflow.keras.batch_norm_fold import fold_all_batch_norms
from aimet_tensorflow.keras.cross_layer_equalization import HighBiasFold, CrossLayerScaling
from aimet_tensorflow.keras.utils.model_transform_utils import replace_relu6_with_relu

Perform Cross Layer Equalization in auto mode step by step

def cross_layer_equalization_auto_stepwise():
    """
    Individual api calls to perform cross layer equalization one step at a time. Pairs to fold and
    scale are found automatically.
    1. Replace Relu6 with Relu
    2. Fold batch norms
    3. Perform cross layer scaling
    4. Perform high bias fold
    """

    # Load the model to equalize
    model = tf.keras.applications.resnet50.ResNet50(weights=None, classes=10)

    # 1. Replace Relu6 layer with Relu
    model_for_cle, _ = replace_relu6_with_relu(model)

    # 2. Fold all batch norms
    folded_pairs, model = fold_all_batch_norms(model_for_cle)

    bn_dict = {}
    for conv_or_linear, bn in folded_pairs:
        bn_dict[conv_or_linear] = bn

    # 3. Perform cross-layer scaling on applicable layer groups
    cls_set_info_list = CrossLayerScaling.scale_model(model_for_cle)

    # 4. Perform high bias fold
    HighBiasFold.bias_fold(cls_set_info_list, bn_dict)

    return model_for_cle

Lower Level APIs for Cross Layer Equalization

API for Batch Norm Folding on subsets of convolution-batchnorm layer pairs

API for Cross Layer Scaling on subset of conv layer groups

API for High bias folding

Custom Datatype used

Code Example for Lower level APIs

Required imports

import tensorflow as tf
from aimet_tensorflow.keras.batch_norm_fold import fold_given_batch_norms
from aimet_tensorflow.keras.cross_layer_equalization import HighBiasFold, CrossLayerScaling
from aimet_tensorflow.keras.utils.model_transform_utils import replace_relu6_with_relu

Perform Cross Layer Equalization in manual mode

def cross_layer_equalization_manual():
    """
    Individual api calls to perform cross layer equalization one step at a time. Pairs to fold and
    scale are provided by the user.
    1. Replace Relu6 with Relu
    2. Fold batch norms
    3. Perform cross layer scaling
    4. Perform high bias fold
    """

    # Load the model to equalize
    model = tf.keras.applications.resnet50.ResNet50(weights=None, classes=10)

    # replace any ReLU6 layers with ReLU
    model_for_cle, _ = replace_relu6_with_relu(model)

    # pick potential pairs of conv and bn ops for fold
    layer_pairs = get_example_layer_pairs_resnet50_for_folding(model_for_cle)

    # fold given layers
    fold_given_batch_norms(model_for_cle, layer_pairs=layer_pairs)

    # Cross Layer Scaling
    # Create a list of consecutive conv layers to be equalized
    consecutive_layer_list = get_consecutive_layer_list_from_resnet50_for_scaling(model_for_cle)

    # invoke api to perform scaling on given list of cls pairs
    scaling_factor_list = CrossLayerScaling.scale_cls_sets(consecutive_layer_list)

    # get info from bn fold and cross layer scaling in format required for high bias fold
    folded_pairs, cls_set_info_list = format_info_for_high_bias_fold(layer_pairs,
                                                                     consecutive_layer_list,
                                                                     scaling_factor_list)

    HighBiasFold.bias_fold(cls_set_info_list, folded_pairs)
    return model_for_cle

Example helper methods to perform CLE in manual mode

Helper to pick layers for batchnorm fold

def get_example_layer_pairs_resnet50_for_folding(model: tf.keras.Model):
    """
    Function to pick example conv-batchnorm layer pairs for folding.
    :param model: Keras model containing conv batchnorm pairs to fold
    :return: pairs of conv and batchnorm layers for batch norm folding in Resnet50 model.
    """

    conv_op_1 = model.layers[2]
    bn_op_1 = model.layers[3]

    conv_op_2 = model.layers[7]
    bn_op_2 = model.layers[8]

    conv_op_3 = model.layers[10]
    bn_op_3 = model.layers[11]

    # make a layer pair list with potential the conv op and bn_op pair along with a flag
    # to indicate if given bn op can be folded upstream or downstream.
    # example of two pairs of conv and bn op  shown below
    layer_pairs = [(conv_op_1, bn_op_1, True),
                   (conv_op_2, bn_op_2, True),
                   (conv_op_3, bn_op_3, True)]

    return layer_pairs

Helper to pick layers for cross layer scaling

def get_consecutive_layer_list_from_resnet50_for_scaling(model: tf.keras.Model):
    """
    helper function to pick example consecutive layer list for scaling.
    :param model: tf.keras.Model
    :return: sample layers for scaling as consecutive_layer_list from Resnet50 model
    """
    conv_op_1 = model.layers[2]
    conv_op_2 = model.layers[7]
    conv_op_3 = model.layers[10]

    consecutive_layer_list = [(conv_op_1, conv_op_2), (conv_op_2, conv_op_3)]
    return consecutive_layer_list

Helper to format data from batchnorm fold and cross layer scaling for usage by high bias fold

def format_info_for_high_bias_fold(layer_pairs, consecutive_layer_list, scaling_factor_list):
    """
    Helper function that formats data from cross layer scaling and bn fold for usage by high bias fold
    :param layer_pairs: info obtained after batchnorm fold
    :param consecutive_layer_list: info obtained after cross layer scaling
    :param scaling_factor_list: scaling params corresponding to consecutive_layer_list
    :return: data formatted for high bias fold
    """

    # convert info after batch norm fold and cross layer scaling for usage by high bias fold api
    folded_pairs = []
    for (conv_op, bn_op_with_meta, _fold_upstream_flag) in layer_pairs:
        folded_pairs.append((conv_op, bn_op_with_meta.op))

    # List that hold a boolean for if there were relu activations between layers of each cross layer scaling set
    is_relu_activation_in_cls_sets = []
    # Note the user is expected to fill in this list manually

    # Convert to a list of cls-set-info elements
    cls_set_info_list = CrossLayerScaling.create_cls_set_info_list(consecutive_layer_list,
                                                                   scaling_factor_list,
                                                                   is_relu_activation_in_cls_sets)

    return folded_pairs, cls_set_info_list