MUCNetV2：內(nèi)存瓶頸和計(jì)算負(fù)載問(wèn)題一舉突破？分類(lèi)&檢測(cè)都有較高性能

論文地址：https://arxiv.org/pdf/2110.15352.pdf

源代碼：https://mcunet.mit.edu

MCUNetV2取得了MCU端新的ImageNet分類(lèi)記錄71.8；更重要的是解鎖了MCU端執(zhí)行稠密預(yù)測(cè)任務(wù)的可能性

一、背景

基于微控制器單元 (MCU) 等微型硬件的物聯(lián)網(wǎng)設(shè)備如今無(wú)處不在。在如此微小的硬件上部署深度學(xué)習(xí)模型將使我們能夠?qū)崿F(xiàn)人工智能的民主化。然而，由于內(nèi)存預(yù)算緊張，微型深度學(xué)習(xí)與移動(dòng)深度學(xué)習(xí)有著根本的不同：普通MCU通常具有小于512kB的SRAM，這對(duì)于部署大多數(shù)現(xiàn)成的深度學(xué)習(xí)網(wǎng)絡(luò)來(lái)說(shuō)太小了。即使對(duì)于Raspberry Pi 4等更強(qiáng)大的硬件，將推理放入L2緩存(1MB)也可以顯著提高能源效率。這些對(duì)具有較小峰值內(nèi)存使用量的高效AI推理提出了新的挑戰(zhàn)。

現(xiàn)有工作采用修剪、量化和神經(jīng)架構(gòu)搜索來(lái)進(jìn)行高效的深度學(xué)習(xí)。但是，這些方法側(cè)重于減少參數(shù)和FLOP的數(shù)量，而不是內(nèi)存瓶頸。緊張的內(nèi)存預(yù)算限制了特征圖/激活大小，限制了我們使用較小的模型容量或較小的輸入圖像大小。實(shí)際上，現(xiàn)有的tinyML工作中使用的輸入分辨率通常很小(<224^2)，這對(duì)于圖像分類(lèi)（例如 ImageNet、VWW）來(lái)說(shuō)可能是可以接受的，但對(duì)于像目標(biāo)檢測(cè)這樣的密集預(yù)測(cè)任務(wù)來(lái)說(shuō)是不可以接受的：如下圖圖所示，性能與圖像分類(lèi)相比，輸入分辨率下物體檢測(cè)的退化速度要快得多。這種限制阻礙了微型深度學(xué)習(xí)在許多現(xiàn)實(shí)生活任務(wù)（例如，行人檢測(cè)中的應(yīng)用）。

研究者對(duì)高效網(wǎng)絡(luò)設(shè)計(jì)中每一層的內(nèi)存使用情況進(jìn)行了深入分析，發(fā)現(xiàn)它們的激活內(nèi)存分布非常不平衡。以MobileNetV2為例，如下圖所示，只有前5個(gè)塊的內(nèi)存峰值較高（>450kB），成為整個(gè)網(wǎng)絡(luò)的內(nèi)存瓶頸。其余13個(gè)塊的內(nèi)存使用率較低，可以輕松容納256kB的MCU。初始內(nèi)存密集階段的峰值內(nèi)存比網(wǎng)絡(luò)其他部分高8倍。這樣的內(nèi)存模式留下了巨大的優(yōu)化空間：如果我們能找到一種方法來(lái)“繞過(guò)”內(nèi)存密集階段，我們可以將整體峰值內(nèi)存減少8倍。

二、前言

由于有限的內(nèi)存，MCU(MicroController Units, MCU)端的TinyDL極具挑戰(zhàn)性。為緩解該問(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)用鋪平了道路 。

三、新框架

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)化，但是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è)備所不期望的。

Reducing Computation Overhead by 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ù)量。

對(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%，模型性能保持同等水平。具體比較見(jiàn)下表：

Joint Neural Architecture and Inference Scheduling Search：

重新分配感受野使我們能夠以最小的計(jì)算/延遲開(kāi)銷(xiāo)享受內(nèi)存減少的好處，但策略因不同的主干而有所不同。在設(shè)計(jì)主干網(wǎng)絡(luò)（例如，使用更大的輸入分辨率）時(shí)，減少的峰值內(nèi)存還允許更大的自由度。為了探索如此大的設(shè)計(jì)空間，研究者建議以自動(dòng)化方式聯(lián)合優(yōu)化神經(jīng)架構(gòu)和推理調(diào)度。給定某個(gè)數(shù)據(jù)集和硬件限制（SRAM 限制、Flash 限制、延遲限制等），我們的目標(biāo)是在滿(mǎn)足所有限制條件的同時(shí)實(shí)現(xiàn)最高準(zhǔn)確度。我們有以下要優(yō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ī)制?？蚣艿耐评硪婊贛CUnet中的TinyEngine擴(kuò)展而來(lái)，除了TinyEngine中已有的優(yōu)化外，還需要確定塊數(shù)量p與模塊數(shù)量n以執(zhí)行patch推理，確保滿(mǎn)足SRAM約束。

Joint Search：需要協(xié)同設(shè)計(jì)骨干優(yōu)化與推理機(jī)制。比如，給定相同約束，可以選更小的模型以layer方式推理，或更大的模塊以patch方式推理。因此，我們兩者納入優(yōu)化并采用進(jìn)化搜索尋找滿(mǎn)足約束的最佳組合。

四、實(shí)驗(yàn)

上圖1給出了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ì)很有必要 。

上圖2給出了on-device profiling的結(jié)果，它不僅與架構(gòu)相關(guān)還與硬件相關(guān)?？梢钥吹剑簆atch推理可以降低SRAM峰值內(nèi)存達(dá)4-6x 。某些模型可能會(huì)因初始階段的低硬件利用導(dǎo)致更長(zhǎng)的推理延遲，但合理的架構(gòu)設(shè)計(jì)可以將額外推理延遲降低到4%，相比內(nèi)存降低完全可以忽略。

上表給出了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%超出了此前同等量化策略方案。

在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ò)展搜索空間使得我們可以在更大的輸入分辨率與更小的模型中作出更佳的選擇。

在不同尺度下，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)。

Memory Distributions of Efficient Models