（附代碼）YOLObile：手機(jī)端19FPS的實(shí)時(shí)目標(biāo)檢測(cè)算法

點(diǎn)擊左上方藍(lán)字關(guān)注我們

全網(wǎng)搜集目標(biāo)檢測(cè)相關(guān)，人工篩選最優(yōu)價(jià)值內(nèi)容

編者薦語(yǔ)

文章提出了一套模型壓縮和編譯結(jié)合的目標(biāo)檢測(cè)加速框架，根據(jù)編譯器的硬件特性而設(shè)計(jì)的剪枝策略能夠在維持高 mAP 的同時(shí)大大提高運(yùn)行速度，壓縮了 14 倍的 YOLOv4 能夠在手機(jī)上達(dá)到 19FPS 的運(yùn)行速度并且依舊維持 49mAP（COCO dataset）的高準(zhǔn)確率。相比 YOLOv3 完整版，該框架快出 7 倍，并且沒(méi)有犧牲準(zhǔn)確率。

轉(zhuǎn)載自 | 機(jī)器之心

隨著近年來(lái) CNN 在目標(biāo)檢測(cè)領(lǐng)域的發(fā)展和創(chuàng)新，目標(biāo)檢測(cè)有了更加廣泛的應(yīng)用。考慮到在實(shí)際場(chǎng)景中的落地需求，目標(biāo)檢測(cè)網(wǎng)絡(luò)往往需要在保持高準(zhǔn)確率的同時(shí)擁有較低的計(jì)算延遲。而現(xiàn)有的目標(biāo)檢測(cè)網(wǎng)絡(luò)，在資源有限的平臺(tái)上，尤其是手機(jī)和嵌入式設(shè)備上部署這類(lèi)應(yīng)用時(shí)，很難同時(shí)實(shí)現(xiàn)高準(zhǔn)確率與實(shí)時(shí)檢測(cè)。

在過(guò)去的幾年中，有很多優(yōu)秀的目標(biāo)檢測(cè)網(wǎng)絡(luò)被相繼提出。其中的 two-stage 檢測(cè)網(wǎng)絡(luò)包括 RCNN 系列和 SPPNet 等，還有 one-stage 檢測(cè)網(wǎng)絡(luò)如 YOLO 系列，SSD 和 Retina-Net 等。相較于 two-stage 網(wǎng)絡(luò)，one-stage 網(wǎng)絡(luò)在犧牲一定準(zhǔn)確率的情況下?lián)Q來(lái)了更快的執(zhí)行速度。即便如此，這些網(wǎng)絡(luò)依然需要較大的計(jì)算量來(lái)達(dá)到可接受的準(zhǔn)確率，這成為了這些網(wǎng)絡(luò)難以在移動(dòng)設(shè)備上實(shí)現(xiàn)實(shí)時(shí)推理的主要阻礙。為此，一些輕量級(jí) (lightweight) 目標(biāo)檢測(cè)網(wǎng)絡(luò)模型被提出，如 SSD-Lite, YOLO-Lite, YOLO-tiny 等，以實(shí)現(xiàn)移動(dòng)設(shè)備上的快速目標(biāo)檢測(cè)。但是這些輕量級(jí)網(wǎng)絡(luò)的解決方案效果依然不理想。

