YOLOv3通道+層剪枝,參數壓縮98%,砍掉48個層,提速2倍

儘管目標檢測新算法層出不窮,但在實際工程項目中不少52CV羣友還是念着YOLOv3的好。將其部署到邊緣設備等時,模型剪枝是非常有必要的,畢竟有原始模型有239M的參數,剪枝後往往也能提速不少。

YOLOv3模型剪枝,瘦身80%,提速100%,精度基本不變

YOLOv3剪枝再升級!

上述剪枝是減少模型通道數,而今天向大家介紹的工程可以實現通道和層的雙向剪枝,在oxford hand 數據集hand檢測問題中,作者實驗中可以實現精度下降很小而參數減少 98%,砍掉 48 個層,提速 2 倍!(不同問題參數減少和提速比例不同,在作者另一個私人項目中,實現了提速 3 倍。)

作者已將其開源,僅需要幾行命令就可以在自己數據集中實現更加靈活和搜索空間更大的剪枝。

本項目以ultralytics/yolov3(https://github.com/ultralytics/yolov3)爲基礎實現,根據論文Learning Efficient Convolutional Networks Through Network Slimming (ICCV 2017)原理基於bn層Gmma係數進行通道剪枝,下面引用了幾種不同的通道剪枝策略,並對原策略進行了改進,提高了剪枝率和精度;在這些工作基礎上,又衍生出了層剪枝,本身通道剪枝已經大大減小了模型參數和計算量,降低了模型對模型資源的佔用,而層剪枝可以進一步減小了計算量,並大大提高了模型推理速度;通過層剪枝和通道剪枝結合,可以壓縮模型的深度和寬度,某種意義上實現了針對不同數據集的小模型搜索。

項目的基本工作流程是,使用yolov3訓練自己數據集,達到理想精度後進行稀疏訓練,稀疏訓練是重中之重,對需要剪枝的層對應的bn gamma係數進行大幅壓縮,理想的壓縮情況如下圖,然後就可以對不重要的通道或者層進行剪枝,剪枝後可以對模型進行微調恢復精度,後續會寫篇博客記錄一些實驗過程及調參經驗,在此感謝行雲大佬(https://github.com/zbyuan)的討論和合作!

更新

1. 增加了對yolov3-spp結構的支持,基礎訓練可以直接使用yolov3-spp.weights初始化權重,各個層剪枝及通道剪枝腳本的使用也和yolov3一致。

2. 增加了多尺度推理支持,train.py和各剪枝腳本都可以指定命令行參數, 如 --img_size 608 .

基礎訓練

環境配置查看requirements.txt,數據準備參考這裏(https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data),預訓練權重可以從darknet官網下載。

用yolov3訓練自己的數據集,修改cfg,配置好data,用yolov3.weights初始化權重。

python train.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/yolov3.weights --epochs 100 --batch-size 32

稀疏訓練

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開始稀疏訓練。

python train.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.weights --epochs 300 --batch-size 32 -sr --s 0.001 --prune 1

通道剪枝策略一

策略源自Lam1360/YOLOv3-model-pruning(https://github.com/Lam1360/YOLOv3-model-pruning),這是一種保守的策略,因爲yolov3中有五組共23處shortcut連接,對應的是add操作,通道剪枝後如何保證shortcut的兩個輸入維度一致,這是必須考慮的問題。而Lam1360/YOLOv3-model-pruning對shortcut直連的層不進行剪枝,避免了維度處理問題,但它同樣實現了較高剪枝率,對模型參數的減小有很大幫助。雖然它剪枝率最低,但是它對剪枝各細節的處理非常優雅,後面的代碼也較多參考了原始項目。在本項目中還更改了它的閾值規則,可以設置更高的剪枝閾值。

python prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --percent 0.85

通道剪枝策略二

策略源自coldlarry/YOLOv3-complete-pruning(https://github.com/coldlarry/YOLOv3-complete-pruning),這個策略對涉及shortcut的卷積層也進行了剪枝,剪枝採用每組shortcut中第一個卷積層的mask,一共使用五種mask實現了五組shortcut相關卷積層的剪枝,進一步提高了剪枝率。本項目中對涉及shortcut的剪枝後激活偏移值處理進行了完善,並修改了閾值規則,可以設置更高剪枝率,當然剪枝率的設置和剪枝後的精度變化跟稀疏訓練有很大關係,這裏再次強調稀疏訓練的重要性。

python shortcut_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --percent 0.6

通道剪枝策略三

策略參考自PengyiZhang/SlimYOLOv3(https://github.com/PengyiZhang/SlimYOLOv3),這個策略的通道剪枝率最高,先以全局閾值找出各卷積層的mask,然後對於每組shortcut,它將相連的各卷積層的剪枝mask取並集,用merge後的mask進行剪枝,這樣對每一個相關層都做了考慮,同時它還對每一個層的保留通道做了限制,實驗中它的剪枝效果最好。在本項目中還對激活偏移值添加了處理,降低剪枝時的精度損失。

python slim_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --global_percent 0.8 --layer_keep 0.01

層剪枝

這個策略是在之前的通道剪枝策略基礎上衍生出來的,針對每一個shortcut層前一個CBL進行評價,對各層的Gmma最高值進行排序,取最小的進行層剪枝。爲保證yolov3結構完整,這裏每剪一個shortcut結構,會同時剪掉一個shortcut層和它前面的兩個卷積層。是的,這裏只考慮剪主幹中的shortcut模塊。但是yolov3中有23處shortcut,剪掉8個shortcut就是剪掉了24個層,剪掉16個shortcut就是剪掉了48個層,總共有69個層的剪層空間;實驗中對簡單的數據集剪掉了較多shortcut而精度降低很少。

python layer_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --shortcuts 12

同時剪層和通道

前面的通道剪枝和層剪枝已經分別壓縮了模型的寬度和深度,可以自由搭配使用,甚至迭代式剪枝,調配出針對自己數據集的一副良藥。這裏整合了一個同時剪層和通道的腳本,方便對比剪枝效果,有需要的可以使用這個腳本進行剪枝。

python layer_channel_prune.py --cfg cfg/my_cfg.cfg --data data/my_data.data --weights weights/last.pt --shortcuts 12 --global_percent 0.8 --layer_keep 0.1

微調finetune

剪枝的效果好不好首先還是要看稀疏情況,而不同的剪枝策略和閾值設置在剪枝後的效果表現也不一樣,有時剪枝後模型精度甚至可能上升,而一般而言剪枝會損害模型精度,這時候需要對剪枝後的模型進行微調,讓精度回升。訓練代碼中默認了前6個epoch進行warmup,這對微調有好處,有需要的可以自行調整超參學習率。

python train.py --cfg cfg/prune_0.85_my_cfg.cfg --data data/my_data.data --weights weights/prune_0.85_last.weights --epochs 100 --batch-size 32

tensorboard實時查看訓練過程

tensorboard --logdir runs

 

案例

使用yolov3-spp訓練oxford hand數據集並剪枝。下載數據集(http://www.robots.ox.ac.uk/~vgg/data/hands/downloads/hand_dataset.tar.gz),解壓到data文件夾,運行converter.py,把得到的train.txt和valid.txt路徑更新在oxfordhand.data中。通過以下代碼分別進行基礎訓練和稀疏訓練:

python train.py --cfg cfg/yolov3-spp-hand.cfg --data data/oxfordhand.data --weights weights/yolov3-spp.weights --batch-size 20 --epochs 100

 

python -c "from models import *; convert('cfg/yolov3.cfg', 'weights/last.pt')"python train.py --cfg cfg/yolov3-spp-hand.cfg --data data/oxfordhand.data --weights weights/converted.weights --batch-size 20 --epochs 300 -sr --s 0.001 --prune 1

訓練的情況如下圖,藍色線是基礎訓練,紅色線是稀疏訓練。其中基礎訓練跑了100個epoch,後半段已經出現了過擬合,最終得到的baseline模型mAP爲0.84;稀疏訓練以s0.001跑了300個epoch,選擇的稀疏類型爲prune 1全局稀疏,爲包括shortcut的剪枝做準備,並且在總epochs的0.7和0.9階段進行了Gmma爲0.1的學習率衰減,稀疏過程中模型精度起伏較大,在學習率降低後精度出現了回升,最終稀疏模型mAP 0.797。

 

再來看看bn的稀疏情況,代碼使用tensorboard記錄了參與稀疏的bn層的Gmma權重變化,下圖左邊看到正常訓練時Gmma總體上分佈在1附近類似正態分佈,右邊可以看到稀疏過程Gmma大部分逐漸被壓到接近0,接近0的通道其輸出值近似於常量,可以將其剪掉。

 

這時候便可以進行剪枝,這裏例子使用layer_channel_prune.py同時進行剪通道和剪層,這個腳本融合了slim_prune剪通道策略和layer_prune剪層策略。Global perent剪通道的全局比例爲0.93,layer keep每層最低保持通道數比例爲0.01,shortcuts剪了16個,相當於剪了48個層(32個CBL,16個shortcut);下圖結果可以看到剪通道後模型掉了一個點,而大小從239M壓縮到5.2M,剪層後mAP掉到0.53,大小壓縮到4.6M,模型參數減少了98%,推理速度也從16毫秒減到6毫秒(Tesla P100測試結果)。

python layer_channel_prune.py --cfg cfg/yolov3-spp-hand.cfg --data data/oxfordhand.data --weights weights/last.pt --global_percent 0.93 --layer_keep 0.01 --shortcuts 16

鑑於模型精度出現了下跌,我們來進行微調,下面是微調50個epoch的結果,精度恢復到了0.793,bn也開始呈正態分佈,這個結果相對於baseline掉了幾個點,但是模型大幅壓縮減少了資源佔用,提高了運行速度。如果想提高精度,可以嘗試降低剪枝率,比如這裏只剪10個shortcut的話,同樣微調50epoch精度可以回到0.81;而想追求速度的話,這裏有個極端例子,全局剪0.95,層剪掉54個,模型壓縮到了2.8M,推理時間降到5毫秒,而mAP降到了0,但是微調50 epochs 後依然回到了0.75。

python train.py --cfg cfg/prune_16_shortcut_prune_0.93_keep_0.01_yolov3-spp-hand.cfg --data data/oxfordhand.data --weights weights/prune_16_shortcut_prune_0.93_keep_0.01_last.weights --batch-size 52 --epochs 50

可以猜測,剪枝得到的cfg是針對該數據集相對合理的結構,而保留的權重可以讓模型快速訓練接近這個結構的能力上限,這個過程類似於一種有限範圍的結構搜索。而不同的訓練策略,稀疏策略,剪枝策略會得到不同的結果,相信即使是這個例子也可以進一步壓縮並保持良好精度。yolov3有衆多優化項目和工程項目,可以利用這個剪枝得到的cfg和weights放到其他項目中做進一步優化和應用。

這裏(https://pan.baidu.com/s/1APUfwO4L69u28Wt9gFNAYw)分享了這個例子的權重和cfg,包括baseline,稀疏,不同剪枝設置後的結果。

工程地址:

https://github.com/tanluren/yolov3-channel-and-layer-pruning

 

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章