AIMET TensorFlow Bias Correction API

Bias Correction API

Bias correction is performed after Cross layer equalization on models. Main api to perform bias correction on entire model is listed below.

aimet_tensorflow.bias_correction.BiasCorrection.correct_bias(reference_model, bias_correct_params, quant_params, data_set, conv_bn_dict=None, perform_only_empirical_bias_corr=True)

Top level function for bias correction

Parameters

  • reference_model (Session) – active tf.compat.v1.Session for the model to be corrected.

  • bias_correct_params (BiasCorrectionParams) – input params for bias correction

  • quant_params (QuantParams) – QuantParams type with params for quantization simulation for bias correction.

  • data_set (DatasetV2) – input data set

  • conv_bn_dict (Optional[Dict[Operation, ConvBnInfoType]]) – Dict of conv and bn with activation info. If None, the function looks for it. This can be obtained on the model with bns and convs using BiasCorrection.find_all_convs_bn_with_activation() api.

  • perform_only_empirical_bias_corr (bool) – a flag to indicate only empirical bias correction is to be performed.

Returns

updated session with corrected bias for given ops

Input Parameter Types

Quantization Params

class aimet_tensorflow.bias_correction.QuantParams(quant_mode='tf_enhanced', round_mode='nearest', use_cuda=True, ops_to_ignore=None)

Quant Params to be passed in by user

Constructor

Parameters

  • quant_mode – Indicates which quantization algorithm should be used, either ‘tf’ or ‘tf_enhanced’. Defaults to ‘tf_enhanced’

  • round_mode – The round scheme to used. One of: ‘nearest’ or ‘stochastic’. Default is ‘nearest’.

  • use_cuda – flag to indicate if GPU is to be used

  • ops_to_ignore – ops to be ignored

Bias Correction Params

aimet_tensorflow.bias_correction.BiasCorrectionParams(batch_size: int, num_quant_samples: int, num_bias_correct_samples: int, input_op_names: List[str], output_op_names: List[str])

Input for bias correction to be passed by the user

Parameters

  • batch_size – input batch size to be used

  • num_quant_samples – samples to be used for quantization

  • num_bias_correct_samples – samples to be used for bias correction

  • input_op_names – list of input op names of the given model

  • output_op_names – list of output op names of the given model

Data Input Type

Format expected is tf.Data.DataSet type

Dataset represents the input data as a tensor or nested structures of tensors, one per input to model along with an iterator that operates on them.

Code Examples for Bias Correction

Required imports

# tensorflow
import tensorflow as tf

from tensorflow.keras.applications.resnet50 import ResNet50

# Cross layer Equalization related imports
from aimet_tensorflow.cross_layer_equalization import equalize_model
from aimet_tensorflow.cross_layer_equalization import GraphSearchUtils, CrossLayerScaling, HighBiasFold
from aimet_tensorflow.batch_norm_fold import fold_all_batch_norms
from aimet_tensorflow.batch_norm_fold import fold_given_batch_norms
from aimet_tensorflow.utils.graph_saver import save_and_load_graph
from aimet_tensorflow.utils.op.conv import BiasUtils

# Bias correction related imports
from aimet_tensorflow.bias_correction import BiasCorrectionParams, QuantParams, BiasCorrection

Only Empirical Bias correction on a given model

def bias_correction_empirical(dataset: tf.data.Dataset):
    """
    Perform bias correction on a given model
    :param dataset: Data passed by user as tf.Dataset type.
    :return: None
    """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(batch_size=1,
                                                  num_quant_samples=10,
                                                  num_bias_correct_samples=10,
                                                  input_op_names=['input_1'],
                                                  output_op_names=['fc1000/Softmax'])

    with sess.as_default():
        # run bias correction on the model
        _new_session = BiasCorrection.correct_bias(sess, bias_correction_params, quant_params, dataset)
    sess.close()

Empirical and Analytical Bias correction on a given model

def bias_correction_empirical_analytical(dataset: tf.data.Dataset):
    """
    Perform bias correction on a given model (mix of empirical and analytical)
    :param dataset: Data passed by user as tf.Dataset type.
    :return: None
    """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(batch_size=1,
                                                  num_quant_samples=10,
                                                  num_bias_correct_samples=10,
                                                  input_op_names=['input_1'],
                                                  output_op_names=['fc1000/Softmax'])

    with sess.as_default():
        # run empirical and analytical bias correction on the model
        _new_session = BiasCorrection.correct_bias(sess, bias_correction_params, quant_params,
                                                   dataset,
                                                   perform_only_empirical_bias_corr=False)
    sess.close()

Empirical and Analytical Bias correction on a given model after performing CLE