為了更好地滿足目標(biāo)檢測(cè)框架的落地需求，CoCoPIE 團(tuán)隊(duì)：美國(guó)東北大學(xué)王言治研究組和威廉瑪麗學(xué)院任彬研究組共同提出了名為 YOLObile 的手機(jī)端目標(biāo)檢測(cè)加速框架。YOLObile 框架通過(guò)「壓縮 - 編譯」協(xié)同設(shè)計(jì)在手機(jī)端實(shí)現(xiàn)了高準(zhǔn)確率實(shí)時(shí)目標(biāo)檢測(cè)。該框架使用一種新提出的名為「block-punched」的權(quán)重剪枝方案，對(duì)模型進(jìn)行有效的壓縮。在編譯器優(yōu)化技術(shù)的協(xié)助下，在手機(jī)端實(shí)現(xiàn)高準(zhǔn)確率的實(shí)時(shí)目標(biāo)檢測(cè)。該研究還提出了一種高效的手機(jī) GPU - 手機(jī) CPU 協(xié)同計(jì)算優(yōu)化方案，進(jìn)一步提高了計(jì)算資源的利用率和執(zhí)行速度。相比 YOLOv4 的原版，加速后的 YOLObile 的運(yùn)行速度提高了 4 倍，并且維持了 49mAP 的準(zhǔn)確率。相比 YOLOv3 完整版，該框架快 7 倍，在手機(jī)上實(shí)現(xiàn)了 19FPS 的實(shí)時(shí)高準(zhǔn)確率目標(biāo)檢測(cè)。同時(shí)準(zhǔn)確率高于 YOLOv3，并沒(méi)有用犧牲準(zhǔn)確率來(lái)提高計(jì)算速度。

• 開(kāi)源的模型與訓(xùn)練代碼：https://github.com/nightsnack/YOLObile

• 論文鏈接：https://arxiv.org/abs/2009.05697

• demo鏈接 · B 站（還有更多應(yīng)用的 demo）：https://www.bilibili.com/video/BV1bz4y1y7CR

研究方法

替換硬件支持性不好的操作符

在原版的 YOLOv4 中，有一些操作符不能夠最大化地利用硬件設(shè)備的執(zhí)行效率，比如帶有指數(shù)運(yùn)算的激活函數(shù)可能會(huì)造成運(yùn)行的延遲增加，成為降低延時(shí)提高效率的瓶頸。該研究把這些操作符相應(yīng)地替換成對(duì)硬件更加友好的版本，還有一些操作符是 ONNX 還未支持的（YOLObile 用 ONNX 作為模型的存儲(chǔ)方式），研究者把它替換成 ONNX 支持的運(yùn)算符。

比如，在 YOLOv4 引入的新模塊 Spatial Attention Module （SAM）中，用了 sigmoid 作為分支的激活函數(shù)，該研究在嘗試把它換成 hard-sigmoid 后發(fā)現(xiàn)：準(zhǔn)確率和直接刪除相比幾乎一致，增加的模塊又會(huì)增加計(jì)算量，所以研究者將其刪除了。

Mish 激活函數(shù)也涉及了指數(shù)運(yùn)算，同時(shí)在 pytorch 上的支持不太友好，會(huì)在訓(xùn)練時(shí)占用很多緩存，同時(shí)它在 pytorch 上也不能夠像 C++ 版本的 YOLOv4 一樣帶來(lái)很大的準(zhǔn)確率提升，而且 ONNX 尚未支持。研究者將其換成了 YOLOv3 的 leaky RELU。

「block-punched」權(quán)重剪枝方案

在 YOLObile 優(yōu)化框架中，作者使用了新提出的名為「block-punched」的權(quán)重剪枝 (weight pruning) 方案。這種剪枝方案旨在在獲得較高剪枝結(jié)構(gòu)自由度的同時(shí)，還能使剪枝后的模型結(jié)構(gòu)較好地利用硬件并行計(jì)算。這樣就從兩方面分別保證了剪枝后模型的準(zhǔn)確率和較高的運(yùn)算速度。

在這種剪枝方案中，每層的權(quán)重矩陣將被劃分為大小相等的多個(gè)塊(block)，因此，每個(gè)塊中將包含來(lái)自 m 個(gè) filter 的 n 個(gè) channel 的權(quán)重。在每個(gè)塊中，被剪枝的位置需要做出如下限定：需要修剪所有 filter 相同位置的一個(gè)或多個(gè)權(quán)重，同時(shí)也修剪所有通道相同位置的一個(gè)或多個(gè)權(quán)重。從另一個(gè)角度來(lái)看，這種剪枝方案將權(quán)重的剪枝位置貫穿了整個(gè)塊中所有的卷積核(kernel)。對(duì)于不同的 block，剪枝的位置和剪掉的 weight 數(shù)量都是靈活可變的。另外這種剪枝策略可以應(yīng)用在卷積核大小 3x3，1x1，5x5 等的卷積層上，也適用于 FC 層。

