caffe加速:合并BatchNorm层和Scale层到Convolution层

Convolution+BatchNorm+Scale+Relu的组合模块在卷积后进行归一化，可以加速训练收敛。但在推理时BatchNorm非常耗时，可以将训练时学习到的BatchNorm+Scale的线性变换参数融合到卷积层，替换原来的Convolution层中weights和bias，实现在不影响准确度的前提下加速预测时间。

# merge bn & scale into convolution and remove dropout layer
# this is a modified version from https://github.com/sanghoon/pva-faster-rcnn/blob/master/tools/gen_merged_model.py which further remove dropout layer

import numpy as np
import sys
import os
import os.path as osp
import google.protobuf as pb
from argparse import ArgumentParser
import sys
caffe_root = ‘/home/yly/ENet/caffe-enet/‘
sys.path.insert（0 caffe_root + ‘python‘） 
import caffe
    
    
def load_and_fill_biases（src_model src_weights dst_model dst_weights）:
    with open（src_model） as f:
        model = caffe.proto.caffe_pb2.NetParameter（）
        pb.text_format.Merge（f.read（） model）
    
    for i layer in enumerate（model.layer）:
        if layer.type == ‘Convolution‘: # or layer.type == ‘Scale‘:
            # Add bias layer if needed
            if layer.convolution_param.bias_term == False:
                layer.convolution_param.bias_term = True
                layer.convolution_param.bias_filler.type = ‘constant‘
                layer.convolution_param.bias_filler.value = 0.0
    
    with open（dst_model ‘w‘） as f:
        f.write（pb.text_format.MessageToString（model））
    
    #caffe.set_mode_gpu（）
    net_src = caffe.Net（src_model src_weights caffe.TEST）
    net_dst = caffe.Net（dst_model caffe.TEST）
    for key in net_src.params.keys（）:
        for i in range（len（net_src.params[key]））:
            net_dst.params[key][i].data[:] = net_src.params[key][i].data[:]
    
    if dst_weights is not None:
        # Store params
        pass
    
    return net_dst
    
    
def merge_conv_and_bn（net i_conv i_bn i_scale）:
    # This is based on Kyeheyon‘s work
    assert（i_conv != None）
    assert（i_bn != None）
    
    def copy_double（data）:
        return np.array（data copy=True dtype=np.double）
    
    key_conv = net._layer_names[i_conv]
    key_bn = net._layer_names[i_bn]
    key_scale = net._layer_names[i_scale] if i_scale else None
    
    # Copy
    bn_mean = copy_double（net.params[key_bn][0].data）
    bn_variance = copy_double（net.params[key_bn][1].data）
    num_bn_samples = copy_double（net.params[key_bn][2].data）
    
    # and Invalidate the BN layer
    net.params[key_bn][0].data[:] = 0
    net.params[key_bn][1].data[:] = 1
    net.params[key_bn][2].data[:] = 1
    if num_bn_samples[0] == 0:
        num_bn_samples[0] = 1
    
    if net.params.has_key（key_scale）:
        print ‘Combine {:s} + {:s} + {:s}‘.format（key_conv key_bn key_scale）
        scale_weight = copy_double（net.params[key_scale][0].data）
        scale_bias = copy_double（net.params[key_scale][1].data）
        net.params[key_scale][0].data[:] = 1
        net.params[key_scale][1].data[:] = 0
    else:
        print ‘Combine {:s} + {:s}‘.format（key_conv key_bn）
        scale_weight = 1
        scale_bias = 0
    
    weight = copy_double（net.params[key_conv][0].data）
    bias = copy_double（net.params[key_conv][1].data）
    alpha = scale_weight / np.sqrt（bn_variance / num_bn_samples[0] + np.