TrainingExtensions/torch/src/python/aimet_torch/cross_layer_equalization.py
def bias_fold(cls_set_info_list: List[ClsSetInfo], bn_layers: Dict[Union[torch.nn.Conv2d, torch.nn.ConvTranspose2d],
                                                                       torch.nn.BatchNorm2d]):
        """
        Folds bias values greater than 3 * sigma to next layer's bias

        :param cls_set_info_list: List of info elements for each cls set
        :param bn_layers: Key: Conv/Linear layer Value: Corresponding folded BN layer
        :return: None
        """
        if not bn_layers:
            logger.info('High Bias folding is not supported for models without BatchNorm Layers')
            return

        for cls_set_info in cls_set_info_list:

            for cls_pair_info in cls_set_info.cls_pair_info_list:

                if (cls_pair_info.layer1.bias is None) or (cls_pair_info.layer2.bias is None) or \
                        (cls_pair_info.layer1 not in bn_layers):
                    continue

                prev_layer_params, curr_layer_params = HighBiasFold.call_mo_high_bias_fold(cls_pair_info, bn_layers)

                prev_layer_bias_shape = cls_pair_info.layer1.weight.shape[0]
                # Transpose weight back to N, C, H, W for transposed Conv2D
                if isinstance(cls_pair_info.layer1, torch.nn.ConvTranspose2d) and cls_pair_info.layer1.groups == 1:
                    prev_layer_bias_shape = cls_pair_info.layer1.weight.shape[1]

                cls_pair_info.layer1.bias.data = torch.from_numpy(np.reshape(prev_layer_params.bias,
                                                                             prev_layer_bias_shape))
                cls_pair_info.layer1.bias.data = cls_pair_info.layer1.bias.data.type(torch.FloatTensor)
                cls_pair_info.layer2.bias.data = torch.from_numpy(np.reshape(curr_layer_params.bias,
                                                                             curr_layer_params.weightShape[0]))
                cls_pair_info.layer2.bias.data = cls_pair_info.layer2.bias.data.type(torch.FloatTensor)

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

        :param cls_set_info_list: List of info elements for each cls set
        :param bn_layers: Key: Conv/Linear layer Value: Corresponding folded BN layer
        :return: None
        
Folds bias values greater than 3 * sigma to next layer's bias
Folds bias values greater than 3 * sigma to next layer ' s bias
def bias_fold(cls_set_info_list: List[ClsSetInfo], bn_layers: Dict[Union[torch.nn.Conv2d, torch.nn.ConvTranspose2d],
                                                                       torch.nn.BatchNorm2d]):
        if not bn_layers:
            logger.info('High Bias folding is not supported for models without BatchNorm Layers')
            return
        for cls_set_info in cls_set_info_list:
            for cls_pair_info in cls_set_info.cls_pair_info_list:
                if (cls_pair_info.layer1.bias is None) or (cls_pair_info.layer2.bias is None) or \
                        (cls_pair_info.layer1 not in bn_layers):
                    continue
                prev_layer_params, curr_layer_params = HighBiasFold.call_mo_high_bias_fold(cls_pair_info, bn_layers)
                prev_layer_bias_shape = cls_pair_info.layer1.weight.shape[0]
                if isinstance(cls_pair_info.layer1, torch.nn.ConvTranspose2d) and cls_pair_info.layer1.groups == 1:
                    prev_layer_bias_shape = cls_pair_info.layer1.weight.shape[1]
                cls_pair_info.layer1.bias.data = torch.from_numpy(np.reshape(prev_layer_params.bias,
                                                                             prev_layer_bias_shape))
                cls_pair_info.layer1.bias.data = cls_pair_info.layer1.bias.data.type(torch.FloatTensor)
                cls_pair_info.layer2.bias.data = torch.from_numpy(np.reshape(curr_layer_params.bias,
                                                                             curr_layer_params.weightShape[0]))
                cls_pair_info.layer2.bias.data = cls_pair_info.layer2.bias.data.type(torch.FloatTensor)
def bias_fold ( cls_set_info_list : List [ ClsSetInfo ] , bn_layers : Dict [ Union [ torch . nn . Conv2d , torch . nn . ConvTranspose2d ] , torch . nn . BatchNorm2d ] ) : if not bn_layers : logger . info ( 'High Bias folding is not supported for models without BatchNorm Layers' ) return for cls_set_info in cls_set_info_list : for cls_pair_info in cls_set_info . cls_pair_info_list : if ( cls_pair_info . layer1 . bias is None ) or ( cls_pair_info . layer2 . bias is None ) or ( cls_pair_info . layer1 not in bn_layers ) : continue prev_layer_params , curr_layer_params = HighBiasFold . call_mo_high_bias_fold ( cls_pair_info , bn_layers ) prev_layer_bias_shape = cls_pair_info . layer1 . weight . shape [ 0 ] if isinstance ( cls_pair_info . layer1 , torch . nn . ConvTranspose2d ) and cls_pair_info . layer1 . groups == 1 : prev_layer_bias_shape = cls_pair_info . layer1 . weight . shape [ 1 ] cls_pair_info . layer1 . bias . data = torch . from_numpy ( np . reshape ( prev_layer_params . bias , prev_layer_bias_shape ) ) cls_pair_info . layer1 . bias . data = cls_pair_info . layer1 . bias . data . type ( torch . FloatTensor ) cls_pair_info . layer2 . bias . data = torch . from_numpy ( np . reshape ( curr_layer_params . bias , curr_layer_params . weightShape [ 0 ] ) ) cls_pair_info . layer2 . bias . data = cls_pair_info . layer2 . bias . data . type ( torch . FloatTensor )
Folds bias values greater than 3 * sigma to next layer's bias
Folds bias values greater than 3 * sigma to next layer ' s bias
"""
        Folds bias values greater than 3 * sigma to next layer's bias

        :param cls_set_info_list: List of info elements for each cls set
        :param bn_layers: Key: Conv/Linear layer Value: Corresponding folded BN layer
        :return: None
        """ # Transpose weight back to N, C, H, W for transposed Conv2D
None
cls_set_info_list
List[ClsSetInfo]
List of info elements for each cls set
List of info elements for each cls set
bn_layers
Dict[Union[torch.nn.Conv2d, torch.nn.ConvTranspose2d],
                                                                       torch.nn.BatchNorm2d]
Conv/Linear layer Value: Corresponding folded BN layer
Conv / Linear layer Value : Corresponding folded BN layer