def bias_correction_after_cle(dataset: tf.data.Dataset):
    """
    Perform bias correction on a given model (mix of empirical and analytical) after
    cross layer equalization.
    :param dataset: Data passed by user as tf.Dataset type.
    :return: None
    """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(batch_size=1,
                                                  num_quant_samples=10,
                                                  num_bias_correct_samples=10,
                                                  input_op_names=['input_1'],
                                                  output_op_names=['fc1000/Softmax'])

    with sess.as_default():

        # store conv bns info before performing CLE
        conv_bn_dict = BiasCorrection.find_all_convs_bn_with_activation(sess,
                                                                        start_op_names=['input_1'],
                                                                        output_op_names=['fc1000/Softmax'])

        # perform CLE
        sess_after_cle = equalize_model(sess, start_op_names=['input_1'], output_op_names=['fc1000/Softmax'])

        # run empirical and analytical bias correction on the model
        _new_session = BiasCorrection.correct_bias(sess_after_cle, bias_correction_params, quant_params,
                                                   dataset,
                                                   conv_bn_dict=conv_bn_dict,
                                                   perform_only_empirical_bias_corr=False)
    sess.close()

Bias Correction Per Layer API

Empirical/ analytical Bias correction can also be performed on a subset of selected layers in a given model using the api listed below.

aimet_tensorflow.bias_correction.BiasCorrection.bias_correction_per_layer(reference_model, corrected_model, bias_correct_params, layer_name_to_be_corrected, data_set)

Helper function to perform empirical bias correction per layer.

Parameters

  • reference_model (Session) – active tensorflow session for reference model

  • corrected_model (Session) – active tensorflow session for corrected model

  • bias_correct_params (BiasCorrectionParams) – bias correction params

  • layer_name_to_be_corrected (str) – name of layer on which bias correction is to be performed

  • quant_params – Quantization specific params from user

Return type

Session

Returns

None, updates corrected model in-place.

aimet_tensorflow.bias_correction.BiasCorrection.analytical_bias_correction_per_layer(corrected_model, layer, preceeding_bn_layer_info, quant_params, is_first_conv=False)

Perform bn based bias correction (analytical bc).

Parameters

  • corrected_model (Session) – active tensorflow session for corrected model

  • layer (Operation) – conv/linear layer to be corrected

  • preceeding_bn_layer_info (ConvBnInfoType) – corresponding preceeding bn/ activation info

  • quant_params (QuantParams) – Quantization specific params from user

  • is_first_conv (bool) – flag to indicate if it’s the first conv layer

Return type

Session

Returns

None, updates corrected_model in place

Code Example for Per-Layer Bias Correction

Empirical Bias correction on one layer

def bias_correction_single_layer_empirical(dataset: tf.data.Dataset):
    """ perform bias correction on one layer """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    bias_correction_params = BiasCorrectionParams(batch_size=1,
                                                  num_quant_samples=10,
                                                  num_bias_correct_samples=10,
                                                  input_op_names=['input_1'],
                                                  output_op_names=['fc1000/Softmax'])

    with sess.as_default():
        # initialize model with zero bias
        sess = BiasUtils.initialize_model_with_bias(sess, bias_correction_params.input_op_names,
                                                    bias_correction_params.output_op_names)

        # pick a layer for bias correction
        example_conv_layer = sess.graph.get_operation_by_name('res2a_branch2a/Conv2D')

        # invoke bias correction of one layer
        BiasCorrection.bias_correction_per_layer(reference_model=sess,
                                                 corrected_model=sess,
                                                 bias_correct_params=bias_correction_params,
                                                 layer_name_to_be_corrected=example_conv_layer.name,
                                                 quant_params=quant_params,
                                                 data_set=dataset)
    sess.close()

Analytical Bias correction on one layer

def bias_correction_single_layer_analytical():
    """ perform analytical bias correction on one layer """

    # load a model
    tf.keras.backend.clear_session()
    _ = ResNet50(weights='imagenet', input_shape=(224, 224, 3))
    sess = tf.compat.v1.keras.backend.get_session()

    # input parameters for bias correction
    # populate required parameters in two data types QuantParams and BiasCorrectParams

    quant_params = QuantParams(quant_mode='tf_enhanced',
                               round_mode='nearest',
                               use_cuda=True,
                               ops_to_ignore=None)

    with sess.as_default():
        # initialize model with zero bias
        sess = BiasUtils.initialize_model_with_bias(sess, ['input_1'], ['fc1000/Softmax'])

        # pick a layer for bias correction
        example_conv_layer = sess.graph.get_operation_by_name('res2a_branch2a/Conv2D')

        # get candidate conv bns in the model
        convs_bn_activation_info_dict = BiasCorrection.find_all_convs_bn_with_activation(sess,
                                                                                         ['input_1'],
                                                                                         ['fc1000/Softmax'])

        # make sure to pick example_conv_layer that has a bn op associated with it
        if example_conv_layer in convs_bn_activation_info_dict.keys():

            preceding_bn_layer_info = convs_bn_activation_info_dict[example_conv_layer]

            # invoke analytical bias correction on this layer
            BiasCorrection.analytical_bias_correction_per_layer(sess,
                                                                example_conv_layer,
                                                                preceding_bn_layer_info,
                                                                quant_params)
    sess.close()