通過(guò)劃分為固定大小的塊來(lái)進(jìn)行剪枝，能夠提高編譯器的并行度，進(jìn)而提高在手機(jī)上運(yùn)行的速率。

1. 在準(zhǔn)確率方面，通過(guò)劃分多個(gè)小區(qū)塊，這種剪枝方法實(shí)現(xiàn)了更加細(xì)粒度的剪枝。相較于傳統(tǒng)的結(jié)構(gòu)化剪枝（剪除整個(gè) filter 或 channel），這種方式具有更高的剪枝結(jié)構(gòu)自由度，從而更好地保持了模型的準(zhǔn)確率。

2. 在硬件表現(xiàn)方面，因?yàn)樵谕恍^(qū)塊中，所有 filter 修剪被修剪的位置相同，所以在并行計(jì)算時(shí)，所有 filter 將統(tǒng)一跳過(guò)讀取相同的輸入數(shù)據(jù)，從而減輕處理這些 filter 的線程之間的內(nèi)存壓力。而限制修剪小區(qū)塊內(nèi)各 channel 的相同位置，確保了所有 channel 共享相同的計(jì)算模式，從而消除處理每個(gè) channel 的線程之間的計(jì)算差異。因此，這種剪枝方案可以大幅度降低在計(jì)算過(guò)程中處理稀疏結(jié)構(gòu)的額外開(kāi)銷(xiāo)，從而達(dá)到更好的加速效果。

Reweighted 算法

作者采用了 Reweighted Regularization 算法來(lái)篩選出剪枝掉的 weight 的位置，reweight 算法依據(jù) weight 的絕對(duì)值大小篩選出相對(duì)重要的位置。整個(gè)剪枝的過(guò)程從依據(jù)預(yù)先訓(xùn)練好的網(wǎng)絡(luò)模型開(kāi)始，然后將 Reweighted 算法融合在訓(xùn)練的過(guò)程中，在訓(xùn)練每 3 到 4 個(gè) epochs 后就能夠得到一組可選的剪枝位置，通過(guò)重復(fù)訓(xùn)練來(lái)調(diào)整剪枝位置，最終得到精確的剪枝模型結(jié)果。

編譯器的優(yōu)化

為了支持「block-punched」剪枝方案，編譯器也需要作出相應(yīng)的調(diào)整和優(yōu)化。為了更好地存儲(chǔ)和調(diào)用經(jīng)過(guò)「block-punched」剪枝后的 weight index，研究者引入了新的存儲(chǔ)方案，將 index 的存儲(chǔ)空間進(jìn)一步壓縮。同時(shí)，通過(guò)對(duì)所有的塊進(jìn)行重新排序，編譯器能夠在更少的內(nèi)存訪問(wèn)次數(shù)下運(yùn)行，進(jìn)而獲得更快的運(yùn)算速度。

手機(jī) GPU 和手機(jī) CPU 協(xié)同計(jì)算的優(yōu)化方案

YOLObile 中還使用了手機(jī) GPU 和手機(jī) CPU 協(xié)同計(jì)算的方式來(lái)進(jìn)一步降低整個(gè)網(wǎng)絡(luò)的運(yùn)算時(shí)間?，F(xiàn)在主流的移動(dòng)端 DNN 推理加速框架，如 TensorFlow-Lite，MNN 和 TVM 都只能支持手機(jī) CPU 或 GPU 單獨(dú)運(yùn)算，因此會(huì)導(dǎo)致潛在的計(jì)算資源浪費(fèi)。

