caffe模型通道剪枝channel pruning

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