caffe模型通道剪枝channel pruning

秒速五厘米 2022-03-22 03:56 266阅读 0赞

deep compression介绍的剪枝：是将权值置0，再通过稀疏存储格式来减小模型大小。

如下， 通过通道剪枝来减少模型大小。

# coding:utf-8
    # by chen yh
    
    import caffe
    import numpy as np
    import shutil
    import matplotlib.pyplot as plt
    
    '''
    These parameters need modification: 
        root: root directory ;
        model: your caffemodel ; 
        prototxt: your prototxt ; 
    
        prune layer: need prune layer name, a list ;
        input layer: input of these layers is output of prune layer,each element is a list ;
        th : thereshold of each prune layer,a list.
    
    Please ensure lenth of prune layer, input layer and th.
    
    picture and get get_sum_l1 functions can help you find suitable threshold. 
    '''
    
    
    def get_prune(net, layer, threshold):  # 返回layer中低于阈值threshold的卷积核的序号
        weight_ori = net.params[layer][0].data
        # bias_ori=net.params[layer][1].data
        
        print("layer:",layer,"shape:",weight_ori.shape)
        #print("weight_ori:",weight_ori)
        sum_l1 = []
        for i in range(weight_ori.shape[0]):
            if (layer == "fc6") or (layer == "fc7") or (layer == "fc8"):
                sum_l1.append((i, np.sum(abs(weight_ori[i, :]))))
            else:
                sum_l1.append((i, np.sum(abs(weight_ori[i, :, :, :]))))  # sum_l1存放每个卷积核的所有权重绝对值之和
        print("sum_l1.shape:",len(sum_l1))
        de_keral = []  # de_keral存放大于阈值的卷积核的序号
        for i in sum_l1:
            if i[1] > threshold:
                de_keral.append(i[0])
    
        print(layer + "层需要prune的卷积核有" + str(weight_ori.shape[0] - len(de_keral)) + "个,保留的卷积核有" + str(len(de_keral)) + "个")
        print("de_keral.shape:",len(de_keral))
        return de_keral
    
    
    def prune(net, pk, lk):  # 输出两个字典,键都是修剪层的layer的名字,值分别是修剪层的weight和bias
        w_new = {}  # 键是layer,值是保存后的weight
        b_new = {}  # 键是layer,值是保存后的bias
        for l in pk.keys():  # 待剪层权重处理 w_n = w[pk[l],:,;,;]
            #print("pk.keys:",l)
            w_old = net.params[l][0].data
            b_old = net.params[l][1].data
            if (l == "fc6") or (l == "fc7"):
                w_n = w_old[pk[l], :]
            else:
                w_n = w_old[pk[l], :, :, :]
            b_n = b_old[pk[l]]
            w_new[l] = w_n
            b_new[l] = b_n
            # net_n.params[l][0].data[...] = w_n
            # net_n.params[l][1].data[...] = b_n
            print("player:",l,",w.shape:",w_new[l].shape,",b.shape",b_new[l].shape)
        for l in lk.keys():  # 以待剪层为输入的层权重处理
            if l not in pk.keys():  # bottom被修剪后本身没有被修剪,所以其权重只需要在原来的net上面取切片,w_n = w[:,lk[l],:,:]
                #print("lk.keys:", l)
                if (l != "conv4_3_norm"):  # 对传统卷积层的处理
                    w_o = net.params[l][0].data
                    b_o = net.params[l][1].data
                    b_new[l] = b_o  # bias保留,因为这些层没有剪卷积核
                    if (l == "fc6") or (l == "fc7") or (l == "fc8"):
                        w_n = w_o[:, lk[l]]
                    else:
                        w_n = w_o[:, lk[l], :, :]
                    w_new[l] = w_n
                else:  # 对特殊层的处理,参数个数不是2
                    w_o = net.params[l][0].data
                    w_n = w_o[lk[l],]
                    w_new[l] = w_n
                print("klayer:",l,",w.shape:",w_new[l].shape,",b.shape",b_new[l].shape)
    
            else:  # pk 和 lk共有的层,也就是这层的bottom和层本身都被修剪过,所以权重不能在原来的net上切片,利用保存了的w_new取切片.
                w_o = w_new[l]
                print("lk.keys else:",l)
                if (l == "fc6") or (l == "fc7"):
                    w_n = w_o[:, lk[l]]
                else:
                    w_n = w_o[:, lk[l], :, :]
                w_new[l] = w_n
                print("llayer:",l,",w.shape:",w_new[l].shape)
    
        return w_new, b_new
    
    
    def get_prototxt(pk, pro_n):  # 复制原来的prototxt,并修改修剪层的num_output,这一段代码有点绕,有空的话优化为几个单独的函数或者弄个类
        with open(pro_n, "r") as p:
            lines = p.readlines()
        k = 0
        with open(pro_n, "w") as p:
            while k < len(lines):  # 遍历所有的lines,此处不宜用for.
                #print("lines[k]:",lines[k])
                if 'name:' in lines[k]:
                    print("lines[k].split:",lines[k].split('"')[1])
                    l_name = lines[k].split('"')[1]  # 获取layer name
                    if l_name in pk.keys():  # 如果name在待修剪层中,则需要修改,下面进入一个找channel的循环块.
                        while True:
                            if "num_output:" in lines[k]:
                                channel_n = "    num_output: " + str(len(pk[l_name])) + "\n"
                                p.write(channel_n)
                                k = k + 1
                                break
                            else:
                                p.write(lines[k])
                                k = k + 1
                    else:  # name不在待修剪层中,直接copy行
                        p.write(lines[k])
                        k = k + 1
    
                else:
                    p.write(lines[k])
                    k = k + 1
        print("deploy_rebirth_prune.prototxt已写好")
    
    
    def savemodel(net, net_n, w_prune, b_prune, path):  # 储存修改后的caffemodel
        for layer in net.params.keys():
            if layer in w_prune.keys():
                print("save model-layer:",layer,"len:",w_prune[layer].shape,"len2:",net_n.params[layer][0].data[...].shape)
                net_n.params[layer][0].data[...] = w_prune[layer]
                if layer in b_prune.keys():
                    net_n.params[layer][1].data[...] = b_prune[layer]
            else:
                weight = net.params[layer]
                for index, w in enumerate(weight):
                    try:
                        net_n.params[layer][index].data[...] = w.data
                    except ValueError:
                        print(layer + "层权重广播出现问题")
    
        net_n.save(path + "age_net_prune.caffemodel")
        print("剪枝结束,保存模型名为age_net_prune.caffemodel")
    
    
    def picture(net, layer):  # 将某一layer所有卷积核的权重绝对值之和排序后画图
        weight = net.params[layer][0].data
    
        sum_l1 = []
        for i in range(weight.shape[0]):
            sum_l1.append(np.sum(abs(weight[i, :, :, :])))
    
        sum_l1.sort()
    
        x = [i for i in range(len(sum_l1))]
        plt.plot(x, sum_l1)
        plt.legend()
        plt.show()
    
    
    def get_sum_l1(net, txt_path, v):  # 定向输出各个层的卷积核的权重绝对值之和到指定文件,v为保存的前多少个值
        with open(txt_path, "w") as t:
            for layer in net.params.keys():
                weight = net.params[layer][0].data
                sum_l1 = []
                try:
                    for i in range(weight.shape[0]):
                        sum_l1.append(np.sum(abs(weight[i, :, :, :])))
                except IndexError:
                    print(layer + "该层非卷积层")
    
                sum_l1.sort()
                t.write(layer + '\n')
                for i in range(v):
                    try:
                        t.write(str(sum_l1[i]) + '  ')
                    except IndexError:
                        print(layer + "层没有" + str(v) + "个参数")
                        break
                t.write("\n\n")
    
    
    if __name__ == "__main__":
    
        root = "/home/xuqiong/code/caffeprune/"
        model = root + "age_net_new.caffemodel"
        prototxt = root + "deploy_age2_new.prototxt"
    
        py = {}  # 键是prune_layer,值是对应的prune的卷积核的序号,也就是p_k
        iy = {}  # 键是以prune_layer为input的layer,值也是对应的p_k
    
        prune_layer = ["conv1", "conv2","fc6","fc7"]
        input_layer = [["conv2"], ["conv3"],["fc7"],["fc8"]]
        th = [2,22,87,4.5]  # al元素的个数保持和prune_layer个数一致,阈值可以自己设
        #prune_layer = ["conv3"]
        #input_layer = [["fc6"]]
        #th = [20]
    
    
        caffe.set_mode_gpu()
        net = caffe.Net(prototxt, model, caffe.TEST)
        pro_n = root + "deploy_age2_prune.prototxt"
        shutil.copyfile(prototxt, pro_n)
        
        #net_n = caffe.Net(pro_n, caffe.TEST)
        #picture(net,"conv3")
        
        for (layer1, layer2, t) in zip(prune_layer, input_layer, th):
            py[layer1] = get_prune(net, layer1, threshold=t)
            print("need prune layer:",layer1)
            for m in layer2:  # 以prune_layer为输入的layer可能有多个,所以input_layer每个元素是一个列表,此处对列表中每一个元素赋值
                iy[m] = py[layer1]
    
        w_prune, b_prune = prune(net, py, iy)
        print("len of w&b:", len(w_prune),len(b_prune))
        #print("prune w & b:",w_prune,b_prune)
        get_prototxt(py, pro_n)
    
        net_n = caffe.Net(pro_n, caffe.TEST)
        savemodel(net, net_n, w_prune, b_prune, root)
    
    '''
        while (raw_input("按1将生成deploy_prune_new.prototxt:")) == "1":
            w_prune, b_prune = prune(net, py, iy)
            get_prototxt(py, pro_n)
    
            while (raw_input("按1将生成剪枝后的模型:")) == "1":
                net_n = caffe.Net(pro_n, caffe.TEST)
                savemodel(net, net_n, w_prune, b_prune, root)
                break
            break
    '''

参考文章：[https://blog.csdn.net/dlyldxwl/article/details/79502829][https_blog.csdn.net_dlyldxwl_article_details_79502829]

[https_blog.csdn.net_dlyldxwl_article_details_79502829]: https://blog.csdn.net/dlyldxwl/article/details/79502829