shrink_bn_caffe.py 文件源码

def generate_nbn_caffemodel(input_prototxt, input_caffemodel, output_prototxt, output_caffemodel, eps = 0.00001):
    input_network = caffe.Net(input_prototxt, input_caffemodel, caffe.TEST)

    f = open(input_caffemodel, 'rb')
    tmp_model = caffe_pb2.NetParameter()
    tmp_model.ParseFromString(f.read())
    f.close()
    layers = tmp_model.layer

    output_network = caffe.Net(output_prototxt, caffe.TEST)

    for i in range(len(layers)):
        if layers[i].type == "Input" or layers[i].type == "Eltwise" or layers[i].type == "Scale" or layers[i].type == "BatchNorm" or layers[i].type == "ImageData" or layers[i].type == "ReLU" or layers[i].type == "Pooling" or layers[i].type == "Split" or layers[i].type == "Concat" or  layers[i].type == "Flatten" or layers[i].type == "SoftmaxWithLoss":
            continue
        elif layers[i].type == "Convolution":
            if not (layers[i+2].type == "Scale" and layers[i+1].type == "BatchNorm"):
                continue
            bn_conv = layers[i+1].name
            scale_conv = layers[i+2].name
            conv_w = input_network.params[layers[i].name][0].data[...]
            print layers[i].name, layers[i+1].name, layers[i+2].name, layers[i+2].scale_param.bias_term, layers[i].convolution_param.bias_term, conv_w.shape
            if layers[i].convolution_param.bias_term:
                # original conv
                conv_b = input_network.params[layers[i].name][1].data[...]
            else:
                conv_b = np.zeros((conv_w.shape[0],), dtype=np.uint8)

            # original batchnormal
            scale = input_network.params[bn_conv][2].data[...]
            mean = input_network.params[bn_conv][0].data[...]
            var = input_network.params[bn_conv][1].data[...]

            # original scale
            scale_w = input_network.params[scale_conv][0].data[...]
            scale_b = input_network.params[scale_conv][1].data[...]
            #print "scale_w:", scale_w

            # calculate
            var = np.sqrt(var/scale+eps)
            conv_b = conv_b-mean/scale
            conv_b = conv_b/var
            var = scale_w/var
            conv_b = scale_w*conv_b
            conv_b = conv_b + scale_b

            for j in range(len(var)):
                output_network.params[layers[i].name][0].data[j] = var[j]*conv_w[j]
            output_network.params[layers[i].name][1].data[...] = conv_b
        else:
            output_network.params[layers[i].name][0].data[...] = input_network.params[layers[i].name][0].data[...]
            output_network.params[layers[i].name][1].data[...] = input_network.params[layers[i].name][1].data[...]    

    output_network.save(output_caffemodel)