YOLObile 提出針對(duì)網(wǎng)絡(luò)中的分支結(jié)構(gòu)，比如 YOLOv4 中大量使用的 Cross Stage Partial (CSP)結(jié)構(gòu)，使用 CPU 來(lái)輔助 GPU 同時(shí)進(jìn)行一些相互無(wú)依賴關(guān)系的分支運(yùn)算，從而更好地利用計(jì)算資源，減少網(wǎng)絡(luò)的運(yùn)算時(shí)間。YOLObile 框架將待優(yōu)化的網(wǎng)絡(luò)分支分為有卷積運(yùn)算分支和無(wú)卷積運(yùn)算分支，并對(duì)這兩種情況分別給出了優(yōu)化方案。

如下圖所示，CSP 是一個(gè)跨越很多殘差 block 的長(zhǎng)分支，在手機(jī)的執(zhí)行過(guò)程中，單一的 GPU 計(jì)算往往是順序地執(zhí)行完上面堆疊殘差 block 的分支后再執(zhí)行下面的 CSP 分支，然后拼接在一起作為下一層的輸入。研究者將卷積層數(shù)更少的 branch2 挪到 CPU 上去，CPU 執(zhí)行時(shí)間少于上面 branch1 在 GPU 上的總運(yùn)算時(shí)間，這個(gè)并行操作能夠有效減少運(yùn)算延遲。當(dāng)然，決定能否將 branch2 轉(zhuǎn)移到 CPU 的因素在于 branch1 的卷積層數(shù)多少，通常 CSP block 會(huì)跨越 8 個(gè)殘差 block，也有的時(shí)候出現(xiàn)只跨越 4 個(gè)殘差 block 的情況，還有在前幾層只跨一個(gè)殘差 block 的情況。對(duì)于只跨 1 個(gè)殘差 block 的情況明顯還是 GPU 順序執(zhí)行更高效，對(duì)于跨越多個(gè)的就需要用實(shí)際測(cè)出的延遲來(lái)做判斷。值得注意的是，轉(zhuǎn)移數(shù)據(jù)到不同處理設(shè)備的時(shí)候，需要加入數(shù)據(jù)傳輸拷貝的時(shí)間。

在 YOLOv4 最后輸出的位置，3 個(gè) YOLO head 輸出部分有很多諸如轉(zhuǎn)置，變形之類(lèi)的非卷積運(yùn)算，這些非卷積運(yùn)算在 CPU 和 GPU 上的運(yùn)行效率相當(dāng)，作者同樣基于運(yùn)行時(shí)間，考慮將部分運(yùn)算符轉(zhuǎn)移到 CPU 上去做，選擇最高效的執(zhí)行方式。

實(shí)驗(yàn)結(jié)果

整個(gè)訓(xùn)練過(guò)程采用 4 塊 RTX2080Ti，訓(xùn)練時(shí)間為 5 天。使用配備高通驍龍 855 CPU 和 Adreno 650 GPU 的三星 Galaxy S20 作為測(cè)試平臺(tái)。訓(xùn)練的數(shù)據(jù)集是 MS COCO，和 YOLOv4 原版保持一致。實(shí)驗(yàn)結(jié)果表明，當(dāng)使用 YOLOv4 為基礎(chǔ)模型進(jìn)行優(yōu)化時(shí)，該研究的優(yōu)化框架可以成功將原模型大小壓縮至 1/14，在未使用 GPU-CPU 協(xié)同計(jì)算優(yōu)化時(shí)，將每秒檢測(cè)幀數(shù)（FPS）提升至 17，且達(dá)到 49（mAP）的準(zhǔn)確率。

從下圖中可以看到，與眾多具有代表性的目標(biāo)檢測(cè)網(wǎng)絡(luò)相比，該研究的優(yōu)化模型在準(zhǔn)確率與速度兩方面同時(shí)具有優(yōu)異的表現(xiàn)，而不再是簡(jiǎn)單的犧牲大幅準(zhǔn)確率來(lái)獲取一定程度的速度提升。

下表展示了 YOLObile 與其他具有代表性的目標(biāo)檢測(cè)網(wǎng)絡(luò)在準(zhǔn)確率與速度方面的具體比較。值得注意的是，作者的 GPU-CPU 協(xié)同計(jì)算優(yōu)化方案可以進(jìn)一步將執(zhí)行速度提高至 19FPS。

