內(nèi)存瓶頸和計(jì)算負(fù)載問(wèn)題一舉突破?韓松團(tuán)隊(duì)提出MUCNetV2:解鎖MCU端新紀(jì)錄!
以下文章來(lái)源于極市平臺(tái) ,作者CV開(kāi)發(fā)者都愛(ài)看的
作者丨h(huán)appy
編輯丨極市平臺(tái)
導(dǎo)讀
本文是MIT韓松團(tuán)隊(duì)繼MCUNet之后關(guān)于MCU 端TinyDL的又一次突破,已被NeurIPS2021接收。所提MCUNetV2取得了MCU端新的ImageNet分類(lèi)記錄71.8;更重要的是解鎖了MCU端執(zhí)行稠密預(yù)測(cè)任務(wù)的可能性,如目標(biāo)檢測(cè)取得了比已有方案高16.9%mAP@VOC的指標(biāo)。
論文地址:https://arxiv.org/abs/2110.15352
Slides地址:https://hanlab.mit.edu/projects/tinyml/mcunet/assets/MCUNetV2/slides.pdf
項(xiàng)目地址:https://hanlab.mit.edu/projects/tinyml/mcunet/
本文是MIT韓松團(tuán)隊(duì)繼MCUNet之后關(guān)于MCU端TinyDL的又一次突破,已被NeurIPS2021接收。針對(duì)CNN在MCU端存在的不平衡內(nèi)存分布問(wèn)題,提出一種patch-by-patch 推理機(jī)制;與此同時(shí),還提出recptive field redistribution 調(diào)整感受野與FLOPs以降低整體計(jì)算負(fù)載。所提MCUNetV2取得了MCU端新的ImageNet分類(lèi)記錄71.8;更重要的是解鎖了MCU端執(zhí)行稠密預(yù)測(cè)任務(wù)的可能性,如目標(biāo)檢測(cè)取得了比已有方案高16.9%mAP@VOC的指標(biāo)。本研究極大程度上解決了TinyDL的內(nèi)存瓶頸問(wèn)題,為圖像分類(lèi)之外的其他視覺(jué)應(yīng)用鋪平了道路 。
Abstract
由于有限的內(nèi)存,MCU(MicroController Units, MCU)端的TinyDL極具挑戰(zhàn)性。我們發(fā)現(xiàn):內(nèi)存瓶頸源于CNN設(shè)計(jì)導(dǎo)致的內(nèi)存分布不平衡 ,即網(wǎng)絡(luò)的前幾個(gè)模塊具有比其他模塊大的多的內(nèi)存占用。
為緩解該問(wèn)題,我們提出一種廣義的patch-by-patch 推理機(jī)制,它僅對(duì)特征圖的局部區(qū)域進(jìn)行處理,大幅降低了峰值內(nèi)存。然而,常規(guī)的實(shí)現(xiàn)方式會(huì)帶來(lái)重疊塊與計(jì)算復(fù)雜問(wèn)題。我們進(jìn)一步提出了recptive field redistribution 調(diào)整感受野與FLOPs以降低整體計(jì)算負(fù)載。人工方式重分布感受野無(wú)疑非常困難!我們采用NAS對(duì)網(wǎng)絡(luò)架構(gòu)與推理機(jī)制進(jìn)行聯(lián)合優(yōu)化得到了本文的MCUNetV2。所提推理機(jī)制能大幅降低峰值內(nèi)存達(dá)4-8倍。
所推MCUNetV2取得了MCU端新的ImageNet分類(lèi)記錄71.8% ;更重要的是,MCUNetV2解鎖了MCU端執(zhí)行稠密預(yù)測(cè)任務(wù)的可能性,如目標(biāo)檢測(cè)取得了比已有方案高16.9%mAP@VOC的指標(biāo)。本研究極大程度上解決了TinyDL的內(nèi)存瓶頸問(wèn)題,為圖像分類(lèi)之外的其他視覺(jué)應(yīng)用鋪平了道路 。本文貢獻(xiàn)可以用下圖描述。
Memory Bottleneck
Imbalanced memory distribution 以MobileNetV2為例,上圖提供了每個(gè)模塊的峰值峰值內(nèi)存占用(注:上述信息通過(guò)int8度量得到)。我們可看到非常清晰的內(nèi)存占用分布不平衡 :前5個(gè)模塊具有非常大的峰值內(nèi)存,超出了MCUs內(nèi)存約束,而其他13個(gè)模塊則能輕易滿足256kB內(nèi)存約束 。我們同樣檢查了其他高效網(wǎng)絡(luò)架構(gòu)(包含針對(duì)MCU而設(shè)計(jì)的MCUNet)并發(fā)現(xiàn)了類(lèi)似的現(xiàn)象。
我們發(fā)現(xiàn):該現(xiàn)象適用于大部分單分支與殘差CNN架構(gòu)。每個(gè)階段,圖像分類(lèi)率降低一半,通道數(shù)僅乘2,故特征尺寸逐漸減小。因此,內(nèi)存瓶頸傾向于出現(xiàn)在網(wǎng)絡(luò)的早期階段 。
Challenges and Opportunities 這種內(nèi)存分布不平衡問(wèn)題極大限制了MCU端的模型容量與輸入分辨率。為適應(yīng)初始階段內(nèi)存密集問(wèn)題,整個(gè)網(wǎng)絡(luò)都需要進(jìn)行縮小,哪怕大部分網(wǎng)絡(luò)已經(jīng)具有非常小的內(nèi)存占用。這也就使得分辨率敏感的任務(wù)(如目標(biāo)檢測(cè))難以落地MCU端。以MobileNetV2的第一個(gè)卷積為例,輸入通道為3,輸出通道維32,stride=2,當(dāng)輸入圖像分辨率為時(shí)需要占用的內(nèi)存為(注:int8量化),這是當(dāng)前MCU所無(wú)法滿足的。另一方面,如果我們能找到一種方式跳過(guò) 內(nèi)存密集階段,我們就可以大幅降低整個(gè)網(wǎng)絡(luò)的峰值內(nèi)存,進(jìn)而有了更大的優(yōu)化空間。
MCUNetV2: Memory-Efficient Patch-based Inference
Patch-based Inference
現(xiàn)有的推理框架(如TFLite Micro, TinyEngine, MicroTVM)采用layer-by-layer方式運(yùn)行。對(duì)于每個(gè)卷積層,推理庫(kù)首先在SRAM中開(kāi)辟輸入與輸出buffer,完成計(jì)算后釋放輸入buffer。這種處理機(jī)制更易于進(jìn)行推理優(yōu)化(比如im2col, tiling),但是SRAM必須保留完整的輸入與輸出buffer。
我們提出一種patch-based inference 機(jī)制打破初始層的內(nèi)存瓶頸問(wèn)題,見(jiàn)下圖。本文所提patch-based推理則在內(nèi)存密集階段以patch-by-patch方式運(yùn)行 (見(jiàn)下圖右)。模型每次僅在小比例區(qū)域(比完整區(qū)域小10倍)進(jìn)行計(jì)算,這可以有效降低峰值內(nèi)存占用。當(dāng)完成該部分后,網(wǎng)絡(luò)的其他具有較小峰值內(nèi)存的模塊則采用常規(guī)的layer-by-layer 方式運(yùn)行。
以上圖stride=1/2的圖示為例,對(duì)于逐層計(jì)算方式,第一層具有大輸入輸出尺寸,導(dǎo)致非常高內(nèi)存占用。通過(guò)空域拆分計(jì)算,我們以patch-by-patch方式開(kāi)辟buffer并進(jìn)行計(jì)算,此時(shí)我們僅需保存一個(gè)塊的buffer而非完整特征圖。
Computation overhead 內(nèi)存節(jié)省的代價(jià)來(lái)自計(jì)算負(fù)載提升。為與逐層推理相同的輸出結(jié)果,非重疊輸出塊需要對(duì)應(yīng)了重疊輸入塊(見(jiàn)上圖b陰影區(qū)域)。這種重復(fù)性的計(jì)算會(huì)提升網(wǎng)絡(luò)整體10-17%計(jì)算量,這是低功耗邊緣設(shè)備所不期望的。
Redistributing the Receptive Field
計(jì)算復(fù)雜度與patch方案初始階段的感受野相關(guān),考慮到patch階段的輸出,越大的感受野對(duì)應(yīng)越大的輸入分辨率,進(jìn)而導(dǎo)致更多的重疊區(qū)域與重復(fù)計(jì)算。對(duì)于MobileNetV2來(lái)說(shuō),如果我們僅考慮下采樣,每個(gè)輸入塊的邊長(zhǎng)為。當(dāng)提升感受野時(shí),每個(gè)輸入塊需要采用尺寸,導(dǎo)致更大的重疊區(qū)域。
我們提出了重分布(redistribute)感受野以降低計(jì)算復(fù)雜度,其基本思想在于:
降低patch階段的感受野;
提升layer階段的感受野。
降低初始階段的感受野有助于降低patch部分的輸入尺寸與重復(fù)計(jì)算。然而,某些任務(wù)會(huì)因感受野較小導(dǎo)致性能下降。因此,我們進(jìn)一步提升layer部分的感受野以補(bǔ)償性能損失。
上圖給出了以MobileNetV2為例調(diào)整前后的架構(gòu)對(duì)比圖。在patch推理階段,我們采用更小的核與更少的模塊,而在layer推理階段提升模塊的數(shù)量。
上表給出了兩種架構(gòu)的性能對(duì)比,可以看到:
對(duì)于兩種架構(gòu)來(lái)說(shuō),patch推理可以降低SRAM峰值達(dá)8倍;
原始MobileNetV2設(shè)計(jì)會(huì)導(dǎo)致patch部分42%的計(jì)算量增加,整體增加10%;而調(diào)整的架構(gòu)輸入尺寸從75下降到了63,整體計(jì)算量?jī)H增加3%,模型性能保持同等水平。
Joint Optimization
感受野重分布使得我們可以以績(jī)效的計(jì)算提升享受大幅內(nèi)存降低,但是需要對(duì)不同架構(gòu)采用逐個(gè)手工調(diào)整。峰值內(nèi)存降低也為架構(gòu)設(shè)計(jì)帶來(lái)了更大的自由度。
為探索更大的設(shè)計(jì)空間,我們采用neural architecture與inference scheduling自動(dòng)聯(lián)合優(yōu)化方案(可參考上圖),旨在獲得滿足所有約束且具有最高精度的架構(gòu)。
Backbone optimization 我們參考MCUNet采用了類(lèi)MnasNet搜索空間,而類(lèi)MobileNetV3搜索空間因Swish激活函數(shù)問(wèn)題導(dǎo)致難以量化而棄用。作者認(rèn)為:最佳的搜索空間配置不僅硬件相關(guān),同樣任務(wù)相關(guān)。因此,我們還將r與w納入搜索空間。
Inference scheduling optimization 給定模型與硬件約束,我們將尋找最佳推理機(jī)制。我們的推理引擎基于MCUnet中的TinyEngine擴(kuò)展而來(lái),除了TinyEngine中已有的優(yōu)化外,我們還需要確定塊數(shù)量p與模塊數(shù)量n以執(zhí)行patch推理,確保滿足SRAM約束。
Joint Search 我們需要協(xié)同設(shè)計(jì)骨干優(yōu)化與推理機(jī)制。比如,給定相同約束,我們可以選更小的模型以layer方式推理,或更大的模塊以patch方式推理。因此,我們兩者納入優(yōu)化并采用進(jìn)化搜索尋找滿足約束的最佳組合。
Experiments
Memory Profiling 內(nèi)存占用是獨(dú)立于推理框架的,為更好的比例,我們研究了兩種分析設(shè)置。
上圖給出了Analytic profiling的結(jié)果,它僅與架構(gòu)有關(guān)而與推理框架無(wú)關(guān)。可以看到:
patch推理可以大幅降低峰值內(nèi)存達(dá)3.7-8.0x,僅增加8-17%計(jì)算量;
對(duì)于MobileNetV2,通過(guò)重分布感受野可以將額外的計(jì)算量從10%降到3%;
內(nèi)存節(jié)省與計(jì)算降低與網(wǎng)絡(luò)架構(gòu)相關(guān),故網(wǎng)絡(luò)架構(gòu)與推理引擎的協(xié)同設(shè)計(jì)很有必要 。
上圖給出了on-device profiling的結(jié)果,它不僅與架構(gòu)相關(guān)還與硬件相關(guān)。可以看到:patch推理可以降低SRAM峰值內(nèi)存達(dá)4-6x 。某些模型可能會(huì)因初始階段的低硬件利用導(dǎo)致更長(zhǎng)的推理延遲,但合理的架構(gòu)設(shè)計(jì)可以將額外推理延遲降低到4%,相比內(nèi)存降低完全可以忽略。
Tiny Image Classification
上表給出了ImageNet上的分類(lèi)性能,從中可以看到:
在256kB SRAM/1MB Flash硬件下,MCUNetV2具有比MCUnet高4.6%top1精度、低18%的峰值SRAM的性能。
在512kB SRAM/2MB Flash硬件下,MCUNet取得了新的ImageNet記錄71.8%(限商用MCU),以3.3%超出了此前同等量化策略方案。
Tiny Object Detection
上表為Pascal VOC數(shù)據(jù)機(jī)上的檢測(cè)性能對(duì)比,從中可以看到:
在H7 MCU上,相比MCUNet,MCUNetV2-H7的性能提升達(dá)16.7%mAP;
在更廉價(jià)的M4MCU上,MCUNetV2-M4仍可提升13.2%mAP指標(biāo),峰值SRAM降低1.9x。
MUCNetV2-M4與MCUNet具有相似的計(jì)算量,但更高的性能。這是因?yàn)椋簆atch推理導(dǎo)致的擴(kuò)展搜索空間使得我們可以在更大的輸入分辨率與更小的模型中作出更佳的選擇。
上表給出了WIDER FACE數(shù)據(jù)集上的性能對(duì)比,從中可以看到:
在不同尺度下,MCUNetV2均優(yōu)于現(xiàn)有方案;
MCUNetV2-L取得了同等檢測(cè)性能,但峰值內(nèi)存被RNNPool-Face-C小3.4x,比LFFD小9.9x;計(jì)算量分別小1.6x與8.4x;
MCUNetV2-S同樣具有比RNNPool-Face-A與EagleEye小1.8x的峰值內(nèi)存。
下圖還給出了檢測(cè)效果的對(duì)比,MCUNetV2的魯棒性更強(qiáng)。
本文僅做學(xué)術(shù)分享,如有侵權(quán),請(qǐng)聯(lián)系刪文。
*博客內(nèi)容為網(wǎng)友個(gè)人發(fā)布,僅代表博主個(gè)人觀點(diǎn),如有侵權(quán)請(qǐng)聯(lián)系工作人員刪除。