卷积层和BN层融合

桃扇骨 2023-01-22 12:53 296阅读 0赞

解释也不错：

[https://my.oschina.net/u/4395251/blog/4769266][https_my.oschina.net_u_4395251_blog_4769266]

跟博士请教，分组卷积可以合并，如果是独立卷积，bn是通道的bn，可能不能合并？

1.  为什么要合并BN层

在训练深度网络模型时，BN（Batch Normalization）层能够加速网络收敛，并且能够控制过拟合，一般放在卷积层之后。BN 层将数据归一化后，能够有效解决梯度消失与梯度爆炸问题。虽然 BN 层在训练时起到了积极作用，然而，在网络前向推断时多了一些层的运算，影响了模型的性能，且占用了更多的内存或者显存空间。目前，很多先进的网络模型（ResNet，MobileNet，Xception，ShuffleNet 等）都使用了BN技术，因此，我们有必要将 BN 层的参数合并到卷积层，来提升模型前向推断的速度。

２.  BN层与卷积层合并的数学原理

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zODc0MDQ2Mw_size_16_color_FFFFFF_t_70][]

则有：

![20190925144712475.png][]

合并后：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zODc0MDQ2Mw_size_16_color_FFFFFF_t_70 1][]

3.  实验结果

机器：显卡 GTX 1080Ti，i7 CPU

本实验对比了Resnet50 模型合并BN层前后的性能，分类精度保持不变，速度显著提升。

模型    CPU前向时间    GPU前向时间  
Resnet50（合并前）    176.17ms    11.03ms  
Resnet50（合并后）    161.69ms    7.3ms  
提升    8.96%    33.27%  
————————————————  
版权声明：本文为CSDN博主「小麦草」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。  
原文链接：https://blog.csdn.net/kangdi7547/article/details/81348254

