YOLOV3剪枝

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

剪枝項(xiàng)目參考https://github.com/tanluren/yolov3-channel-and-layer-pruning

主要思路

• 1、利用batch normalization中的縮放因子γ 作為重要性因子，即γ越小，所對(duì)應(yīng)的channel不太重要，就可以裁剪（pruning）。

• 2、約束γ的大小，在目標(biāo)方程中增加一個(gè)關(guān)于γ的L1正則項(xiàng)，使其稀疏化，這樣可以做到在訓(xùn)練中自動(dòng)剪枝，這是以往模型壓縮所不具備的。

剪枝過程

img

這里寫圖片描述

分為三部分，第一步，訓(xùn)練；第二步，剪枝；第三步，微調(diào)剪枝后的模型，循環(huán)執(zhí)行

YOLOV3剪枝源碼

1、正常剪枝

這部分分析來自該倉庫https://github.com/coldlarry/YOLOv3-complete-pruning，但是更新的倉庫也可以完成正常剪枝，prune.py。

使用了正常剪枝模式，不對(duì)short cut層（需要考慮add操作的維度一致問題）及上采樣層（無BN）進(jìn)行裁剪。

• 1、找到需要裁剪的BN層的對(duì)應(yīng)的索引。

• 2、每次反向傳播前，將L1正則產(chǎn)生的梯度添加到BN層的梯度中。

• 3、設(shè)置裁剪率進(jìn)行裁剪。

• 將需要裁剪的層的BN層的γ參數(shù)的絕對(duì)值提取到一個(gè)列表并從小到大進(jìn)行排序，若裁剪率0.8，則列表中0.8分位數(shù)的值為裁剪閾值。

• 將小于裁剪閾值的通道的γ置為0，驗(yàn)證裁剪后的map（并沒有將β置為0）。

• 4、創(chuàng)建新的模型結(jié)構(gòu)，β合并到下一個(gè)卷積層中BN中的running_mean計(jì)算。

• 5、生成新的模型文件。

2、優(yōu)化的正常剪枝

slim_prune.py，在正常剪枝模式的基礎(chǔ)上，完成對(duì)shortcut層的剪枝，同時(shí)避免裁剪掉整個(gè)層。

• 1、找到需要裁剪的BN層的對(duì)應(yīng)的索引。

# 解析模型文件
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前一個(gè)CBL不剪
                ignore_idx.add(i)

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

        elif module_def['type'] == 'shortcut':
            # 根據(jù)cfg中的from層獲得shortcut的卷積層對(duì)應(yīng)的索引
            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

• 2、每次反向傳播前，將L1正則產(chǎn)生的梯度添加到BN層的梯度中。

