作者｜機器之心編輯部
來源｜機器之心
本文提出了一套模型壓縮和編譯結合的目標檢測加速框架，根據編譯器的硬件特性而設計的剪枝策略能夠在維持高 mAP 的同時大大提高運行速度，壓縮了 14 倍的 YOLOv4 能夠在手機上達到 19FPS 的運行速度並且依舊維持 49mAP（COCO dataset）的高準確率。相比 YOLOv3 完整版，該框架快出 7 倍，並且沒有犧牲準確率。該框架由美國東北大學王言治研究組和威廉瑪麗學院任彬研究組共同提出。
隨着近年來 CNN 在目標檢測領域的發展和創新，目標檢測有了更加廣泛的應用。考慮到在實際場景中的落地需求，目標檢測網絡往往需要在保持高準確率的同時擁有較低的計算延遲。而現有的目標檢測網絡，在資源有限的平臺上，尤其是手機和嵌入式設備上部署這類應用時，很難同時實現高準確率與實時檢測。
在過去的幾年中，有很多優秀的目標檢測網絡被相繼提出。其中的 two-stage 檢測網絡包括 RCNN 系列和 SPPNet 等，還有 one-stage 檢測網絡如 YOLO 系列，SSD 和 Retina-Net 等。相較於 two-stage 網絡，one-stage 網絡在犧牲一定準確率的情況下換來了更快的執行速度。即便如此，這些網絡依然需要較大的計算量來達到可接受的準確率，這成爲了這些網絡難以在移動設備上實現實時推理的主要阻礙。爲此，一些輕量級 (lightweight) 目標檢測網絡模型被提出，如 SSD-Lite, YOLO-Lite, YOLO-tiny 等，以實現移動設備上的快速目標檢測。但是這些輕量級網絡的解決方案效果依然不理想。
爲了更好地滿足目標檢測框架的落地需求，CoCoPIE 團隊：美國東北大學王言治研究組和威廉瑪麗學院任彬研究組共同提出了名爲 YOLObile 的手機端目標檢測加速框架。YOLObile 框架通過「壓縮 - 編譯」協同設計在手機端實現了高準確率實時目標檢測。該框架使用一種新提出的名爲「block-punched」的權重剪枝方案，對模型進行有效的壓縮。在編譯器優化技術的協助下，在手機端實現高準確率的實時目標檢測。該研究還提出了一種高效的手機 GPU - 手機 CPU 協同計算優化方案，進一步提高了計算資源的利用率和執行速度。
相比 YOLOv4 的原版，加速後的 YOLObile 的運行速度提高了 4 倍，並且維持了 49mAP 的準確率。相比 YOLOv3 完整版，該框架快 7 倍，在手機上實現了 19FPS 的實時高準確率目標檢測。同時準確率高於 YOLOv3，並沒有用犧牲準確率來提高計算速度。
論文標題：
YOLObile: Real-Time Object Detection on Mobile Devices via Compression-Compilation Co-Design
論文鏈接：
https://arxiv.org/abs/2009.05697
代碼鏈接：
https://github.com/nightsnack/YOLObile
demo鏈接：
https://www.bilibili.com/video/BV1bz4y1y7CR
研究方法
替換硬件支持性不好的操作符
在原版的 YOLOv4 中，有一些操作符不能夠最大化地利用硬件設備的執行效率，比如帶有指數運算的激活函數可能會造成運行的延遲增加，成爲降低延時提高效率的瓶頸。該研究把這些操作符相應地替換成對硬件更加友好的版本，還有一些操作符是 ONNX 還未支持的（YOLObile 用 ONNX 作爲模型的存儲方式），研究者把它替換成 ONNX 支持的運算符。
比如，在 YOLOv4 引入的新模塊 Spatial Attention Module （SAM）中，用了 sigmoid 作爲分支的激活函數，該研究在嘗試把它換成 hard-sigmoid 後發現：準確率和直接刪除相比幾乎一致，增加的模塊又會增加計算量，所以研究者將其刪除了。
Mish 激活函數也涉及了指數運算，同時在 pytorch 上的支持不太友好，會在訓練時佔用很多緩存，同時它在 pytorch 上也不能夠像 C++ 版本的 YOLOv4 一樣帶來很大的準確率提升，而且 ONNX 尚未支持。研究者將其換成了 YOLOv3 的 leaky RELU。
「block-punched」權重剪枝方案
在 YOLObile 優化框架中，作者使用了新提出的名爲「block-punched」的權重剪枝 (weight pruning) 方案。這種剪枝方案旨在在獲得較高剪枝結構自由度的同時，還能使剪枝後的模型結構較好地利用硬件並行計算。這樣就從兩方面分別保證了剪枝後模型的準確率和較高的運算速度。
在這種剪枝方案中，每層的權重矩陣將被劃分爲大小相等的多個塊(block)，因此，每個塊中將包含來自 m 個 filter 的 n 個 channel 的權重。在每個塊中，被剪枝的位置需要做出如下限定：需要修剪所有 filter 相同位置的一個或多個權重，同時也修剪所有通道相同位置的一個或多個權重。
從另一個角度來看，這種剪枝方案將權重的剪枝位置貫穿了整個塊中所有的卷積核(kernel)。對於不同的 block，剪枝的位置和剪掉的 weight 數量都是靈活可變的。另外這種剪枝策略可以應用在卷積核大小 3x3，1x1，5x5 等的卷積層上，也適用於 FC 層。
通過劃分爲固定大小的塊來進行剪枝，能夠提高編譯器的並行度，進而提高在手機上運行的速率。
1. 在準確率方面，通過劃分多個小區塊，這種剪枝方法實現了更加細粒度的剪枝。相較於傳統的結構化剪枝（剪除整個 filter 或 channel），這種方式具有更高的剪枝結構自由度，從而更好地保持了模型的準確率。
2. 在硬件表現方面，因爲在同一小區塊中，所有 filter 修剪被修剪的位置相同，所以在並行計算時，所有 filter 將統一跳過讀取相同的輸入數據，從而減輕處理這些 filter 的線程之間的內存壓力。而限制修剪小區塊內各 channel 的相同位置，確保了所有 channel 共享相同的計算模式，從而消除處理每個 channel 的線程之間的計算差異。因此，這種剪枝方案可以大幅度降低在計算過程中處理稀疏結構的額外開銷，從而達到更好的加速效果。
Reweighted 算法
作者採用了 Reweighted Regularization 算法來篩選出剪枝掉的 weight 的位置，reweight 算法依據 weight 的絕對值大小篩選出相對重要的位置。整個剪枝的過程從依據預先訓練好的網絡模型開始，然後將 Reweighted 算法融合在訓練的過程中，在訓練每 3 到 4 個 epochs 後就能夠得到一組可選的剪枝位置，通過重複訓練來調整剪枝位置，最終得到精確的剪枝模型結果。
編譯器的優化
爲了支持「block-punched」剪枝方案，編譯器也需要作出相應的調整和優化。爲了更好地存儲和調用經過「block-punched」剪枝後的 weight index，研究者引入了新的存儲方案，將 index 的存儲空間進一步壓縮。同時，通過對所有的塊進行重新排序，編譯器能夠在更少的內存訪問次數下運行，進而獲得更快的運算速度。
手機 GPU 和手機 CPU 協同計算的優化方案
YOLObile 中還使用了手機 GPU 和手機 CPU 協同計算的方式來進一步降低整個網絡的運算時間。現在主流的移動端 DNN 推理加速框架，如 TensorFlow-Lite，MNN 和 TVM 都只能支持手機 CPU 或 GPU 單獨運算，因此會導致潛在的計算資源浪費。
YOLObile 提出針對網絡中的分支結構，比如 YOLOv4 中大量使用的 Cross Stage Partial (CSP)結構，使用 CPU 來輔助 GPU 同時進行一些相互無依賴關係的分支運算，從而更好地利用計算資源，減少網絡的運算時間。YOLObile 框架將待優化的網絡分支分爲有卷積運算分支和無卷積運算分支，並對這兩種情況分別給出了優化方案。
如下圖所示，CSP 是一個跨越很多殘差 block 的長分支，在手機的執行過程中，單一的 GPU 計算往往是順序地執行完上面堆疊殘差 block 的分支後再執行下面的 CSP 分支，然後拼接在一起作爲下一層的輸入。研究者將卷積層數更少的 branch2 挪到 CPU 上去，CPU 執行時間少於上面 branch1 在 GPU 上的總運算時間，這個並行操作能夠有效減少運算延遲。
當然，決定能否將 branch2 轉移到 CPU 的因素在於 branch1 的卷積層數多少，通常 CSP block 會跨越 8 個殘差 block，也有的時候出現只跨越 4 個殘差 block 的情況，還有在前幾層只跨一個殘差 block 的情況。對於只跨 1 個殘差 block 的情況明顯還是 GPU 順序執行更高效，對於跨越多個的就需要用實際測出的延遲來做判斷。值得注意的是，轉移數據到不同處理設備的時候，需要加入數據傳輸拷貝的時間。
在 YOLOv4 最後輸出的位置，3 個 YOLO head 輸出部分有很多諸如轉置，變形之類的非卷積運算，這些非卷積運算在 CPU 和 GPU 上的運行效率相當，作者同樣基於運行時間，考慮將部分運算符轉移到 CPU 上去做，選擇最高效的執行方式。
實驗結果
整個訓練過程採用 4 塊 RTX2080Ti，訓練時間爲 5 天。使用配備高通驍龍 855 CPU 和 Adreno 650 GPU 的三星 Galaxy S20 作爲測試平臺。訓練的數據集是 MS COCO，和 YOLOv4 原版保持一致。
實驗結果表明，當使用 YOLOv4 爲基礎模型進行優化時，該研究的優化框架可以成功將原模型大小壓縮至 1/14，在未使用 GPU-CPU 協同計算優化時，將每秒檢測幀數（FPS）提升至 17，且達到 49（mAP）的準確率。
從下圖中可以看到，與衆多具有代表性的目標檢測網絡相比，該研究的優化模型在準確率與速度兩方面同時具有優異的表現，而不再是簡單的犧牲大幅準確率來獲取一定程度的速度提升。
下表展示了 YOLObile 與其他具有代表性的目標檢測網絡在準確率與速度方面的具體比較。值得注意的是，作者的 GPU-CPU 協同計算優化方案可以進一步將執行速度提高至 19FPS。
相比於其他 one-stage 目標檢測模型和 light-weight 模型，經過剪枝後的 YOLObile 更快且更精確，在準確率和速度上實現了帕累托最優。
消融實驗
不同剪枝策略比較
表格中列出了在相同壓縮倍率的情況下，Structured Pruning, Unstructured Pruning, 和文中提出的 Block-Punched Pruning 這三種剪枝策略的準確率和運行效率的比較。
可以看到，Unstructured Pruning 在壓縮 8 倍的狀態下能夠維持較高的準確率，但是和前文描述的一樣，它的運行效率較低，即使在壓縮 8 倍的情況下，運行時間只提高了不到 2 倍。Structured Pruning 後模型能夠高效率地在硬件設備上執行，但是準確率相比 Unstructured Pruning 有大幅度的下降。而根據編譯器設計的 Block-Punched 剪枝策略，能夠維持較高的準確率，並且達到和 Structured Pruning 相當的執行效率。
和 Pattern-Based Pruning 的比較
Pattern-Based Pruning 只能用在 3x3 的層，爲了達到較高的壓縮倍率，就需要對 3x3 的卷積層壓縮掉更多的 weight。
下圖是一個關於不同倍率下 pattern-based pruning 和 block-punched pruning 的比較，可以看出 pattern-based pruning 在較低倍率壓縮的情況下，能夠擁有較高的準確率和執行效率，但是受限制於只適用於 3x3 卷積層，即使將所有的 3x3 層全部壓縮也只能達到 5 倍多的壓縮率，但是準確率就非常低了。block-based pruning 由於可以將所有的層都進行壓縮，在高倍率的情況下依然能夠維持較高的準確率。
Block-Punched 剪枝中分塊大小的比較
該研究對四種不同的塊大小進行實驗，爲了實現較高的硬件並行度，每個模塊中的 channel 數固定爲 4，這與 GPU 和 CPU 矢量寄存器的長度相同。在每個塊中使用不同數量的 filter 來評估準確性和速度。
如圖 5.3 所示，與較小的塊相比，較大的塊可以更好地利用硬件並行性，並可以實現更高的推理速度。但是其粗略的修剪粒度導致準確性下降。較小的塊可以獲得較高的精度，但會犧牲推理速度。根據結果，研究者將 8×4（8 個連續 filter 的 4 個連續 channel）視爲移動設備上的所需塊大小，這在精度和速度之間取得了良好的平衡。
各層的剪枝比例
YOLOv4 在卷積層中僅包含 3×3 和 1×1 的 kernel。該研究提出這兩種類型的卷積層在修剪過程中具有不同的敏感度並進行了兩組實驗。在相同數量的 FLOPs 下，將一組中的所有層平均修剪，另一組中，3×3 卷積層的壓縮率比 1×1 卷積層高 1.15 倍。如上表所示，均勻修剪的模型和不均勻修剪的模型相比，準確性和速度都較低。
????
現在，在「知乎」也能找到我們了
進入知乎首頁搜索「PaperWeekly」
點擊「關注」訂閱我們的專欄吧
關於PaperWeekly
PaperWeekly 是一個推薦、解讀、討論、報道人工智能前沿論文成果的學術平臺。如果你研究或從事 AI 領域，歡迎在公衆號後臺點擊「交流羣」，小助手將把你帶入 PaperWeekly 的交流羣裏。
