# -*- coding:utf-8 -*- # 通过Kmeans聚类的方法来量化权重 import numpy as np import matplotlib.pyplot as plt import scipy.cluster.vq as scv import pickle import os os.environ['GLOG_minloglevel'] = '2' import caffe import time # 获得各层的量化码表 def kmeans_net(net, layers, num_c=16, initials=None): # net: 网络 # layers: 需要量化的层 # num_c: 各层的量化级别 # initials: 初始聚类中心 codebook = {} # 量化码表 if type(num_c) == type(1): num_c = [num_c] * len(layers) else: assert len(num_c) == len(layers) # 对各层进行聚类分析 print "==============Perform K-means=============" for idx, layer in enumerate(layers): print "Eval layer:", layer W = net.params[layer][0].data.flatten() W = W[np.where(W != 0)] # 筛选不为0的权重 # 默认情况下，聚类中心为线性分布中心 if initials is None: # Default: uniform sample min_W = np.min(W) max_W = np.max(W) initial_uni = np.linspace(min_W, max_W, num_c[idx] - 1) codebook[layer], _ = scv.kmeans(W, initial_uni) elif type(initials) == type(np.array([])): codebook[layer], _ = scv.kmeans(W, initials) elif initials == 'random': codebook[layer], _ = scv.kmeans(W, num_c[idx] - 1) else: raise Exception # 将0权重值附上 codebook[layer] = np.append(0.0, codebook[layer]) print "codebook size:", len(codebook[layer]) return codebook # 随机量化权重值 def stochasitc_quantize2(W, codebook): # mask插入新维度：(W.shape,1) mask = W[:, np.newaxis] - codebook mask_neg = mask mask_neg[mask_neg > 0.0] -= 99999.0 max_neg = np.max(mask_neg, axis=1) max_code = np.argmax(mask_neg, axis=1) mask_pos = mask mask_pos += 99999.0 min_code = np.argmin(mask_pos, axis=1) min_pos = np.min(mask_pos, axis=1) rd = np.random.uniform(low=0.0, high=1.0, size=(len(W))) thresh = min_pos.astype(np.float32) / (min_pos - max_neg) max_idx = thresh < rd min_idx = thresh >= rd codes = np.zeros(W.shape) codes[max_idx] += min_code[max_idx] codes[min_idx] += max_code[min_idx] return codes.astype(np.int) # 得到网络的量化权重值 def quantize_net(net, codebook): layers = codebook.keys() codes_W = {} print "================Perform quantization==============" for layer in layers: print "Quantize layer:", layer W = net.params[layer][0].data codes, _ = scv.vq(W.flatten(), codebook[layer]) # 根据码表得到量化权重值 # codes = stochasitc_quantize2(W.flatten(), codebook[layer]) # 采用随机量化的方式 codes = np.reshape(codes, W.shape) codes_W[layer] = np.array(codes, dtype=np.uint32) # 将量化后的权重保存到网络中 W_q = np.reshape(codebook[layer][codes], W.shape) np.copyto(net.params[layer][0].data, W_q) return codes_W def quantize_net_with_dict(net, layers, codebook, use_stochastic=False, timing=False): start_time = time.time() codeDict = {} # 记录各个量化中心所处的位置 maskCode = {} # 各层量化结果 for layer in layers: print "Quantize layer:", layer W = net.params[layer][0].data if use_stochastic: codes = stochasitc_quantize2(W.flatten(), codebook[layer]) else: codes, _ = scv.vq(W.flatten(), codebook[layer]) W_q = np.reshape(codebook[layer][codes], W.shape) net.params[layer][0].data[...] = W_q maskCode[layer] = np.reshape(codes, W.shape) codeBookSize = len(codebook[layer]) a = maskCode[layer].flatten() b = xrange(len(a)) codeDict[layer] = {} for i in xrange(len(a)): codeDict[layer].setdefault(a[i], []).append(b[i]) if timing: print "Update codebook time:%f" % (time.time() - start_time) return codeDict, maskCode def static_vars(**kwargs): def decorate(func): for k in kwargs: setattr(func, k, kwargs[k]) return func return decorate @static_vars(step_cache={}, step_cache2={}, count=0) def update_codebook_net(net, codebook, codeDict, maskCode, args, update_layers=None, snapshot=None): start_time = time.time() extra_lr = args['lr'] # 基础学习速率 decay_rate = args['decay_rate'] # 衰减速率 momentum = args['momentum'] # 遗忘因子 update_method = args['update'] # 更新方法 smooth_eps = 0 normalize_flag = args['normalize_flag'] # 是否进行归一化 if update_method == 'rmsprop': extra_lr /= 100 # 对码表与量化结果的初始化 if update_codebook_net.count == 0: step_cache2 = update_codebook_net.step_cache2 step_cache = update_codebook_net.step_cache if update_method == 'adadelta': for layer in update_layers: step_cache2[layer] = {} for code in xrange(1, len(codebook[layer])): step_cache2[layer][code] = 0.0 smooth_eps = 1e-8 for layer in update_layers: step_cache[layer] = {} for code in xrange(1, len(codebook[layer])): step_cache[layer][code] = 0.0 update_codebook_net.count = 1 else: # 读入上次运算的结果 step_cache2 = update_codebook_net.step_cache2 step_cache = update_codebook_net.step_cache update_codebook_net.count += 1 # 所有层名 total_layers = net.params.keys() if update_layers is None: # 所有层都需要进行更新 update_layers = total_layers # 权重码表的更新 for layer in total_layers: if layer in update_layers: diff = net.params[layer][0].diff.flatten() # 误差梯度 codeBookSize = len(codebook[layer]) dx = np.zeros((codeBookSize)) # 编码表的误差更新 for code in xrange(1, codeBookSize): indexes = codeDict[layer][code] # codeDict保存属于某编码的权重的序号 #diff_ave = np.sum(diff[indexes]) / len(indexes) diff_ave = np.sum(diff[indexes]) # 统计该编码所有的误差更新和 # 针对于不同方法进行更新 if update_method == 'sgd': dx[code] = -extra_lr * diff_ave elif update_method == 'momentum': if code in step_cache[layer]: dx[code] = momentum * step_cache[layer][code] - (1 - momentum) * extra_lr * diff_ave step_cache[layer][code] = dx elif update_method == 'rmsprop': if code in step_cache[layer]: step_cache[layer][code] = decay_rate * step_cache[layer][code] + (1.0 - decay_rate) * diff_ave ** 2 dx[code] = -(extra_lr * diff_ave) / np.sqrt(step_cache[layer][code] + 1e-6) elif update_method == 'adadelta': if code in step_cache[layer]: step_cache[layer][code] = step_cache[layer][code] * decay_rate + (1.0 - decay_rate) * diff_ave ** 2 dx[code] = -np.sqrt((step_cache2[layer][code] + smooth_eps) / (step_cache[layer][code] + smooth_eps)) * diff_ave step_cache2[layer][code] = step_cache2[layer][code] * decay_rate + (1.0 - decay_rate) * (dx[code] ** 2) # 是否需要进行归一化更新参数 if normalize_flag: codebook[layer] += extra_lr * np.sqrt(np.mean(codebook[layer] ** 2)) / np.sqrt(np.mean(dx ** 2)) * dx else: codebook[layer] += dx else: pass # maskCode保存编码结果 W2 = codebook[layer][maskCode[layer]] net.params[layer][0].data[...] = W2 # 量化后权重值 print "Update codebook time:%f" % (time.time() - start_time) # 保存量化结果 def store_all(net, codebook, dir_t, idx=0): net.save(dir_t + 'caffemodel%d' % idx) # 量化网络及码表 pickle.dump(codebook, open(dir_t + 'codebook%d' % idx, 'w')) # 恢复权重值 def recover_all(net, dir_t, idx=0): layers = net.params.keys() net.copy_from(dir_t + 'caffemodel%d' % idx) codebook = pickle.load(open(dir_t + 'codebook%d' % idx)) maskCode = {} codeDict = {} for layer in layers: W = net.params[layer][0].data # 码表结果 codes, _ = scv.vq(W.flatten(), codebook[layer]) # 编码结果重新排列 maskCode[layer] = np.reshape(codes, W.shape) codeBookSize = len(codebook[layer]) a = maskCode[layer].flatten() b = xrange(len(a)) codeDict[layer] = {} for i in xrange(len(a)): # codeDict保存每个码有哪些位置，而maskCode保存每个位置属于哪个码 codeDict[layer].setdefault(a[i], []).append(b[i]) return codebook, maskCode, codeDict def analyze_log(fileName): data = open(fileName, "r") y = [] for line in data: y.append(float(line.split()[0])) return y # 读入测试数据 def parse_caffe_log(log): lines = open(log).readlines() try: res = map(lambda x: float(x.split()[-1]), lines[-3:-1]) except Exception as e: print e res = [0.0, 0.0] return res # 检测量化后网络的精度 def test_quantize_accu(test_net): test_iter = 100 test_loss = 0 accuracy = 0 for test_it in range(test_iter): # 进行一次测试 test_net.forward() # 计算test loss test_loss += test_net.blobs['loss'].data # 计算test accuracy accuracy += test_net.blobs['accuracy'].data return (test_loss / test_iter), (accuracy / test_iter) def save_quantize_net(codebook, maskcode, net_filename, total_layers): # 编码 quantizeNet = {} for layer in total_layers: quantizeNet[layer+'_codebook'] = np.float32(codebook[layer]) quantizeNet[layer + '_maskcode'] = np.int8(maskcode[layer]) np.savez(net_filename,quantizeNet) # 保存修剪量化的网络参数 def save_pruned_quantize_net(codebook, maskcode, net_filename, total_layers): # W_flatten: 扁平化的权重矩阵 # num_level: 量化级别 quantizeNet = {} for layer in total_layers: W_flatten = maskCode[layer].flatten() indx = 0 num_level = 8 csc_W = [] csc_indx = [] for n in range(len(W_flatten)): if W_flatten[n]!=0 or indx == 2**num_level: csc_W.append(W_flatten[n]) csc_indx.append(indx) indx = 0 else: indx += 1 if indx!=0: csc_W.append(0) csc_indx.append(indx-1) print max(csc_indx) quantizeNet[layer + '_codebook'] = np.float32(codebook[layer]) quantizeNet[layer + '_maskcode_W'] = np.array(csc_W, dtype=np.int8) print max(csc_indx) quantizeNet[layer + '_maskcode_indx'] = np.array(csc_indx, dtype=np.int8) np.savez(net_filename, quantizeNet) caffe.set_mode_gpu() caffe.set_device(0) caffe_root = '../../' model_dir = caffe_root + 'models/mnist/' deploy = model_dir + 'deploy.prototxt' solver_file = model_dir + 'solver.prototxt' # model_name = 'LeNet5_Mnist_shapshot_iter_10000' model_name = 'LeNet5_Mnist_shapshot_iter_10000_pruned' caffemodel = model_dir + model_name + '.caffemodel' dir_t = '/weight_quantize/' # 运行测试命令 args = dict(lr=0.01, decay_rate = 0.0009, momentum = 0.9, update = 'adadelta', normalize_flag = False) start_time = time.time() solver = caffe.SGDSolver(solver_file) solver.net.copy_from(caffemodel) # 需要量化的权重 total_layers = ['conv1','conv2','ip1','ip2'] num_c = 2 ** 8 # 量化级别，由8位整数表示 codebook = kmeans_net(solver.test_nets[0], total_layers, num_c) codeDict, maskCode = quantize_net_with_dict(solver.test_nets[0], total_layers, codebook) quantize_net_caffemodel = model_dir + model_name + '_quantize.caffemodel' solver.test_nets[0].save(quantize_net_caffemodel) quantize_net_npz = model_dir + model_name + '_quantize_net' save_pruned_quantize_net(codebook, maskCode, quantize_net_npz , total_layers) # 迭代训练编码表 accuracys = [] co_iters = 40 ac_iters = 10 for i in xrange(2500): if (i % (co_iters + ac_iters) == 0 and i > 0): # 重新量化 # 导入训练后的 codebook = kmeans_net(solver.net, total_layers, num_c) codeDict, maskCode = quantize_net_with_dict(solver.net, total_layers, codebook) solver.net.save(quantize_net_caffemodel) solver.test_nets[0].copy_from(quantize_net_caffemodel) _, accu = test_quantize_accu(solver.test_nets[0]) accuracys.append(accu) solver.step(1) if (i % (co_iters + ac_iters) < co_iters): # 码表更新 update_codebook_net(solver.net, codebook, codeDict, maskCode, args=args, update_layers=total_layers) print "Iter:%d, Time cost:%f" % (i, time.time() - start_time) plt.plot(accuracys, 'r.-') plt.show()