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

aimet_tensorflow.keras.batch_norm_fold.fold_all_batch_norms(model)¶

Fold all batch_norm layers in a model into corresponding conv layers

Parameters: model (Tuple[Model, Layer]) – model to find all batch norms for

API for Cross Layer Scaling

Under development

API for High Bias Folding

aimet_tensorflow.keras.cross_layer_equalization.HighBiasFold.bias_fold(cls_set_info_list, bn_layers)¶

Folds bias values greater than 3 * sigma to next layer’s bias

Parameters

cls_set_info_list (List[ClsSetInfo]) – List of info elements for each cls set
bn_layers (Dict[Conv2D, BatchNormalization]) – Key: Conv/Linear layer Value: Corresponding folded BN layer

Code Examples for Higher Level APIs¶

Required imports

import tensorflow as tf
from aimet_tensorflow.keras.batch_norm_fold import fold_all_batch_norms

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
    """

    # Note: Cross layer scaling and high bias fold auto stepwise functions still under development.

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

    # 1. Replace Relu6 layer with Relu
    replace_relu6_with_relu(model)

    # 2. Fold all batch norms
    fold_all_batch_norms(model)

Lower Level APIs for Cross Layer Equalization¶

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

aimet_tensorflow.keras.batch_norm_fold.fold_given_batch_norms(model, layer_pairs)¶

Fold a given set of batch_norm layers into conv layers

Parameters

model (Tuple[Model, Layer]) – model to fold selected batchnorms for
layer_pairs (List[Union[Tuple[Union[Conv2D, Dense, Conv2DTranspose, DepthwiseConv2D], BatchNormalization, bool], Tuple[BatchNormalization, Union[Conv2D, Dense, Conv2DTranspose, DepthwiseConv2D], bool]]]) – Tuple of conv, bn layers and is_batch_norm_second flag

API for Cross Layer Scaling on subset of conv layer groups

aimet_tensorflow.keras.cross_layer_equalization.CrossLayerScaling.scale_cls_sets(cls_sets)¶

Scale each cls set

Parameters: cls_sets (List[Union[Tuple[Conv2D, Conv2D], Tuple[Conv2D, DepthwiseConv2D, Conv2D]]]) – Cls sets to scale
Return type: List[Union[ndarray, Tuple[ndarray, ndarray]]]
Returns: List of scale factors corresponding to each scaled cls set

API for High bias folding

aimet_tensorflow.keras.cross_layer_equalization.HighBiasFold.bias_fold(cls_set_info_list, bn_layers)

Folds bias values greater than 3 * sigma to next layer’s bias

Parameters

cls_set_info_list (List[ClsSetInfo]) – List of info elements for each cls set
bn_layers (Dict[Conv2D, BatchNormalization]) – Key: Conv/Linear layer Value: Corresponding folded BN layer

Custom Datatype used¶

class aimet_tensorflow.keras.cross_layer_equalization.ClsSetInfo(cls_pair_1, cls_pair_2=None)¶

This class hold information about the layers in a CLS set, along with corresponding scaling factors and other information like if there is a ReLU activation function between the CLS set layers

Constructor takes 2 pairs if Depth-wise separable layer is being folded

Parameters

cls_pair_1 (ClsSetLayerPairInfo) – Pair between two conv or conv and depth-wise conv
cls_pair_2 (Optional[ClsSetLayerPairInfo]) – Pair between depth-wise conv and point-wise conv

class ClsSetLayerPairInfo(layer1, layer2, scale_factor, relu_activation_between_layers)¶

Models a pair of layers that were scaled using CLS. And related information.

Parameters

layer1 (Conv2D) – Layer whose bias is folded
layer2 (Conv2D) – Layer to which bias of previous layer’s bias is folded
scale_factor (ndarray) – Scale Factor found from Cross Layer Scaling to scale BN parameters
relu_activation_between_layers (bool) – If the activation between layer1 and layer2 is Relu

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 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
    """

    # Note: Cross layer scaling and high bias fold manual functions still under development.

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

    # replace any ReLU6 layers with ReLU
    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)

    # fold given layers
    fold_given_batch_norms(model, 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)

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

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):
    """
    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