• 3、設(shè)置裁剪率進(jìn)行裁剪。

  • 將需要裁剪的層的BN層的γ參數(shù)的絕對(duì)值提取到一個(gè)列表并從小到大進(jìn)行排序，若裁剪率0.8，則列表中0.8分位數(shù)的值為裁剪閾值。

    # 提取需要裁剪的層的BN參數(shù)
    bn_weights = gather_bn_weights(model.module_list, prune_idx)
    # 排序
    sorted_bn, sorted_index = torch.sort(bn_weights)
    # 分位數(shù)索引
    thresh_index = int(len(bn_weights) * opt.global_percent)
    thresh = sorted_bn[thresh_index].cuda()
    

  • 將小于裁剪閾值的通道提取出來；如果整層的通道γ均低于閾值，為了避免整層被裁剪，保留該層中γ值最大的幾個(gè)（根據(jù)layer_keep參數(shù)進(jìn)行設(shè)置，最小為1）通道。

    def obtain_filters_mask(model, thresh, CBL_idx, prune_idx):
        pruned = 0
        total = 0
        num_filters = []
        filters_mask = []
        for idx in CBL_idx:
            bn_module = model.module_list[idx][1]
            # 如果該層需要裁剪，則先確定裁剪后的最小通道數(shù)min_channel_num，然后根據(jù)裁剪閾值進(jìn)行通道裁剪確定mask，如果整層的通道γ均低于閾值，為了避免整層被裁剪，留該層中γ值最大的幾個(gè)（根據(jù)layer_keep參數(shù)進(jìn)行設(shè)置，最小為1）通道。
            if idx in prune_idx:
                weight_copy = bn_module.weight.data.abs().clone()
                
                channels = weight_copy.shape[0]
                min_channel_num = int(channels * opt.layer_keep) if int(channels * opt.layer_keep) > 0 else 1
                mask = weight_copy.gt(thresh).float()
                
                if int(torch.sum(mask)) < min_channel_num: 
                    _, sorted_index_weights = torch.sort(weight_copy,descending=True)
                    mask[sorted_index_weights[:min_channel_num]]=1. 
                remain = int(mask.sum())
                pruned = pruned + mask.shape[0] - remain
    
                print(f'layer index: {idx:>3d} \t total channel: {mask.shape[0]:>4d} \t '
                        f'remaining channel: {remain:>4d}')
            # 如果該層不需要裁剪，則全部保留
            else:
                mask = torch.ones(bn_module.weight.data.shape)
                remain = mask.shape[0]
    
            total += mask.shape[0]
            num_filters.append(remain)
            filters_mask.append(mask.clone())
    

  • 合并shortcut層的mask，采用取并集的策略。

    def merge_mask(model, CBLidx2mask, CBLidx2filters):
        # 最后一層開始遍歷
        for i in range(len(model.module_defs) - 1, -1, -1):
            mtype = model.module_defs[i]['type']
            if mtype == 'shortcut':
                if model.module_defs[i]['is_access']:
                    continue
    
                Merge_masks =  []
                layer_i = i
                # 循環(huán)的停止條件是到網(wǎng)絡(luò)的feature map 發(fā)生下采樣時(shí)
                while mtype == 'shortcut':
                    # 標(biāo)志為true
                    model.module_defs[layer_i]['is_access'] = True
    
                    # 如果前一層為卷積層，添加該層上一層卷積層通道的mask
                    if model.module_defs[layer_i-1]['type'] == 'convolutional':
                        bn = int(model.module_defs[layer_i-1]['batch_normalize'])
                        if bn:
                            Merge_masks.append(CBLidx2mask[layer_i-1].unsqueeze(0))
                    # 找到和該層shortcut連接的層的索引
                    layer_i = int(model.module_defs[layer_i]['from'])+layer_i
                    mtype = model.module_defs[layer_i]['type']
    
                    # 如果和shortcut連接的層為卷積層，則添加該層通道的mask；否則進(jìn)入下一次while循環(huán)
                    if mtype == 'convolutional':
                        bn = int(model.module_defs[layer_i]['batch_normalize'])
                        if bn:
                            Merge_masks.append(CBLidx2mask[layer_i].unsqueeze(0))
    
                # 綜合考慮所有feature map 大小相同（即通道數(shù)相同，不發(fā)生下采樣）的shortcut層對(duì)應(yīng)的卷積層通道的mask，只要一個(gè)為true則全部不剪裁
                if len(Merge_masks) > 1:
                    Merge_masks = torch.cat(Merge_masks, 0)
                    merge_mask = (torch.sum(Merge_masks, dim=0) > 0).float()
                else:
                    merge_mask = Merge_masks[0].float()
    
                layer_i = i
                mtype = 'shortcut'
                # 更新新的merge_mask
                while mtype == 'shortcut':
    
                    if model.module_defs[layer_i-1]['type'] == 'convolutional':
                        bn = int(model.module_defs[layer_i-1]['batch_normalize'])
                        if bn:
                            CBLidx2mask[layer_i-1] = merge_mask
                            CBLidx2filters[layer_i-1] = int(torch.sum(merge_mask).item())
    
                    layer_i = int(model.module_defs[layer_i]['from'])+layer_i
                    mtype = model.module_defs[layer_i]['type']
    
                    if mtype == 'convolutional':
                        bn = int(model.module_defs[layer_i]['batch_normalize'])
                        if bn:
                            CBLidx2mask[layer_i] = merge_mask
                            CBLidx2filters[layer_i] = int(torch.sum(merge_mask).item())
    
    

