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

我愛計算機視覺發表於2019-11-27

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

比如之前52CV曾經分享過:

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

YOLOv3剪枝再升級!

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

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

工程地址:

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

歡迎大家給大佬加Star!


以下內容為原作者52CV群友漂的投稿。

本專案以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)的討論和合作!

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

更新

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通道+層剪枝,引數壓縮98%,砍掉48個層,提速2倍!


歡迎使用和測試,有問題或者交流實驗過程可以發issue或者q我1806380874

案例

使用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。


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

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

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


這時候便可以進行剪枝,這裡例子使用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

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

鑑於模型精度出現了下跌,我們來進行微調,下面是微調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

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

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

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


工程地址:

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

歡迎大家給大佬加Star!

相關文章