TrainingExtensions/tensorflow/src/python/aimet_tensorflow/keras/batch_norm_fold.py
def _fold_selected_batch_norms(
        model: (tf.keras.Model, tf.keras.layers.Layer), layer_pairs: List[PAIR_TYPE]):
    """
    Fold a given set of batch_norm layers into conv layers

    :param model: model to fold selected batchnorm for
    :param layer_pairs: Tuple of conv, bn layers and is_batch_norm_second flag
    """

    list_of_bn_layers = []
    for pair in layer_pairs:

        conv_linear, batchnorm, is_batch_norm_second = pair
        assert isinstance(
            conv_linear,
            (tf.keras.layers.Conv2D, tf.keras.layers.Dense, tf.keras.layers.DepthwiseConv2D))

        list_of_bn_layers.append(batchnorm)

        #  check flag
        is_bias_valid = False

        if conv_linear.bias is not None:
            is_bias_valid = True

        bn_params = _get_bn_params(batchnorm)
        weight_tensor = _get_weight_tensor_transpose_reshape(conv_linear)
        bias_tensor = _get_bias_tensor(conv_linear)

        # Updated weight and bias
        bias = libpymo.fold(bn_params, weight_tensor, bias_tensor, is_bias_valid,
                            is_batch_norm_second)

        if isinstance(conv_linear, tf.keras.layers.DepthwiseConv2D):
            # Depthwise conv layers in TF have outputs(Noc) set to 1.
            # we send in format [Nic, Noc, kh, kw]
            numpy_weight_reshaped = np.reshape(weight_tensor.data, weight_tensor.shape)\
                .transpose((2, 3, 0, 1))
        elif isinstance(conv_linear, tf.keras.layers.Dense):
            # o, i - convert to i , o
            numpy_weight_reshaped = np.reshape(
                weight_tensor.data,
                [weight_tensor.shape[0], weight_tensor.shape[1]]).transpose(1, 0)
        else:
            # conv2D case
            # we sent in format [Noc, Nic, kh, kw]
            numpy_weight_reshaped = np.reshape(weight_tensor.data, weight_tensor.shape)\
                .transpose((2, 3, 1, 0))

        # update bias tensor, even in case there was no existing bias add op in given conv2D op.
        bias_tensor_shape = [weight_tensor.shape[0]]
        numpy_bias_reshaped = np.reshape(bias, bias_tensor_shape)
        conv_linear.set_weights([numpy_weight_reshaped.data, numpy_bias_reshaped])

        BNUtils.modify_bn_params_to_make_as_passthrough(batchnorm)

    _delete_all_bns_from_model(model, list_of_bn_layers)

    Fold a given set of batch_norm layers into conv layers

    :param model: model to fold selected batchnorm for
    :param layer_pairs: Tuple of conv, bn layers and is_batch_norm_second flag
    
Fold a given set of batch_norm layers into conv layers
Fold a given set of batch_norm layers into conv layers
def _fold_selected_batch_norms(
        model: (tf.keras.Model, tf.keras.layers.Layer), layer_pairs: List[PAIR_TYPE]):
    list_of_bn_layers = []
    for pair in layer_pairs:
        conv_linear, batchnorm, is_batch_norm_second = pair
        assert isinstance(
            conv_linear,
            (tf.keras.layers.Conv2D, tf.keras.layers.Dense, tf.keras.layers.DepthwiseConv2D))
        list_of_bn_layers.append(batchnorm)
        is_bias_valid = False
        if conv_linear.bias is not None:
            is_bias_valid = True
        bn_params = _get_bn_params(batchnorm)
        weight_tensor = _get_weight_tensor_transpose_reshape(conv_linear)
        bias_tensor = _get_bias_tensor(conv_linear)
        bias = libpymo.fold(bn_params, weight_tensor, bias_tensor, is_bias_valid,
                            is_batch_norm_second)
        if isinstance(conv_linear, tf.keras.layers.DepthwiseConv2D):
            numpy_weight_reshaped = np.reshape(weight_tensor.data, weight_tensor.shape)\
                .transpose((2, 3, 0, 1))
        elif isinstance(conv_linear, tf.keras.layers.Dense):
            numpy_weight_reshaped = np.reshape(
                weight_tensor.data,
                [weight_tensor.shape[0], weight_tensor.shape[1]]).transpose(1, 0)
        else:
            numpy_weight_reshaped = np.reshape(weight_tensor.data, weight_tensor.shape)\
                .transpose((2, 3, 1, 0))
        bias_tensor_shape = [weight_tensor.shape[0]]
        numpy_bias_reshaped = np.reshape(bias, bias_tensor_shape)
        conv_linear.set_weights([numpy_weight_reshaped.data, numpy_bias_reshaped])
        BNUtils.modify_bn_params_to_make_as_passthrough(batchnorm)
    _delete_all_bns_from_model(model, list_of_bn_layers)