• 4、驗(yàn)證裁剪模型之后的MAP。

• 5、實(shí)際裁剪模型參數(shù)，β合并到下一個(gè)卷積層中BN中的running_mean計(jì)算。驗(yàn)證MAP，比較模型參數(shù)量及inference速度

• 6、創(chuàng)建新的模型結(jié)構(gòu)，保存新的cfg及weights。

3、層剪枝

和優(yōu)化的正常剪枝類似。這個(gè)策略是在之前的通道剪枝策略基礎(chǔ)上衍生出來的，針對(duì)每一個(gè)shortcut層前一個(gè)CBL進(jìn)行評(píng)價(jià)，對(duì)各層的Gmma均值進(jìn)行排序，取最小的進(jìn)行層剪枝。為保證yolov3結(jié)構(gòu)完整，這里每剪一個(gè)shortcut結(jié)構(gòu)，會(huì)同時(shí)剪掉一個(gè)shortcut層和它前面的兩個(gè)卷積層。是的，這里只考慮剪主干中的shortcut模塊。但是yolov3中有23處shortcut，剪掉8個(gè)shortcut就是剪掉了24個(gè)層，剪掉16個(gè)shortcut就是剪掉了48個(gè)層，總共有69個(gè)層的剪層空間；實(shí)驗(yàn)中對(duì)簡(jiǎn)單的數(shù)據(jù)集剪掉了較多shortcut而精度降低很少。

稀疏策略

scale參數(shù)默認(rèn)0.001，根據(jù)數(shù)據(jù)集，mAP,BN分布調(diào)整，數(shù)據(jù)分布廣類別多的，或者稀疏時(shí)掉點(diǎn)厲害的適當(dāng)調(diào)小s;-sr用于開啟稀疏訓(xùn)練；--prune 0適用于prune.py，--prune 1 適用于其他剪枝策略。稀疏訓(xùn)練就是精度和稀疏度的博弈過程，如何尋找好的策略讓稀疏后的模型保持高精度同時(shí)實(shí)現(xiàn)高稀疏度是值得研究的問題，大的s一般稀疏較快但精度掉的快，小的s一般稀疏較慢但精度掉的慢；配合大學(xué)習(xí)率會(huì)稀疏加快，后期小學(xué)習(xí)率有助于精度回升。
注意：訓(xùn)練保存的pt權(quán)重包含epoch信息，可通過python -c "from models import *; convert('cfg/yolov3.cfg', 'weights/last.pt')"轉(zhuǎn)換為darknet weights去除掉epoch信息，使用darknet weights從epoch 0開始稀疏訓(xùn)練。

1、恒定s

這是一開始的策略，也是默認(rèn)的策略。在整個(gè)稀疏過程中，始終以恒定的s給模型添加額外的梯度，因?yàn)榱Χ缺容^均勻，往往壓縮度較高。但稀疏過程是個(gè)博弈過程，我們不僅想要較高的壓縮度，也想要在學(xué)習(xí)率下降后恢復(fù)足夠的精度，不同的s最后稀疏結(jié)果也不同，想要找到合適的s往往需要較高的時(shí)間成本。

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

2、全局s衰減

關(guān)鍵代碼是下面這句，在epochs的0.5階段s衰減100倍。前提是0.5之前權(quán)重已經(jīng)完成大幅壓縮，這時(shí)對(duì)s衰減有助于精度快速回升，但是相應(yīng)的bn會(huì)出現(xiàn)一定膨脹，降低壓縮度，有利有弊，可以說是犧牲較大的壓縮度換取較高的精度，同時(shí)減少尋找s的時(shí)間成本。當(dāng)然這個(gè)0.5和100可以自己調(diào)整。注意也不能為了在前半部分加快壓縮bn而大大提高s，過大的s會(huì)導(dǎo)致模型精度下降厲害，且s衰減后也無法恢復(fù)。如果想使用這個(gè)策略，可以在prune_utils.py中的BNOptimizer把下面這句取消注釋。

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

