内存瓶颈和计算负载问题一举突破？韩松团队提出MUCNetV2：解锁MCU端新纪录！

以下文章来源于极市平台 ，作者CV开发者都爱看的

作者丨happy

编辑丨极市平台

导读

本文是MIT韩松团队继MCUNet之后关于MCU 端TinyDL的又一次突破，已被NeurIPS2021接收。所提MCUNetV2取得了MCU端新的ImageNet分类记录71.8；更重要的是解锁了MCU端执行稠密预测任务的可能性，如目标检测取得了比已有方案高16.9%mAP@VOC的指标。 

论文地址：https://arxiv.org/abs/2110.15352

Slides地址：https://hanlab.mit.edu/projects/tinyml/mcunet/assets/MCUNetV2/slides.pdf

项目地址：https://hanlab.mit.edu/projects/tinyml/mcunet/

本文是MIT韩松团队继MCUNet之后关于MCU端TinyDL的又一次突破，已被NeurIPS2021接收。针对CNN在MCU端存在的不平衡内存分布问题，提出一种patch-by-patch 推理机制；与此同时，还提出recptive field redistribution 调整感受野与FLOPs以降低整体计算负载。所提MCUNetV2取得了MCU端新的ImageNet分类记录71.8；更重要的是解锁了MCU端执行稠密预测任务的可能性，如目标检测取得了比已有方案高16.9%mAP@VOC的指标。本研究极大程度上解决了TinyDL的内存瓶颈问题，为图像分类之外的其他视觉应用铺平了道路 。

Abstract

由于有限的内存，MCU(MicroController Units, MCU)端的TinyDL极具挑战性。我们发现：内存瓶颈源于CNN设计导致的内存分布不平衡 ，即网络的前几个模块具有比其他模块大的多的内存占用。

为缓解该问题，我们提出一种广义的patch-by-patch 推理机制，它仅对特征图的局部区域进行处理，大幅降低了峰值内存。然而，常规的实现方式会带来重叠块与计算复杂问题。我们进一步提出了recptive field redistribution 调整感受野与FLOPs以降低整体计算负载。人工方式重分布感受野无疑非常困难！我们采用NAS对网络架构与推理机制进行联合优化得到了本文的MCUNetV2。所提推理机制能大幅降低峰值内存达4-8倍。

所推MCUNetV2取得了MCU端新的ImageNet分类记录71.8% ；更重要的是，MCUNetV2解锁了MCU端执行稠密预测任务的可能性，如目标检测取得了比已有方案高16.9%mAP@VOC的指标。本研究极大程度上解决了TinyDL的内存瓶颈问题，为图像分类之外的其他视觉应用铺平了道路 。本文贡献可以用下图描述。

Memory Bottleneck

Imbalanced memory distribution  以MobileNetV2为例，上图提供了每个模块的峰值峰值内存占用(注：上述信息通过int8度量得到)。我们可看到非常清晰的内存占用分布不平衡 ：前5个模块具有非常大的峰值内存，超出了MCUs内存约束，而其他13个模块则能轻易满足256kB内存约束 。我们同样检查了其他高效网络架构(包含针对MCU而设计的MCUNet)并发现了类似的现象。

我们发现：该现象适用于大部分单分支与残差CNN架构。每个阶段，图像分类率降低一半，通道数仅乘2，故特征尺寸逐渐减小。因此，内存瓶颈倾向于出现在网络的早期阶段 。

Challenges and Opportunities  这种内存分布不平衡问题极大限制了MCU端的模型容量与输入分辨率。为适应初始阶段内存密集问题，整个网络都需要进行缩小，哪怕大部分网络已经具有非常小的内存占用。这也就使得分辨率敏感的任务(如目标检测)难以落地MCU端。以MobileNetV2的第一个卷积为例，输入通道为3，输出通道维32，stride=2，当输入图像分辨率为时需要占用的内存为(注：int8量化)，这是当前MCU所无法满足的。另一方面，如果我们能找到一种方式跳过 内存密集阶段，我们就可以大幅降低整个网络的峰值内存，进而有了更大的优化空间。

MCUNetV2: Memory-Efficient Patch-based Inference

Patch-based Inference

现有的推理框架(如TFLite Micro, TinyEngine, MicroTVM)采用layer-by-layer方式运行。对于每个卷积层，推理库首先在SRAM中开辟输入与输出buffer，完成计算后释放输入buffer。这种处理机制更易于进行推理优化(比如im2col, tiling)，但是SRAM必须保留完整的输入与输出buffer。

我们提出一种patch-based inference 机制打破初始层的内存瓶颈问题，见下图。本文所提patch-based推理则在内存密集阶段以patch-by-patch方式运行 (见下图右)。模型每次仅在小比例区域(比完整区域小10倍)进行计算，这可以有效降低峰值内存占用。当完成该部分后，网络的其他具有较小峰值内存的模块则采用常规的layer-by-layer 方式运行。

