YOLOV3剪枝源码阅读---模型部署加速

2020-01-07  本文已影响0人  _从前从前_

YOLOV3剪枝

论文：Network Slimming-Learning Efficient Convolutional Networks through Network Slimming

剪枝项目参考https://github.com/tanluren/yolov3-channel-and-layer-pruning

主要思路

剪枝过程

img 这里写图片描述

分为三部分，第一步，训练；第二步，剪枝；第三步，微调剪枝后的模型，循环执行

YOLOV3剪枝源码

1、正常剪枝

这部分分析来自该仓库https://github.com/coldlarry/YOLOv3-complete-pruning，但是更新的仓库也可以完成正常剪枝，prune.py。

使用了正常剪枝模式，不对short cut层（需要考虑add操作的维度一致问题）及上采样层（无BN）进行裁剪。

2、优化的正常剪枝

slim_prune.py，在正常剪枝模式的基础上，完成对shortcut层的剪枝，同时避免裁剪掉整个层。

# 解析模型文件
def parse_module_defs2(module_defs):
    CBL_idx = []
    Conv_idx = []
    shortcut_idx=dict()
    shortcut_all=set()
    ignore_idx = set()
    for i, module_def in enumerate(module_defs):
        if module_def['type'] == 'convolutional':
            # 如果是卷积层中的BN层则将该层索引添加到CBL_idx
            if module_def['batch_normalize'] == '1': 
                CBL_idx.append(i)
            else:
                Conv_idx.append(i)
            if module_defs[i+1]['type'] == 'maxpool':
                #spp前一个CBL不剪
                ignore_idx.add(i)

        elif module_def['type'] == 'upsample':
            #上采样层前的卷积层不裁剪
            ignore_idx.add(i - 1)

        elif module_def['type'] == 'shortcut':
            # 根据cfg中的from层获得shortcut的卷积层对应的索引
            identity_idx = (i + int(module_def['from']))
            # 如果shortcut连接的是卷积层则直接添加索引
            if module_defs[identity_idx]['type'] == 'convolutional':

                #ignore_idx.add(identity_idx)
                shortcut_idx[i-1]=identity_idx
                shortcut_all.add(identity_idx)
            # 如果shortcut连接的是shortcut层，则添加前一层卷积层的索引
            elif module_defs[identity_idx]['type'] == 'shortcut':

                #ignore_idx.add(identity_idx - 1)
                shortcut_idx[i-1]=identity_idx-1
                shortcut_all.add(identity_idx-1)
            shortcut_all.add(i-1)


        # 得到需要剪枝的BN层的索引
    prune_idx = [idx for idx in CBL_idx if idx not in ignore_idx]

    return CBL_idx, Conv_idx, prune_idx,shortcut_idx,shortcut_all

3、层剪枝

和优化的正常剪枝类似。这个策略是在之前的通道剪枝策略基础上衍生出来的，针对每一个shortcut层前一个CBL进行评价，对各层的Gmma均值进行排序，取最小的进行层剪枝。为保证yolov3结构完整，这里每剪一个shortcut结构，会同时剪掉一个shortcut层和它前面的两个卷积层。是的，这里只考虑剪主干中的shortcut模块。但是yolov3中有23处shortcut，剪掉8个shortcut就是剪掉了24个层，剪掉16个shortcut就是剪掉了48个层，总共有69个层的剪层空间；实验中对简单的数据集剪掉了较多shortcut而精度降低很少。

稀疏策略

scale参数默认0.001，根据数据集，mAP,BN分布调整，数据分布广类别多的，或者稀疏时掉点厉害的适当调小s;-sr用于开启稀疏训练；--prune 0适用于prune.py，--prune 1 适用于其他剪枝策略。稀疏训练就是精度和稀疏度的博弈过程，如何寻找好的策略让稀疏后的模型保持高精度同时实现高稀疏度是值得研究的问题，大的s一般稀疏较快但精度掉的快，小的s一般稀疏较慢但精度掉的慢；配合大学习率会稀疏加快，后期小学习率有助于精度回升。
注意：训练保存的pt权重包含epoch信息，可通过python -c "from models import *; convert('cfg/yolov3.cfg', 'weights/last.pt')"转换为darknet weights去除掉epoch信息，使用darknet weights从epoch 0开始稀疏训练。

1、恒定s

这是一开始的策略，也是默认的策略。在整个稀疏过程中，始终以恒定的s给模型添加额外的梯度，因为力度比较均匀，往往压缩度较高。但稀疏过程是个博弈过程，我们不仅想要较高的压缩度，也想要在学习率下降后恢复足够的精度，不同的s最后稀疏结果也不同，想要找到合适的s往往需要较高的时间成本。