3、局部s衰減

關(guān)鍵代碼是下面兩句，在epochs的0.5階段開始對(duì)85%的通道保持原力度壓縮，15%的通道進(jìn)行s衰減100倍。這個(gè)85%是個(gè)先驗(yàn)知識(shí)，是由策略一稀疏后嘗試剪通道幾乎不掉點(diǎn)的最大比例，幾乎不掉點(diǎn)指的是相對(duì)稀疏后精度；如果微調(diào)后還是不及baseline，或者說達(dá)不到精度要求，就可以使用策略三進(jìn)行局部s衰減，從中間開始重新稀疏，這可以在犧牲較小壓縮度情況下提高較大精度。如果想使用這個(gè)策略可以在train.py中把下面這兩句取消注釋，并根據(jù)自己策略一情況把0.85改為自己的比例，還有0.5和100也是可調(diào)的。策略二和三不建議一起用，除非你想做組合策略。

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

知識(shí)蒸餾

參考論文：Learning Efficient Object Detection Models with Knowledge Distillation。

核心思想：

• 對(duì)于obj和分類損失：將學(xué)生模型和老師模型的obj和分類的輸出展開為一維向量,計(jì)算KL散度損失
• 對(duì)于Box損失：將學(xué)生模型xywh的輸出x_offset，y_offset，w/grid_cell_w，h/grid_cell_h（這里是否有數(shù)量級(jí)的問題，似乎用歸一化的歐式距離更好）分別和老師模型的輸出、target計(jì)算L2距離，如果學(xué)生模型的輸出，如果學(xué)生和老師更遠(yuǎn)，學(xué)生會(huì)再向target學(xué)習(xí)，而不是向老師學(xué)習(xí)。這時(shí)候老師的輸出是hard label。

image.png

def distillation_loss2(model, targets, output_s, output_t):
    '''
    :param model:  學(xué)生模型
    :param targets: 標(biāo)簽
    :param output_s: 學(xué)生模型的輸出
    :param output_t: 老師模型的輸出
    :return: 附加Loss
    '''
    reg_m = 0.0
    T = 3.0
    Lambda_cls, Lambda_box = 0.0001, 0.001
        
    # KL 損失，衡量兩個(gè)分布的差異
    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])
        # 標(biāo)簽轉(zhuǎn)換
    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

自己項(xiàng)目實(shí)驗(yàn)的總結(jié)：

一、測(cè)試環(huán)境

• 宿主機(jī)：Ubuntu 16.04， Docker環(huán)境：Ubuntu 16.04.6 LTS
• CPU：32 Intel(R) Xeon(R) CPU E5-2630 v3 @ 2.40GHz
• GPU：GTX 1080 8G
• CUDA_CUDNN：10.1.243，7.6.3.30
• Python：3.6.9
• onnx：1.4.1
• tensorrt：6.0.1.5
• pytorch：1.3.0

二、測(cè)試模型（YOLOV3）

模型采用pytorch框架訓(xùn)練、剪枝，轉(zhuǎn)換為darknet模型，再轉(zhuǎn)換成ONNX模型，最后轉(zhuǎn)換為tensorrt模型。將以下六個(gè)模型依次命名為model1、model2、...、model6。

1. Pytorch模型：model.pt
2. Pytorch剪枝且微調(diào)模型：model_075.weights（轉(zhuǎn)換為darknet格式）
3. Tensorrt 模型：model_416416.trt（由model1轉(zhuǎn)換）
4. Tensorrt 模型：model_256416.trt（由model1轉(zhuǎn)換）
5. Tensorrt 模型：model_075_256416.trt（由model2轉(zhuǎn)換）
6. Tensorrt 模型：model_075_416416.trt（由model2轉(zhuǎn)換）