相比于其他 one-stage 目標(biāo)檢測(cè)模型和 light-weight 模型，經(jīng)過(guò)剪枝后的 YOLObile 更快且更精確，在準(zhǔn)確率和速度上實(shí)現(xiàn)了帕累托最優(yōu)。

消融實(shí)驗(yàn)

不同剪枝策略比較

表格中列出了在相同壓縮倍率的情況下，Structured Pruning, Unstructured Pruning, 和文中提出的 Block-Punched Pruning 這三種剪枝策略的準(zhǔn)確率和運(yùn)行效率的比較。可以看到，Unstructured Pruning 在壓縮 8 倍的狀態(tài)下能夠維持較高的準(zhǔn)確率，但是和前文描述的一樣，它的運(yùn)行效率較低，即使在壓縮 8 倍的情況下，運(yùn)行時(shí)間只提高了不到 2 倍。Structured Pruning 后模型能夠高效率地在硬件設(shè)備上執(zhí)行，但是準(zhǔn)確率相比 Unstructured Pruning 有大幅度的下降。而根據(jù)編譯器設(shè)計(jì)的 Block-Punched 剪枝策略，能夠維持較高的準(zhǔn)確率，并且達(dá)到和 Structured Pruning 相當(dāng)?shù)膱?zhí)行效率。

和 Pattern-Based Pruning 的比較

Pattern-Based Pruning 只能用在 3x3 的層，為了達(dá)到較高的壓縮倍率，就需要對(duì) 3x3 的卷積層壓縮掉更多的 weight。下圖是一個(gè)關(guān)于不同倍率下 pattern-based pruning 和 block-punched pruning 的比較，可以看出 pattern-based pruning 在較低倍率壓縮的情況下，能夠擁有較高的準(zhǔn)確率和執(zhí)行效率，但是受限制于只適用于 3x3 卷積層，即使將所有的 3x3 層全部壓縮也只能達(dá)到 5 倍多的壓縮率，但是準(zhǔn)確率就非常低了。block-based pruning 由于可以將所有的層都進(jìn)行壓縮，在高倍率的情況下依然能夠維持較高的準(zhǔn)確率。

Block-Punched 剪枝中分塊大小的比較

該研究對(duì)四種不同的塊大小進(jìn)行實(shí)驗(yàn)，為了實(shí)現(xiàn)較高的硬件并行度，每個(gè)模塊中的 channel 數(shù)固定為 4，這與 GPU 和 CPU 矢量寄存器的長(zhǎng)度相同。在每個(gè)塊中使用不同數(shù)量的 filter 來(lái)評(píng)估準(zhǔn)確性和速度。如圖 5.3 所示，與較小的塊相比，較大的塊可以更好地利用硬件并行性，并可以實(shí)現(xiàn)更高的推理速度。但是其粗略的修剪粒度導(dǎo)致準(zhǔn)確性下降。較小的塊可以獲得較高的精度，但會(huì)犧牲推理速度。根據(jù)結(jié)果，研究者將 8×4（8 個(gè)連續(xù) filter 的 4 個(gè)連續(xù) channel）視為移動(dòng)設(shè)備上的所需塊大小，這在精度和速度之間取得了良好的平衡。

各層的剪枝比例

YOLOv4 在卷積層中僅包含 3×3 和 1×1 的 kernel。該研究提出這兩種類(lèi)型的卷積層在修剪過(guò)程中具有不同的敏感度并進(jìn)行了兩組實(shí)驗(yàn)。在相同數(shù)量的 FLOPs 下，將一組中的所有層平均修剪，另一組中，3×3 卷積層的壓縮率比 1×1 卷積層高 1.15 倍。如上表所示，均勻修剪的模型和不均勻修剪的模型相比，準(zhǔn)確性和速度都較低。

END

雙一流大學(xué)研究生團(tuán)隊(duì)創(chuàng)建，專(zhuān)注于目標(biāo)檢測(cè)與深度學(xué)習(xí)，希望可以將分享變成一種習(xí)慣！

整理不易，點(diǎn)贊三連↓