bn_module.weight.grad.data.add_(s * torch.sign(bn_module.weight.data))

2、全局s衰减

关键代码是下面这句，在epochs的0.5阶段s衰减100倍。前提是0.5之前权重已经完成大幅压缩，这时对s衰减有助于精度快速回升，但是相应的bn会出现一定膨胀，降低压缩度，有利有弊，可以说是牺牲较大的压缩度换取较高的精度，同时减少寻找s的时间成本。当然这个0.5和100可以自己调整。注意也不能为了在前半部分加快压缩bn而大大提高s，过大的s会导致模型精度下降厉害，且s衰减后也无法恢复。如果想使用这个策略，可以在prune_utils.py中的BNOptimizer把下面这句取消注释。

# s = s if epoch <= opt.epochs * 0.5 else s * 0.01

3、局部s衰减

关键代码是下面两句，在epochs的0.5阶段开始对85%的通道保持原力度压缩，15%的通道进行s衰减100倍。这个85%是个先验知识，是由策略一稀疏后尝试剪通道几乎不掉点的最大比例，几乎不掉点指的是相对稀疏后精度；如果微调后还是不及baseline，或者说达不到精度要求，就可以使用策略三进行局部s衰减，从中间开始重新稀疏，这可以在牺牲较小压缩度情况下提高较大精度。如果想使用这个策略可以在train.py中把下面这两句取消注释，并根据自己策略一情况把0.85改为自己的比例，还有0.5和100也是可调的。策略二和三不建议一起用，除非你想做组合策略。

#if opt.sr and opt.prune==1 and epoch > opt.epochs * 0.5:
    # idx2mask = get_mask2(model, prune_idx, 0.85)

知识蒸馏

参考论文：Learning Efficient Object Detection Models with Knowledge Distillation。

核心思想：

image.png 
def distillation_loss2(model, targets, output_s, output_t):
    '''
    :param model:  学生模型
    :param targets: 标签
    :param output_s: 学生模型的输出
    :param output_t: 老师模型的输出
    :return: 附加Loss
    '''
    reg_m = 0.0
    T = 3.0
    Lambda_cls, Lambda_box = 0.0001, 0.001
        
    # KL 损失，衡量两个分布的差异
    criterion_st = torch.nn.KLDivLoss(reduction='sum')
    ft = torch.cuda.FloatTensor if output_s[0].is_cuda else torch.Tensor
    lcls, lbox = ft([0]), ft([0])
        # 标签转换
    tcls, tbox, indices, anchor_vec = build_targets(model, targets)
    reg_ratio, reg_num, reg_nb = 0, 0, 0
    for i, (ps, pt) in enumerate(zip(output_s, output_t)):  # layer index, layer predictions
        b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx

        nb = len(b)
        if nb:  # number of targets
            pss = ps[b, a, gj, gi]  # prediction subset corresponding to targets
            pts = pt[b, a, gj, gi]

            psxy = torch.sigmoid(pss[:, 0:2])  # pxy = pxy * s - (s - 1) / 2,  s = 1.5  (scale_xy)
            psbox = torch.cat((psxy, torch.exp(pss[:, 2:4]) * anchor_vec[i]), 1).view(-1, 4)  # predicted box

            ptxy = torch.sigmoid(pts[:, 0:2])  # pxy = pxy * s - (s - 1) / 2,  s = 1.5  (scale_xy)
            ptbox = torch.cat((ptxy, torch.exp(pts[:, 2:4]) * anchor_vec[i]), 1).view(-1, 4)  # predicted box


            l2_dis_s = (psbox - tbox[i]).pow(2).sum(1)
            l2_dis_s_m = l2_dis_s + reg_m
            l2_dis_t = (ptbox - tbox[i]).pow(2).sum(1)
            l2_num = l2_dis_s_m > l2_dis_t
            lbox += l2_dis_s[l2_num].sum()
            reg_num += l2_num.sum().item()
            reg_nb += nb

        output_s_i = ps[..., 4:].view(-1, model.nc + 1)
        output_t_i = pt[..., 4:].view(-1, model.nc + 1)
        lcls += criterion_st(nn.functional.log_softmax(output_s_i/T, dim=1), nn.functional.softmax(output_t_i/T,dim=1))* (T*T) / ps.size(0)

    if reg_nb:
        reg_ratio = reg_num / reg_nb

    return lcls * Lambda_cls + lbox * Lambda_box, reg_ratio

自己项目实验的总结：

一、测试环境

二、测试模型（YOLOV3）