def _fold_selected_batch_norms ( model : ( tf . keras . Model , tf . keras . layers . Layer ) , layer_pairs : List [ PAIR_TYPE ] ) : list_of_bn_layers = [ ] for pair in layer_pairs : conv_linear , batchnorm , is_batch_norm_second = pair assert isinstance ( conv_linear , ( tf . keras . layers . Conv2D , tf . keras . layers . Dense , tf . keras . layers . DepthwiseConv2D ) ) list_of_bn_layers . append ( batchnorm ) is_bias_valid = False if conv_linear . bias is not None : is_bias_valid = True bn_params = _get_bn_params ( batchnorm ) weight_tensor = _get_weight_tensor_transpose_reshape ( conv_linear ) bias_tensor = _get_bias_tensor ( conv_linear ) bias = libpymo . fold ( bn_params , weight_tensor , bias_tensor , is_bias_valid , is_batch_norm_second ) if isinstance ( conv_linear , tf . keras . layers . DepthwiseConv2D ) : numpy_weight_reshaped = np . reshape ( weight_tensor . data , weight_tensor . shape ) . transpose ( ( 2 , 3 , 0 , 1 ) ) elif isinstance ( conv_linear , tf . keras . layers . Dense ) : numpy_weight_reshaped = np . reshape ( weight_tensor . data , [ weight_tensor . shape [ 0 ] , weight_tensor . shape [ 1 ] ] ) . transpose ( 1 , 0 ) else : numpy_weight_reshaped = np . reshape ( weight_tensor . data , weight_tensor . shape ) . transpose ( ( 2 , 3 , 1 , 0 ) ) bias_tensor_shape = [ weight_tensor . shape [ 0 ] ] numpy_bias_reshaped = np . reshape ( bias , bias_tensor_shape ) conv_linear . set_weights ( [ numpy_weight_reshaped . data , numpy_bias_reshaped ] ) BNUtils . modify_bn_params_to_make_as_passthrough ( batchnorm ) _delete_all_bns_from_model ( model , list_of_bn_layers )
Fold a given set of batch_norm layers into conv layers
Fold a given set of batch_norm layers into conv layers
"""
    Fold a given set of batch_norm layers into conv layers

    :param model: model to fold selected batchnorm for
    :param layer_pairs: Tuple of conv, bn layers and is_batch_norm_second flag
    """ #  check flag # Updated weight and bias # Depthwise conv layers in TF have outputs(Noc) set to 1. # we send in format [Nic, Noc, kh, kw] # o, i - convert to i , o # conv2D case # we sent in format [Noc, Nic, kh, kw] # update bias tensor, even in case there was no existing bias add op in given conv2D op.
model
(tf.keras.Model, tf.keras.layers.Layer)
model to fold selected batchnorm for
model to fold selected batchnorm for
layer_pairs
List[PAIR_TYPE]
Tuple of conv, bn layers and is_batch_norm_second flag
Tuple of conv bn layers and is_batch_norm_second flag