三、測(cè)試結(jié)果

通過采用不同的輸入形狀（矩形和正方形）、不同的輸入尺寸（320、416、608）、不同Batch Size（1、4、8、16）下對(duì)同一張圖片采取循環(huán)推理1000次取平均時(shí)間，比較模型推理速度的差異。

注1：矩形輸入指不改變寬高比的情況下，最長邊resize到320、416、608，短邊采用(128,128,128)填充至32的倍數(shù)。如輸入尺寸為（720,1080）的720P圖像，320、416、608分別對(duì)應(yīng)（192, 320）、（256, 416）、（352, 608）。

注2：model3-6括號(hào)中的時(shí)間代表純前向推理時(shí)間，batch size 均為1。

四、測(cè)試結(jié)論

1. batch size 為 1 不能有效的利用顯卡資源。通過多batch size 比較時(shí)間，矩形比正方形加速40%。(輸入圖片是16：9，矩形相對(duì)于正方形計(jì)算量少了43%）恰好證明了結(jié)論2。

2. 當(dāng)輸入面積（N * H * W）大于一定值時(shí)，速度瓶頸才會(huì)出現(xiàn)在FLOPs上。詳情請(qǐng)見，Roofline Model與深度學(xué)習(xí)模型的性能分析，模型的推理速度是由計(jì)算量為A且訪存量為B的模型在算力為C且?guī)挒镈的計(jì)算平臺(tái)所能達(dá)到的理論性能上限E是多少決定的。

3. Tensorrt 可以降低模型的訪存量，加速25%，但是似乎對(duì)模型計(jì)算量的降低非常有限。

4. 通道剪枝可以降低模型的計(jì)算量和訪存量，加速35%，但是在batch size 為 1時(shí)加速不明顯，說明這次通道剪枝并沒有優(yōu)化模型的計(jì)算密度。和和的裁剪比例及大小有關(guān)。

5. 若顯卡的算力足夠，提高batch size 為 1時(shí)的推理速度，核心是優(yōu)化模型的訪存量；提高多batch size時(shí)的推理速度核心是優(yōu)化模型的計(jì)算量。

QA:

Q1：在做了矩形框輸入和剪枝后，為什么batch size 為1時(shí)，320/416/608三種輸入分辨率的前向推理的速度基本相同？

image.png

猜測(cè)：在計(jì)算量未達(dá)到GPU的瓶頸時(shí)，僅提高輸入的分辨率而不改變模型，雖然計(jì)算密度基本不變（略微提高），但是實(shí)際帶寬利用率會(huì)變大，所以FLOPS和FLOPs都會(huì)以相同的比例增長，因此最后的前向推理時(shí)間相同。

其中 表示計(jì)算密度， 表示輸入通道數(shù)， 表示輸出通道數(shù)， 表示卷積核的平面尺寸， 表示輸出特征圖的平面尺寸。僅當(dāng) 改變時(shí)，計(jì)算密度基本不變，但是表示實(shí)際帶寬利用率變大，使得在相同的計(jì)算密度下FLOPs增大。

Q2：多batch size進(jìn)行推理，為什么加速明顯？
猜測(cè)：在多batch size中，F(xiàn)LOPS成倍增長但是時(shí)間卻沒有成倍增長，所以FLOPs提高?？赡芎蚎1一樣能夠提高實(shí)際帶寬利用率，并且由于模型參數(shù)均一致可以放入緩存中，同時(shí)降低訪存量，提高了計(jì)算密度。

Q3：通道剪枝之后的模型，在320/416和未進(jìn)行剪枝對(duì)比，batch size 為 1時(shí)前向推理速度基本相同？
猜測(cè)：FLOPS降低但是時(shí)間基本相同，說明FLOPs也同比例降低。