模型采用pytorch框架训练、剪枝，转换为darknet模型，再转换成ONNX模型，最后转换为tensorrt模型。将以下六个模型依次命名为model1、model2、...、model6。

  1. Pytorch模型：model.pt
  2. Pytorch剪枝且微调模型：model_075.weights（转换为darknet格式）
  3. Tensorrt 模型：model_416416.trt（由model1转换）
  4. Tensorrt 模型：model_256416.trt（由model1转换）
  5. Tensorrt 模型：model_075_256416.trt（由model2转换）
  6. Tensorrt 模型：model_075_416416.trt（由model2转换）

三、测试结果

通过采用不同的输入形状（矩形和正方形）、不同的输入尺寸（320、416、608）、不同Batch Size（1、4、8、16）下对同一张图片采取循环推理1000次取平均时间，比较模型推理速度的差异。

注1：矩形输入指不改变宽高比的情况下，最长边resize到320、416、608，短边采用(128,128,128)填充至32的倍数。如输入尺寸为（720,1080）的720P图像，320、416、608分别对应（192, 320）、（256, 416）、（352, 608）。

注2：model3-6括号中的时间代表纯前向推理时间，batch size 均为1。

Model rec(1/4/8/16) [ms] Square(1/4/8/16) [ms]
Model1（416） 14.4/32.8/59.7/108.9 17.9/51.9/97.6/166.6
Model1（608） 21.3/63.8/122.5/196.9 31.1/105.3/205.5/322.2
Model1（320） 14.1/22.2/38.2/64.5 14.8/32.9/60.7/106.7
Model2（416） 13.9/21.6/39.9/76.6 14.4/33.4/63.9/124.7
Model2（608） 13.9/40.2/76.8/151.3 19.2/67.4/131.1/259.5
Model2（320） 13.9/14.0/24.1/44.9 14.3/20.9/38.3/73.6
Model3和Model4（416） 10.8（9.4） 16.1（14.0）
Model5和Model6（416） 8.4（6.4） 12.1（9.7）

四、测试结论

  1. batch size 为 1 不能有效的利用显卡资源。通过多batch size 比较时间，矩形比正方形加速40%。(输入图片是16：9，矩形相对于正方形计算量少了43%）恰好证明了结论2。

  2. 当输入面积（N * H * W）大于一定值时，速度瓶颈才会出现在FLOPs上。详情请见，Roofline Model与深度学习模型的性能分析，模型的推理速度是由计算量为A且访存量为B的模型在算力为C且带宽为D的计算平台所能达到的理论性能上限E是多少决定的。

  3. Tensorrt 可以降低模型的访存量，加速25%，但是似乎对模型计算量的降低非常有限。

  4. 通道剪枝可以降低模型的计算量和访存量，加速35%，但是在batch size 为 1时加速不明显，说明这次通道剪枝并没有优化模型的计算密度。和K_i和K_o的裁剪比例及大小有关。

  5. 若显卡的算力足够，提高batch size 为 1时的推理速度，核心是优化模型的访存量；提高多batch size时的推理速度核心是优化模型的计算量。

QA:

Q1：在做了矩形框输入和剪枝后，为什么batch size 为1时，320/416/608三种输入分辨率的前向推理的速度基本相同？


image.png

猜测：在计算量未达到GPU的瓶颈时，仅提高输入的分辨率而不改变模型，虽然计算密度基本不变（略微提高），但是实际带宽利用率会变大，所以FLOPS和FLOPs都会以相同的比例增长，因此最后的前向推理时间相同。
pi = \frac {K_i \times K_o \times M^2 \times F^2}{K_i \times M_2 + K_o \times F^2 }
其中\pi 表示计算密度，K_i 表示输入通道数，K_o 表示输出通道数，M^2 表示卷积核的平面尺寸，F^2 表示输出特征图的平面尺寸。仅当F^2 改变时，计算密度基本不变，但是\beta表示实际带宽利用率变大，使得在相同的计算密度下FLOPs增大。

Q2：多batch size进行推理，为什么加速明显？
猜测：在多batch size中，FLOPS成倍增长但是时间却没有成倍增长，所以FLOPs提高。可能和Q1一样能够提高实际带宽利用率，并且由于模型参数均一致可以放入缓存中，同时降低访存量，提高了计算密度。

pi = \frac {B \times K_i \times K_o \times M^2 \times F^2 }{K_i \times M_2 + B \times K_o \times F^2 }

Q3：通道剪枝之后的模型，在320/416和未进行剪枝对比，batch size 为 1时前向推理速度基本相同？
猜测：FLOPS降低但是时间基本相同，说明FLOPs也同比例降低。

上一篇下一篇

猜你喜欢

热点阅读