caffe模型通道剪枝channel pruning

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

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

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

  1. # coding:utf-8
  2. # by chen yh
  4. import caffe
  5. import numpy as np
  6. import shutil
  7. import matplotlib.pyplot as plt
  9. '''
  10. These parameters need modification: 
  11.     root: root directory ;
  12.     model: your caffemodel ; 
  13.     prototxt: your prototxt ; 
  15.     prune layer: need prune layer name, a list ;
  16.     input layer: input of these layers is output of prune layer,each element is a list ;
  17.     th : thereshold of each prune layer,a list.
  19. Please ensure lenth of prune layer, input layer and th.
  21. picture and get get_sum_l1 functions can help you find suitable threshold. 
  22. '''
  25. def get_prune(net, layer, threshold):  # 返回layer中低于阈值threshold的卷积核的序号
  26.     weight_ori = net.params[layer][0].data
  27.     # bias_ori=net.params[layer][1].data
  29.     print("layer:",layer,"shape:",weight_ori.shape)
  30.     #print("weight_ori:",weight_ori)
  31.     sum_l1 = []
  32.     for i in range(weight_ori.shape[0]):
  33.         if (layer == "fc6") or (layer == "fc7") or (layer == "fc8"):
  34.             sum_l1.append((i, np.sum(abs(weight_ori[i, :]))))
  35.         else:
  36.             sum_l1.append((i, np.sum(abs(weight_ori[i, :, :, :]))))  # sum_l1存放每个卷积核的所有权重绝对值之和
  37.     print("sum_l1.shape:",len(sum_l1))
  38.     de_keral = []  # de_keral存放大于阈值的卷积核的序号
  39.     for i in sum_l1:
  40.         if i[1] > threshold:
  41.             de_keral.append(i[0])
  43.     print(layer + "层需要prune的卷积核有" + str(weight_ori.shape[0] - len(de_keral)) + "个,保留的卷积核有" + str(len(de_keral)) + "个")
  44.     print("de_keral.shape:",len(de_keral))
  45.     return de_keral
  48. def prune(net, pk, lk):  # 输出两个字典,键都是修剪层的layer的名字,值分别是修剪层的weight和bias
  49.     w_new = {}  # 键是layer,值是保存后的weight
  50.     b_new = {}  # 键是layer,值是保存后的bias
  51.     for l in pk.keys():  # 待剪层权重处理 w_n = w[pk[l],:,;,;]
  52.         #print("pk.keys:",l)
  53.         w_old = net.params[l][0].data
  54.         b_old = net.params[l][1].data
  55.         if (l == "fc6") or (l == "fc7"):
  56.             w_n = w_old[pk[l], :]
  57.         else:
  58.             w_n = w_old[pk[l], :, :, :]
  59.         b_n = b_old[pk[l]]
  60.         w_new[l] = w_n
  61.         b_new[l] = b_n
  62.         # net_n.params[l][0].data[...] = w_n
  63.         # net_n.params[l][1].data[...] = b_n
  64.         print("player:",l,",w.shape:",w_new[l].shape,",b.shape",b_new[l].shape)
  65.     for l in lk.keys():  # 以待剪层为输入的层权重处理
  66.         if l not in pk.keys():  # bottom被修剪后本身没有被修剪,所以其权重只需要在原来的net上面取切片,w_n = w[:,lk[l],:,:]
  67.             #print("lk.keys:", l)
  68.             if (l != "conv4_3_norm"):  # 对传统卷积层的处理
  69.                 w_o = net.params[l][0].data
  70.                 b_o = net.params[l][1].data
  71.                 b_new[l] = b_o  # bias保留,因为这些层没有剪卷积核
  72.                 if (l == "fc6") or (l == "fc7") or (l == "fc8"):
  73.                     w_n = w_o[:, lk[l]]
  74.                 else:
  75.                     w_n = w_o[:, lk[l], :, :]
  76.                 w_new[l] = w_n
  77.             else:  # 对特殊层的处理,参数个数不是2
  78.                 w_o = net.params[l][0].data
  79.                 w_n = w_o[lk[l],]
  80.                 w_new[l] = w_n
  81.             print("klayer:",l,",w.shape:",w_new[l].shape,",b.shape",b_new[l].shape)
  83.         else:  # pk 和 lk共有的层,也就是这层的bottom和层本身都被修剪过,所以权重不能在原来的net上切片,利用保存了的w_new取切片.
  84.             w_o = w_new[l]
  85.             print("lk.keys else:",l)
  86.             if (l == "fc6") or (l == "fc7"):
  87.                 w_n = w_o[:, lk[l]]
  88.             else:
  89.                 w_n = w_o[:, lk[l], :, :]
  90.             w_new[l] = w_n
  91.             print("llayer:",l,",w.shape:",w_new[l].shape)
  93.     return w_new, b_new
  96. def get_prototxt(pk, pro_n):  # 复制原来的prototxt,并修改修剪层的num_output,这一段代码有点绕,有空的话优化为几个单独的函数或者弄个类
  97.     with open(pro_n, "r") as p:
  98.         lines = p.readlines()
  99.     k = 0
  100.     with open(pro_n, "w") as p:
  101.         while k < len(lines):  # 遍历所有的lines,此处不宜用for.
  102.             #print("lines[k]:",lines[k])
  103.             if 'name:' in lines[k]:
  104.                 print("lines[k].split:",lines[k].split('"')[1])
  105.                 l_name = lines[k].split('"')[1]  # 获取layer name
  106.                 if l_name in pk.keys():  # 如果name在待修剪层中,则需要修改,下面进入一个找channel的循环块.
  107.                     while True:
  108.                         if "num_output:" in lines[k]:
  109.                             channel_n = "    num_output: " + str(len(pk[l_name])) + "\n"
  110.                             p.write(channel_n)
  111.                             k = k + 1
  112.                             break
  113.                         else:
  114.                             p.write(lines[k])
  115.                             k = k + 1
  116.                 else:  # name不在待修剪层中,直接copy行
  117.                     p.write(lines[k])
  118.                     k = k + 1
  120.             else:
  121.                 p.write(lines[k])
  122.                 k = k + 1
  123.     print("deploy_rebirth_prune.prototxt已写好")
  126. def savemodel(net, net_n, w_prune, b_prune, path):  # 储存修改后的caffemodel
  127.     for layer in net.params.keys():
  128.         if layer in w_prune.keys():
  129.             print("save model-layer:",layer,"len:",w_prune[layer].shape,"len2:",net_n.params[layer][0].data[...].shape)
  130.             net_n.params[layer][0].data[...] = w_prune[layer]
  131.             if layer in b_prune.keys():
  132.                 net_n.params[layer][1].data[...] = b_prune[layer]
  133.         else:
  134.             weight = net.params[layer]
  135.             for index, w in enumerate(weight):
  136.                 try:
  137.                     net_n.params[layer][index].data[...] = w.data
  138.                 except ValueError:
  139.                     print(layer + "层权重广播出现问题")
  141.     net_n.save(path + "age_net_prune.caffemodel")
  142.     print("剪枝结束,保存模型名为age_net_prune.caffemodel")
  145. def picture(net, layer):  # 将某一layer所有卷积核的权重绝对值之和排序后画图
  146.     weight = net.params[layer][0].data
  148.     sum_l1 = []
  149.     for i in range(weight.shape[0]):
  150.         sum_l1.append(np.sum(abs(weight[i, :, :, :])))
  152.     sum_l1.sort()
  154.     x = [i for i in range(len(sum_l1))]
  155.     plt.plot(x, sum_l1)
  156.     plt.legend()
  157.     plt.show()
  160. def get_sum_l1(net, txt_path, v):  # 定向输出各个层的卷积核的权重绝对值之和到指定文件,v为保存的前多少个值
  161.     with open(txt_path, "w") as t:
  162.         for layer in net.params.keys():
  163.             weight = net.params[layer][0].data
  164.             sum_l1 = []
  165.             try:
  166.                 for i in range(weight.shape[0]):
  167.                     sum_l1.append(np.sum(abs(weight[i, :, :, :])))
  168.             except IndexError:
  169.                 print(layer + "该层非卷积层")
  171.             sum_l1.sort()
  172.             t.write(layer + '\n')
  173.             for i in range(v):
  174.                 try:
  175.                     t.write(str(sum_l1[i]) + '  ')
  176.                 except IndexError:
  177.                     print(layer + "层没有" + str(v) + "个参数")
  178.                     break
  179.             t.write("\n\n")
  182. if __name__ == "__main__":
  184.     root = "/home/xuqiong/code/caffeprune/"
  185.     model = root + "age_net_new.caffemodel"
  186.     prototxt = root + "deploy_age2_new.prototxt"
  188.     py = {}  # 键是prune_layer,值是对应的prune的卷积核的序号,也就是p_k
  189.     iy = {}  # 键是以prune_layer为input的layer,值也是对应的p_k
  191.     prune_layer = ["conv1", "conv2","fc6","fc7"]
  192.     input_layer = [["conv2"], ["conv3"],["fc7"],["fc8"]]
  193.     th = [2,22,87,4.5]  # al元素的个数保持和prune_layer个数一致,阈值可以自己设
  194.     #prune_layer = ["conv3"]
  195.     #input_layer = [["fc6"]]
  196.     #th = [20]
  199.     caffe.set_mode_gpu()
  200.     net = caffe.Net(prototxt, model, caffe.TEST)
  201.     pro_n = root + "deploy_age2_prune.prototxt"
  202.     shutil.copyfile(prototxt, pro_n)
  204.     #net_n = caffe.Net(pro_n, caffe.TEST)
  205.     #picture(net,"conv3")
  207.     for (layer1, layer2, t) in zip(prune_layer, input_layer, th):
  208.         py[layer1] = get_prune(net, layer1, threshold=t)
  209.         print("need prune layer:",layer1)
  210.         for m in layer2:  # 以prune_layer为输入的layer可能有多个,所以input_layer每个元素是一个列表,此处对列表中每一个元素赋值
  211.             iy[m] = py[layer1]
  213.     w_prune, b_prune = prune(net, py, iy)
  214.     print("len of w&b:", len(w_prune),len(b_prune))
  215.     #print("prune w & b:",w_prune,b_prune)
  216.     get_prototxt(py, pro_n)
  218.     net_n = caffe.Net(pro_n, caffe.TEST)
  219.     savemodel(net, net_n, w_prune, b_prune, root)
  221. '''
  222.     while (raw_input("按1将生成deploy_prune_new.prototxt:")) == "1":
  223.         w_prune, b_prune = prune(net, py, iy)
  224.         get_prototxt(py, pro_n)
  226.         while (raw_input("按1将生成剪枝后的模型:")) == "1":
  227.             net_n = caffe.Net(pro_n, caffe.TEST)
  228.             savemodel(net, net_n, w_prune, b_prune, root)
  229.             break
  230.         break
  231. '''

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

发表评论

表情：
评论列表 （有 0 条评论，345人围观）

还没有评论，来说两句吧...

相关阅读

    相关 GO语言：channel通道

    通道可以被认为是Goroutines通信的管道。类似于管道中的水从一端到另一端的流动，数据可以从一端发送到另一端，通过通道接收。 在前面讲Go语言的并发时候，我们就说过，当多

    系统管理员系统管理员/ 0 赞/ 125 阅读