[https://github.com/vietnamican/conv-bn-merge/blob/main/convbnmerge/convbnmerge.py][https_github.com_vietnamican_conv-bn-merge_blob_main_convbnmerge_convbnmerge.py]

合并代码1：

[https://blog.csdn.net/zym19941119/article/details/84640433][https_blog.csdn.net_zym19941119_article_details_84640433]

import torch
    import os
    from collections import OrderedDict
    import cv2
    import numpy as np
    import torchvision.transforms as transforms
    
    
    """  Parameters and variables  """
    IMAGENET = '/home/zym/ImageNet/ILSVRC2012_img_val_256xN_list.txt'
    LABEL = '/home/zym/ImageNet/synset.txt'
    TEST_ITER = 10
    SAVE = False
    TEST_AFTER_MERGE = True
    
    
    """  Functions  """
    def merge(params, name, layer):
        # global variables
        global weights, bias
        global bn_param
    
        if layer == 'Convolution':
            # save weights and bias when meet conv layer
            if 'weight' in name:
                weights = params.data
                bias = torch.zeros(weights.size()[0])
            elif 'bias' in name:
                bias = params.data
            bn_param = {}
    
        elif layer == 'BatchNorm':
            # save bn params
            bn_param[name.split('.')[-1]] = params.data
    
            # running_var is the last bn param in pytorch
            if 'running_var' in name:
                # let us merge bn ~
                tmp = bn_param['weight'] / torch.sqrt(bn_param['running_var'] + 1e-5)
                weights = tmp.view(tmp.size()[0], 1, 1, 1) * weights
                bias = tmp*(bias - bn_param['running_mean']) + bn_param['bias']
    
                return weights, bias
    
        return None, None
    
    
    """  Main functions  """
    # import pytorch model
    import models.shufflenetv2.shufflenetv2_merge as shufflenetv2
    pytorch_net = shufflenetv2.ShuffleNetV2().eval()
    model_path = shufflenetv2.weight_file
    
    # load weights
    print('Finding trained model weights...')
    try:
        for file in os.listdir(model_path):
            if 'pth' in file:
                print('Loading weights from %s ...' % file)
                trained_weights = torch.load(os.path.join(model_path, file))
                # pytorch_net.load_state_dict(trained_weights)
                print('Weights load success')
                break
    except:
        raise ValueError('No trained model found or loading error occurs')
    
    # go through pytorch net
    print('Going through pytorch net weights...')
    new_weights = OrderedDict()
    inner_product_flag = False
    for name, params in trained_weights.items():
        if len(params.size()) == 4:
            _, _ = merge(params, name, 'Convolution')
            prev_layer = name
        elif len(params.size()) == 1 and not inner_product_flag:
            w, b = merge(params, name, 'BatchNorm')
            if w is not None:
                new_weights[prev_layer] = w
                new_weights[prev_layer.replace('weight', 'bias')] = b
        else:
            # inner product layer
            # if meet inner product layer,
            # the next bias weight can be misclassified as 'BatchNorm' layer as len(params.size()) == 1
            new_weights[name] = params
            inner_product_flag = True
    
    # align names in new_weights with pytorch model
    # after move BatchNorm layer in pytorch model,
    # the layer names between old model and new model will mis-align
    print('Aligning weight names...')
    pytorch_net_key_list = list(pytorch_net.state_dict().keys())
    new_weights_key_list = list(new_weights.keys())
    assert len(pytorch_net_key_list) == len(new_weights_key_list)
    for index in range(len(pytorch_net_key_list)):
        new_weights[pytorch_net_key_list[index]] = new_weights.pop(new_weights_key_list[index])
    
    # save new weights
    if SAVE:
        torch.save(new_weights, model_path + '/' + file.replace('.pth', '_merged.pth'))
    
    # test merged pytorch model
    if TEST_AFTER_MERGE:
        try:
            pytorch_net.load_state_dict(new_weights)
            print('Pytorch net load weights success~')
        except:
            raise ValueError('Load new weights error')
    
        print('-' * 50)
        with open(LABEL) as f:
            labels = f.read().splitlines()
        with open(IMAGENET) as f:
            images = f.read().splitlines()
            for _ in range(TEST_ITER):
                # cv2 default chann el is BGR
                image_path, label = images[np.random.randint(0, len(images))].split(' ')
                # image_path, label = images[0].split(' ')
                input_image = cv2.imread(image_path)
                input_image = cv2.resize(input_image, (224, 224))
                input_image = transforms.Compose([transforms.ToTensor(),
                                                  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                                       std=[0.229, 0.224, 0.225])
                                                  ])(input_image)
                input_image = input_image.view(1, 3, 224, 224)
                output_logits = pytorch_net(input_image)
                _, index = output_logits.max(dim=1)
                print('true label: \t%s' % labels[int(label)])
                print('predict label:\t%s' % labels[int(index)])
                print('-' * 50)

合并代码2：

[https://github.com/owphoo/pytorch\_merge\_bn/blob/master/pytorch\_merge\_bn.py][https_github.com_owphoo_pytorch_merge_bn_blob_master_pytorch_merge_bn.py]

import torch
    import os
    from collections import OrderedDict
    import numpy as np
    
    global merged
    merged = True
    def merge(params, name, layer, deconv_layer_names=['deconv']):
        # global variables
        global weights, bias
        global bn_param
        global merged
    
        is_deconv = False
        for deconv_name in deconv_layer_names:
            if deconv_name in name:
                is_deconv = True
                break 
        if layer == 'Convolution':
            # save weights and bias when meet conv layer
            if 'weight' in name:
                weights = params.data
                bias = torch.zeros(weights.size()[0], device=weights.device)
                if is_deconv:
                    bias = torch.zeros(weights.size()[1], device=weights.device)
                else:
                    bias = torch.zeros(weights.size()[0], device=weights.device)
                merged = False
            elif 'bias' in name:
                bias = params.data
            bn_param = {}
    
        elif layer == 'BatchNorm':
            # save bn params
            bn_param[name.split('.')[-1]] = params.data
    
            # running_var is the last bn param in pytorch
            if 'running_var' in name:
                # merge bn
                tmp = bn_param['weight'] / torch.sqrt(bn_param['running_var'] + 1e-5)
                if is_deconv:
                    weights = (tmp.view(tmp.size()[0], 1, 1, 1) * weights.permute(1,0,2,3)).permute(1,0,2,3)
                else:
                    weights = tmp.view(tmp.size()[0], 1, 1, 1) * weights
                bias = tmp * (bias - bn_param['running_mean']) + bn_param['bias']
    
                return weights, bias
    
        return None, None
    
    
    import sys
    if __name__ == '__main__':
        if len(sys.argv) != 2:
            print('Usage: python pytorch_merge_bn.py YOU_MODEL')
            sys.exit(-1)
        model_path = sys.argv[1]
        print('input model: ', model_path)
        checkpoint = torch.load(model_path)
    
        trained_weights = checkpoint['net_state_dict']
    
        '''
        ## conv_bn_relu module
        #           NAME           |           SIZE
        #   conv4.0.weight              torch.Size([128, 256, 3, 3])
        #   conv4.1.weight              torch.Size([256])
        #   conv4.1.bias                torch.Size([256])
        #   conv4.1.running_mean        torch.Size([256])
        #   conv4.1.running_var         torch.Size([256])
        ## deconv_bn_relu module
        #           NAME           |           SIZE
        #   deconv4.0.weight             torch.Size([256, 128, 4, 4])
        #   deconv4.1.weight             torch.Size([128])
        #   deconv4.1.bias               torch.Size([128])
        #   deconv4.1.running_mean       torch.Size([128])
        #   deconv4.1.running_var        torch.Size([128])
        '''
    
        # check it in your net modules
        deconv_layer_names = ['deconv4', 'deconv3', 'deconv2', 'deconv1']
        temp = []
        for deconv_name in deconv_layer_names:
            temp.append(deconv_name + '.0')
            temp.append(deconv_name + '.1')
        deconv_layer_names = temp
    
        # go through pytorch net
        new_weights = OrderedDict()
        inner_product_flag = False
        for name, params in trained_weights.items():
            print ('name: ', name, params.size())
            if len(params.size()) == 4:
                _, _ = merge(params, name, 'Convolution', deconv_layer_names=deconv_layer_names)
                prev_layer = name
                # print('prev1: ', prev_layer)
            elif len(params.size()) == 1 and not inner_product_flag:
                w, b = merge(params, name, 'BatchNorm', deconv_layer_names=deconv_layer_names)
                # print('prev2: ', prev_layer)
                if w is not None:
                    new_weights[prev_layer] = w
                    new_weights[prev_layer.replace('weight', 'bias')] = b
                # mergebn
                merged = True
            else:
                # inner product layer (TODO, inner product layer may have bn module)
                if name.find('num_batches_tracked') == -1:
                    new_weights[name] = params
                    inner_product_flag = True
                else:
                    pass
    
        # for the last conv/deconv if it has no bn module 
        if merged is False:
            new_weights[prev_layer] = weights
            new_weights[prev_layer.replace('weight', 'bias')] = bias
    
        checkpoint['net_state_dict'] = new_weights
    
        # save new weights
        model_name = model_path[model_path.rfind('/')+1:]
        model_path = model_path[:model_path.rfind('/')]
        if model_path.find('/') == -1:
            model_path = './'
        torch.save(checkpoint, model_path + '/merge_bn_' + model_name)

[https_my.oschina.net_u_4395251_blog_4769266]: https://my.oschina.net/u/4395251/blog/4769266
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zODc0MDQ2Mw_size_16_color_FFFFFF_t_70]: /images/20221020/7e7f6ed17eda472ca16fc3cbadb245ad.png
[20190925144712475.png]: /images/20221020/617a01c6cc084addb5500ce17927ee6e.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dlaXhpbl8zODc0MDQ2Mw_size_16_color_FFFFFF_t_70 1]: /images/20221020/230dd63f7c7141c18a03c9ebfeec7886.png
[https_github.com_vietnamican_conv-bn-merge_blob_main_convbnmerge_convbnmerge.py]: https://github.com/vietnamican/conv-bn-merge/blob/main/convbnmerge/convbnmerge.py
[https_blog.csdn.net_zym19941119_article_details_84640433]: https://blog.csdn.net/zym19941119/article/details/84640433
[https_github.com_owphoo_pytorch_merge_bn_blob_master_pytorch_merge_bn.py]: https://github.com/owphoo/pytorch_merge_bn/blob/master/pytorch_merge_bn.py