YOLOV3剪枝

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

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

主要思路

• 1、利用batch normalization中的縮放因子γ 作為重要性因子，即γ越小轮洋，所對應(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秕磷。

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

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

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

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

• 將需要裁剪的層的BN層的γ參數(shù)的絕對值提取到一個(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ǔ)上，完成對shortcut層的剪枝箩绍，同時(shí)避免裁剪掉整個(gè)層孔庭。

• 1、找到需要裁剪的BN層的對應(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的卷積層對應(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ù)的絕對值提取到一個(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層對應(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ǔ)上衍生出來的，針對每一個(gè)shortcut層前一個(gè)CBL進(jìn)行評價(jià)搓幌，對各層的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)中對簡單的數(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í)對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階段開始對85%的通道保持原力度壓縮殷蛇，15%的通道進(jìn)行s衰減100倍。這個(gè)85%是個(gè)先驗(yàn)知識(shí)橄浓，是由策略一稀疏后嘗試剪通道幾乎不掉點(diǎn)的最大比例粒梦，幾乎不掉點(diǎn)指的是相對稀疏后精度；如果微調(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夫植。

核心思想：

• 對于obj和分類損失：將學(xué)生模型和老師模型的obj和分類的輸出展開為一維向量,計(jì)算KL散度損失
• 對于Box損失：將學(xué)生模型xywh的輸出x_offset，y_offset油讯，w/grid_cell_w详民，h/grid_cell_h（這里是否有數(shù)量級的問題，似乎用歸一化的歐式距離更好）分別和老師模型的輸出陌兑、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é)：

一心肪、測試環(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

二硬鞍、測試模型（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)換）

三、測試結(jié)果

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

注1：矩形輸入指不改變寬高比的情況下罚随，最長邊resize到320玉工、416、608淘菩，短邊采用(128,128,128)填充至32的倍數(shù)遵班。如輸入尺寸為（720,1080）的720P圖像，320潮改、416狭郑、608分別對應(yīng)（192, 320）、（256, 416）汇在、（352, 608）翰萨。

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

四亩鬼、測試結(jié)論

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

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

3. Tensorrt 可以降低模型的訪存量咆瘟，加速25%嚼隘，但是似乎對模型計(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

猜測：在計(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)行推理，為什么加速明顯潜必？
猜測：在多batch size中靴姿，F(xiàn)LOPS成倍增長但是時(shí)間卻沒有成倍增長，所以FLOPs提高磁滚》鹣牛可能和Q1一樣能夠提高實(shí)際帶寬利用率，并且由于模型參數(shù)均一致可以放入緩存中垂攘，同時(shí)降低訪存量维雇，提高了計(jì)算密度。

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