import sys
import caffe
import numpy as np
percent = 0.5
np.set_printoptions(threshold=sys.maxsize)
MODEL_FILE = '********.prototxt' #模型文件
PRETRAIN_FILE = '*********.caffemodel' #参数文件
params_txt = 'params.txt' #保存参数文件
pf = open(params_txt, 'w')
total = 0 #统计BN参数的个数
net = caffe.Net(MODEL_FILE, PRETRAIN_FILE, caffe.TEST) #加载模型和参数
print("load model successful")
for param_name in net.params.keys():
weight = net.params[param_name][0].data
bias = net.params[param_name][1].data
if 'bn' in param_name:
total += weight.shape[0]
print("the number of bn is :", total)
bn = []
index = 0
small_num = 0
for param_name in net.params.keys():
weight = net.params[param_name][0].data #gama
bias = net.params[param_name][1].data #beta
if 'bn' in param_name:
size = weight.shape[0]
weight.shape = (-1, 1)
for w in weight:
bn.append(abs(w))
y = sorted(bn) #根据BN 的gama值进行排序
thre_index = int(total * percent) #根据设置的裁剪比例来确定gama阈值
print("thread is :", thre) │
# 这里并没有裁剪,只是对要裁剪的gama置0,测试了一下结果
# 可能会出现某些layer的channel被全部裁剪的情况,对于这种情况,保留其gama绝对值最大的那个channel
pruned = 0 │
cfg = [] │
cfg_mask = [] │
for param_name in net.params.keys(): │
weight = net.params[param_name][0].data │
bias = net.params[param_name][1].data │
if 'bn' in param_name: │
size = weight.shape[0] │
print(param_name) │
num = 0 │
index = 0 │
max_id = 0 │
max_gama = 0 │
max_b = 0 │
│
for w in weight: │
if abs(w) > max_gama: │
max_id = index │
max_gama = w │
max_b = net.params[param_name][1].data[index] │
│
if abs(w) < thre: │
num +=1 │
small_num +=1 │
net.params[param_name][0].data[index] = 0 │
net.params[param_name][1].data[index] = 0 │
│
index +=1 │
if num == size: │
net.params[param_name][0].data[max_id] = max_gama │
net.params[param_name][1].data[max_id] = max_b │
│
print(size,num) │
print(small_num)#
net.save('at_least_one_new.caffemodel')