Channel Pruning for Accelerating Very Deep Neural Networks代码详解

完整的剪枝过程主要包括以下几个步骤：
1、权重分解以及channel pruning
2、对剪枝后的模型fine_tune，恢复表征能力
3、将微调后的caffemodel作为faster的pre_trained模型，再次进行fine_tune

以下对主要的函数进行讲解，文章的后面会对部分小函数进行解释。

---

网络裁剪

通过以下指令进行pruning：
./run.sh python3 train.py -action c3 -caffe [GPU0]

if __name__ == '__main__':
    args = parse_args()
    cfgs.set_nBatches(dcfgs.nBatches)
    dcfgs.dic.option=1
    DEBUG = 1

    if args.action == cfgs.Action.c3:
        c3()#调用这个函数

c3函数是控制整个流程的函数：

def c3(pt=cfgs.vgg.model,model=cfgs.vgg.weights): 
    dcfgs.splitconvrelu=True
    cfgs.accname='accuracy@5' 
    def solve(pt, model):
        net = Net(pt, model=model)
        net.load_frozen() # step1生成的frozen文件中加载
        '''调用剪枝的R3主函数，进行空间分解、通道分解、通道剪枝，返回修建后的权重，以及剪枝后的网络结构'''
        WPQ, new_pt = net.R3()
        return {
  "WPQ": WPQ, "new_pt": new_pt}

    def stepend(new_pt, model, WPQ):
        net = Net(new_pt, model=model)
        net.WPQ = WPQ
        net.finalmodel(save=False) #将weight加载进WPQ
        net.dis_memory()
        new_pt, new_model = net.save(prefix='3c')#存为3c为前缀的新模型
        print('caffe test -model',new_pt, '-weights',new_model)
        return {
  "final": None}

    worker = Worker()
    outputs = worker.do(step0, pt=pt, model=model)#执行step0函数
    pt = outputs['pt']
    outputs = worker.do(step1,**outputs)#执行step1函数
    outputs['pt'] = mem_pt(pt)
    outputs = worker.do(solve, **outputs)#执行solve函数
    printstage("saving")
    outputs = worker.do(stepend, model=model, **outputs)#执行stepend函数

Step 0 网络预处理

def step0(pt, model):
    net = Net(pt, model=model, noTF=1)  
    # WPQ存的是剪枝后的权重，稍后会存在caffemodel中
    WPQ, pt, model = net.preprocess_resnet() #对relu bathnorm scale层进行处理
    return {
  "WPQ": WPQ, "pt": pt, "model": model}

Step 1 数据准备

def step1(pt, model, WPQ, check_exist=False):
    net = Net(pt, model, noTF=1)
    model = net.finalmodel(WPQ) #将WPQ加载进model中（后续有函数详解）
    convs = net.convs  #convs等于proto中所有conv的name list  不去掉最后一层conv

    #在整个数据集中选取batch=5000张图，每个特征图中只选取10个点来计算
    #将这50000个点对应的特征图值 以及通道数  做成freeze.pikcle的内存形式
    net.freeze_images(check_exist=check_exist, convs=convs)
    return {
  "model":model}

Step 2 solve 网络剪枝主函数

def solve(pt, model):
    net = Net(pt, model=model)
    net.load_frozen() # step1生成的frozen文件中加载
    '''调用剪枝的R3主函数，进行空间分解、通道分解、通道剪枝，返回修建后的权重，以及剪枝后的网络结构'''
    WPQ, new_pt = net.R3()
    return {
  "WPQ": WPQ, "new_pt": new_pt}

Step 3 stepend 保存最终的caffemodel

def stepend(new_pt, model, WPQ):
    net = Net(new_pt, model=model)
    net.WPQ = WPQ
    net.finalmodel(save=False) #将weight加载进WPQ
    net.dis_memory()
    new_pt, new_model = net.save(prefix='3c')#存为3c为前缀的新模型
    print('caffe test -model',new_pt, '-weights',new_model)
    return {
  "final": None}

上述函数的具体实现

R3函数

被step3的solve函数调用，所有空间分解、通道分解、通道剪枝的函数都放在R3函数中：

def R3(self): 
       speed_ratio = dcfgs.dic.keep  #3  也就是3倍加速
       prefix += str(int(speed_ratio)+1)+'x'
       DEBUG = True
       convs= self.convs
       self.WPQ = dict()
       self.selection = dict()
       self._mem = True
       end = 5 #即只考虑[1,5)组卷积组
       alldic = ['conv%d_1' % i for i in range(1,end)] + ['conv%d_2' % i for i in range(3, end)]
       pooldic = ['conv1_2', 'conv2_2']#, 'conv3_3']  #后面有pool的conv层
       #经验的要保留的通道数目        #原始    faster    3c VGG
       rankdic = {
  'conv1_1': 17,  #64      #12      #64
                  'conv1_2': 17,  #64      #64      #V22 H22 P58
                  'conv2_1': 37,  #128     #21      #V49 H49 P117
                  'conv2_2': 47,  #128     #128     #V62 H62 P117
                  'conv3_1': 83,  #256     #73      #V110 H110 P237
                  'conv3_2': 89,  #256     #58      #V118 H118 P242
                  'conv3_3': 106, #256     #256     #V114 H114 P256
                  'conv4_1': 175, #512     #121     #V233 H233 P475
                  'conv4_2': 192, #512     #166     #V256 H256 P457
                  'conv4_3': 227, #512     #512     #V302 H302 P512
                  'conv5_1': 398, #512     #512     #V398 H398 P512
                  'conv5_2': 390, #512     #512     #V390 H390 P512
                  'conv5_3': 379} #512     #512     #V379 H379 P512

       c_ratio = 1.15

       def getX(name):
           x = self.extract_XY(self.bottom_names[name][0], name)
           return np.rollaxis(x.reshape((-1, 3, 3, x.shape[1])), 3, 1).copy()

       def setConv(c, d):
           if c in self.selection:
               self.param_data(c)[:,self.selection[c],:,:] = d
           else:
               self.set_param_data(c, d)

       t = Timer()
       #zip  将两个参数按元祖组合成list 返回    所以得到的是类似[(conv1,conv2),(conv2,conv3),......(conv5,pool5)]
       #conv, convnext分别代表这一层conv，和其下一层conv或pool
       for conv, convnext in zip(convs[1:], convs[2:]+['pool5']): 
           #因为空间分解、通道分解会产生新的conv_V\H\P层，代替原来的conv层
           conv_V = underline(conv, 'V')   #conv后加下划线V                       
           conv_H = underline(conv, 'H')                         
           conv_P = underline(conv, 'P')
           W_shape = self.param_shape(conv)
           d_c = int(W_shape[0] / c_ratio)
           rank = rankdic[conv]
           d_prime = rank
           if d_c < rank: d_c = rank

           '''spatial decomposition空间分解'''
           if True:
               t.tic()
               weights = self.param_data(conv)
               if conv in self.selection:
                   weights = weights[:,self.selection[conv],:,:]
               if 1:
                   Y = self._feats_dict[conv] - self.param_b_data(conv)
                   X = getX(conv)
                   if conv in self.selection:
                       X = X[:,self.selection[conv],:,:]

                   #执行的是SVD分解，左奇异值是V，右奇异值*对角是H
                   #V变成了3*1的低维卷积层，H变成了相同维数的1*3的卷积层