以上图stride=1/2的图示为例，对于逐层计算方式，第一层具有大输入输出尺寸，导致非常高内存占用。通过空域拆分计算，我们以patch-by-patch方式开辟buffer并进行计算，此时我们仅需保存一个块的buffer而非完整特征图。

Computation overhead  内存节省的代价来自计算负载提升。为与逐层推理相同的输出结果，非重叠输出块需要对应了重叠输入块(见上图b阴影区域)。这种重复性的计算会提升网络整体10-17%计算量，这是低功耗边缘设备所不期望的。

Redistributing the Receptive Field

计算复杂度与patch方案初始阶段的感受野相关，考虑到patch阶段的输出，越大的感受野对应越大的输入分辨率，进而导致更多的重叠区域与重复计算。对于MobileNetV2来说，如果我们仅考虑下采样，每个输入块的边长为。当提升感受野时，每个输入块需要采用尺寸，导致更大的重叠区域。

我们提出了重分布(redistribute)感受野以降低计算复杂度，其基本思想在于：

降低patch阶段的感受野；

提升layer阶段的感受野。

降低初始阶段的感受野有助于降低patch部分的输入尺寸与重复计算。然而，某些任务会因感受野较小导致性能下降。因此，我们进一步提升layer部分的感受野以补偿性能损失。

上图给出了以MobileNetV2为例调整前后的架构对比图。在patch推理阶段，我们采用更小的核与更少的模块，而在layer推理阶段提升模块的数量。

上表给出了两种架构的性能对比，可以看到：

对于两种架构来说，patch推理可以降低SRAM峰值达8倍；

原始MobileNetV2设计会导致patch部分42%的计算量增加，整体增加10%；而调整的架构输入尺寸从75下降到了63，整体计算量仅增加3%，模型性能保持同等水平。

Joint Optimization

感受野重分布使得我们可以以绩效的计算提升享受大幅内存降低，但是需要对不同架构采用逐个手工调整。峰值内存降低也为架构设计带来了更大的自由度。

为探索更大的设计空间，我们采用neural architecture与inference scheduling自动联合优化方案(可参考上图)，旨在获得满足所有约束且具有最高精度的架构。

Backbone optimization  我们参考MCUNet采用了类MnasNet搜索空间，而类MobileNetV3搜索空间因Swish激活函数问题导致难以量化而弃用。作者认为：最佳的搜索空间配置不仅硬件相关，同样任务相关。因此，我们还将r与w纳入搜索空间。

Inference scheduling optimization  给定模型与硬件约束，我们将寻找最佳推理机制。我们的推理引擎基于MCUnet中的TinyEngine扩展而来，除了TinyEngine中已有的优化外，我们还需要确定块数量p与模块数量n以执行patch推理，确保满足SRAM约束。

Joint Search  我们需要协同设计骨干优化与推理机制。比如，给定相同约束，我们可以选更小的模型以layer方式推理，或更大的模块以patch方式推理。因此，我们两者纳入优化并采用进化搜索寻找满足约束的最佳组合。

Experiments

Memory Profiling  内存占用是独立于推理框架的，为更好的比例，我们研究了两种分析设置。

上图给出了Analytic profiling的结果，它仅与架构有关而与推理框架无关。可以看到：

patch推理可以大幅降低峰值内存达3.7-8.0x，仅增加8-17%计算量；

对于MobileNetV2，通过重分布感受野可以将额外的计算量从10%降到3%；

内存节省与计算降低与网络架构相关，故网络架构与推理引擎的协同设计很有必要 。

上图给出了on-device profiling的结果，它不仅与架构相关还与硬件相关。可以看到：patch推理可以降低SRAM峰值内存达4-6x 。某些模型可能会因初始阶段的低硬件利用导致更长的推理延迟，但合理的架构设计可以将额外推理延迟降低到4%，相比内存降低完全可以忽略。

Tiny Image Classification

上表给出了ImageNet上的分类性能，从中可以看到：

在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数据机上的检测性能对比，从中可以看到：

在H7 MCU上，相比MCUNet，MCUNetV2-H7的性能提升达16.7%mAP；

在更廉价的M4MCU上，MCUNetV2-M4仍可提升13.2%mAP指标，峰值SRAM降低1.9x。

MUCNetV2-M4与MCUNet具有相似的计算量，但更高的性能。这是因为：patch推理导致的扩展搜索空间使得我们可以在更大的输入分辨率与更小的模型中作出更佳的选择。

上表给出了WIDER FACE数据集上的性能对比，从中可以看到：

在不同尺度下，MCUNetV2均优于现有方案；

MCUNetV2-L取得了同等检测性能，但峰值内存被RNNPool-Face-C小3.4x，比LFFD小9.9x；计算量分别小1.6x与8.4x；

MCUNetV2-S同样具有比RNNPool-Face-A与EagleEye小1.8x的峰值内存。

下图还给出了检测效果的对比，MCUNetV2的鲁棒性更强。

本文仅做学术分享，如有侵权，请联系删文。