MCUNetV2：面向微型深度学习的内存高效分块推理方法

Ji Lin, Wei-Ming Chen, Han Cai, Chuang Gan, and Song Han. 2021. MCUNetV2: memory-efficient patch-based inference for tiny deep learning. In Proceedings of the 35th International Conference on Neural Information Processing Systems (NIPS '21). Curran Associates Inc., Red Hook, NY, USA, Article 180, 2346–2358.

1. 引言与动机

在物联网时代，基于微控制器单元（MCU）的微型硬件无处不在。在这些微型硬件上部署深度学习模型将使我们能够实现人工智能的民主化。然而，由于内存预算极其紧张，微型深度学习与移动深度学习有着根本性的不同：一个常见的MCU通常具有小于512KB的SRAM，这对于部署大多数现成的深度学习网络来说太小了。即使对于更强大的硬件如Raspberry Pi 4，将推理适配到L2缓存（1MB）中也能显著提高能效。这些都对具有小峰值内存使用的高效AI推理提出了新的挑战。

现有工作采用剪枝、量化和神经架构搜索来实现高效深度学习。然而，这些方法专注于减少参数数量和FLOPs，而不是内存瓶颈。紧张的内存预算限制了特征图/激活值的大小，限制我们使用小的模型容量或小的输入图像尺寸。实际上，现有tinyML工作中使用的输入分辨率通常很小（< 224²），这对于图像分类（例如ImageNet、VWW）可能是可以接受的，但对于密集预测任务如目标检测就不够了。

图2描述：论文中的图2展示了目标检测相比图像分类对输入分辨率更加敏感。图中显示了不同输入分辨率（160、224、288、352）下的Pascal VOC mAP和ImageNet Top-1准确率。可以清楚地看到，随着分辨率降低，VOC mAP的下降速度明显快于ImageNet Top-1准确率，这说明了目标检测任务对输入分辨率的高度依赖性。

2. 深入理解微型深度学习的内存瓶颈

2.1 内存分布的不平衡性

通过系统分析CNN模型的内存瓶颈，论文发现了一个关键现象。以MobileNetV2为例，提供了int8量化下的每块峰值内存使用情况。

图1描述：图1展示了MobileNetV2具有非常不平衡的内存使用分布。峰值内存由前5个具有高峰值内存的块决定，而后面的块都共享小的内存使用。通过使用分块推理（4×4块），能够显著减少前5个块的内存使用，将整体峰值内存减少8倍，从1372kB降至172kB，适配256kB内存预算的MCU。重要的是，两种设置的模型架构和准确率都没有改变。我们也可以观察到内存使用分布不平衡的清晰模式，前5个块的峰值内存很大，超过了MCU的内存约束，而其余13个块轻松适配256kB内存约束。第三个块的内存使用比网络其余部分大8倍，成为内存瓶颈。通过检查其他高效网络设计，发现这种现象在不同CNN backbone中相当普遍，即使对于专门为内存受限微控制器设计的模型也是如此。

这种情况适用于大多数单分支或残差CNN设计，原因在于层级结构：每个阶段后，图像分辨率下采样一半，导致像素减少4倍，而通道数仅增加2倍或更小的比例，导致激活值大小递减。因此，内存瓶颈倾向于出现在网络的早期阶段，之后峰值内存使用要小得多。

2.2 挑战与机遇

不平衡的内存分布显著限制了在MCU上可执行的模型容量和输入分辨率。为了容纳初始内存密集阶段，整个网络需要缩小，即使网络的大部分已经具有小的内存使用。它还使分辨率敏感任务（如目标检测）变得困难，因为高分辨率输入将导致大的初始峰值内存。

考虑MobileNetV2中的第一个卷积层，输入通道3，输出通道32，步长2，在224×224分辨率的图像上运行它需要的内存为：

$$\text{Memory} = 3 \times 224^2 + 32 \times 112^2 = 539\text{kB}$$

即使在int8量化时，这也无法适配微控制器。另一方面，如果我们能找到一种方法来"绕过"初始内存密集阶段，我们可以大大减少整个网络的峰值内存，为优化留下巨大的空间。

3. MCUNetV2：内存高效的基于分块的推理

3.1 通过基于分块的推理打破内存瓶颈

论文提出通过分块推理来打破初始层的内存瓶颈。现有的深度学习推理框架（如TensorFlow Lite Micro、TinyEngine、microTVM等）使用逐层执行。对于每个卷积层，推理库首先在SRAM中分配输入和输出激活缓冲区，并在整个层计算完成后释放输入缓冲区。这种实现使推理优化变得容易（如im2col、tiling等），但SRAM必须保存整个输入和输出激活。

图3描述：图3展示了分块推理如何显著减少执行卷积层序列所需的峰值内存。研究了两个卷积层（步长1和2）。在逐层计算下（a），第一个卷积具有大的输入/输出激活大小，主导了峰值内存需求。通过分块计算（b），分配缓冲区来托管最终输出激活，并逐块计算结果。只需要存储来自一个块的激活而不是整个特征图，减少了峰值内存。第一个输入是图像，可以从压缩格式（如JPEG）部分解码，不需要完全存储。

3.2 计算开销分析

显著的内存节省以计算开销为代价。为了保持与逐层推理相同的输出结果，不重叠的输出块对应于输入图像中的重叠块（图3(b)中的阴影区域）。这是因为核大小>1的卷积滤波器有助于增加感受野。输出块上的边界像素依赖于来自相邻块的输入。这种重复计算即使在最优超参数选择下也可以将整体网络计算增加10-17%，这对于低功耗边缘设备是不希望的。

3.3 通过重新分配感受野减少计算开销

计算开销与基于分块的初始阶段的感受野相关。考虑基于分块阶段的输出，它在输入图像上的感受野越大，每个块的分辨率就越大，导致更大的重叠区域和重复计算。

对于MobileNetV2，如果只考虑下采样，每个输入块的边长为224/4 = 56。但当考虑增加的感受野时，每个输入块必须使用75×75的形状，导致大的重叠区域。

论文提出重新分配CNN的感受野（RF）以减少计算开销。基本思想是：

1. 减少基于分块的初始阶段的感受野
2. 增加后期阶段的感受野

减少初始阶段的RF有助于减少每个输入块的大小和重复计算。然而，如果整体RF较小，某些任务可能会有性能下降（例如检测大物体）。因此，进一步增加后期阶段的RF以补偿性能损失。

图4描述：图4展示了重新分配的MobileNetV2（MbV2-RD）具有减少的分块推理阶段感受野和增加的逐层阶段感受野。两个网络具有相同的性能水平，但MbV2-RD在基于分块的推理下具有更小的开销。移动倒置块表示为MB{扩展比}{核大小}。虚线边框表示stride=2。

3.4 联合神经架构和推理调度搜索

重新分配感受野使我们能够以最小的计算/延迟开销享受内存减少的好处，但策略因不同backbone而异。减少的峰值内存也为设计backbone架构提供了更大的自由度（例如，使用更大的输入分辨率）。

为了探索这样一个大的设计空间，论文提出以自动化方式联合优化神经架构和推理调度。给定某个数据集和硬件约束（SRAM限制、Flash限制、延迟限制等），目标是在满足所有约束的同时实现最高准确率。

优化的维度包括：

Backbone优化：

• 每个倒置残差块的不同核大小 k[] ∈ {3, 5, 7}
• 不同的扩展比 e[] ∈ {3, 4, 6}
• 每个阶段的不同块数 d[] ∈ {2, 3, 4}
• 输入分辨率 r ∈ [96, 256]
• 全局宽度乘数 w 和分块宽度缩放 w[] ∈ {0.5, 0.75, 1.0}

推理调度优化：

• 块数 p（将图像分割成 p×p 个重叠块）
• 执行基于分块推理的块数 n

4. 实验结果与深入分析

4.1 现有网络的峰值内存减少

图5描述：图5展示了分析性能分析结果。分块推理显著减少了推理峰值内存3.7-8.0倍，同时仅产生8-17%的计算开销。内存减少和计算开销与网络设计相关。对于MobileNetV2，通过重新分配感受野可以将计算开销从10%减少到3%。所有网络采用224²的输入分辨率和4×4块。

图6描述：图6展示了设备上的测量结果。分块推理在MCU上运行时将测量的峰值SRAM使用减少了4-6倍。某些网络的延迟开销可能很大，但通过适当的架构设计（MbV2-RD）可以将其减少到4%。由于MCU内存限制，必须缩小宽度w和分辨率以适配。

4.2 MCUNetV2用于微型图像分类

在ImageNet上，MCUNetV2在256kB SRAM/1MB Flash下实现了64.9%的Top-1准确率，在512kB SRAM/2MB Flash下实现了71.8%的创纪录准确率。

图7描述：图7左图展示了MCUNetV2在视觉唤醒词（VWW）准确率vs峰值SRAM权衡方面具有更好的表现。与MCUNet相比，MCUNetV2在4.0倍更小的峰值内存下实现了更好的准确率。它在<32kB内存下实现了>90%的准确率，便于部署在极小的硬件上。右图展示了基于分块的方法扩展了可以适配MCU的搜索空间，允许更好的准确率vs延迟权衡。

4.3 MCUNetV2用于微型目标检测

目标检测对较小的输入分辨率敏感。当前最先进的方法由于分辨率瓶颈无法在MCU上实现良好的检测性能。MCUNetV2打破了检测器的内存瓶颈，将mAP提高了两位数。

在Pascal VOC上，MCUNetV2-H7将mAP提高了16.9%（从51.4%到68.3%）。它还可以缩小规模以适配更便宜的Cortex-M4 MCU（仅256kB SRAM），同时仍将mAP提高13.2%，峰值SRAM减少1.9倍。

4.4 超参数分析

图8描述：图8展示了基于分块推理的超参数研究。左图显示峰值内存通常随着更多块以分块方式执行（n）和更大的块数（p）而下降。MobileNetV2的最优索引是n*=5，其中特征图下采样8倍。右图显示将输入图像分割成更多块（更大的p）会导致更大的计算开销。感受野重新分配减少了开销（MbV2-RD）。

图9描述：图9展示了VWW上的MCUNetV2架构。颜色代表核大小；每个块的高度代表扩展比。名称是MB{扩展比}{核大小}x{核大小}。虚线边框的块具有stride=2。底部的{}x{}表示特征图分辨率。

5. 附录：数学推导与技术细节

A. 分块推理的内存

对于一个具有输入特征图大小 $H_{in} \times W_{in} \times C_{in}$ 和输出特征图大小 $H_{out} \times W_{out} \times C_{out}$ 的卷积层，传统逐层推理的峰值内存需求为：

$$M_{layer} = H_{in} \times W_{in} \times C_{in} + H_{out} \times W_{out} \times C_{out}$$

对于分块推理，假设我们将输入分成 $p \times p$ 个块，每个输出块的大小为 $\frac{H_{out}}{p} \times \frac{W_{out}}{p}$。考虑到感受野 $RF$ 和步长 $s$，每个输入块的实际大小为：

$$H_{patch} = \frac{H_{out}}{p} \times s + (RF - 1)$$

$$W_{patch} = \frac{W_{out}}{p} \times s + (RF - 1)$$

因此，分块推理的峰值内存为：

$$M_{patch} = H_{patch} \times W_{patch} \times C_{in} + \frac{H_{out}}{p} \times \frac{W_{out}}{p} \times C_{out}$$

内存减少比例为：

$$\text{Reduction} = \frac{M_{layer}}{M_{patch}} \approx p^2 \times \frac{1}{1 + \frac{(RF-1) \times p}{H_{out}}}$$

当 $RF \ll \frac{H_{out}}{p}$ 时，内存减少接近 $p^2$ 倍。

B. 计算开销分析

由于输入块的重叠，计算开销可以表示为：

$$\text{Overhead} = \frac{\text{Total\_Patch\_Area}}{\text{Original\_Area}} - 1$$

对于 $p \times p$ 个块，总的块面积为：

$$\text{Total\_Patch\_Area} = p^2 \times H_{patch} \times W_{patch}$$

原始面积为 $H_{in} \times W_{in}$，因此：

$$\text{Overhead} = \frac{p^2 \times (H_{out}/p \times s + RF - 1)^2}{(H_{out} \times s + RF - 1)^2} - 1$$

简化后：

$$\text{Overhead} \approx \frac{2(RF - 1) \times p \times s}{H_{out} \times s + RF - 1}$$

这表明计算开销与感受野 $RF$ 成正比，与块数 $p$ 成正比，这解释了为什么减少初始阶段的感受野可以减少计算开销。

C. 感受野重新分配的优化

给定总的感受野要求 $RF_{total}$，我们需要在分块阶段和逐层阶段之间分配：

$$RF_{total} = RF_{patch} \times RF_{perlayer}$$

优化问题可以表述为：

$$\min_{RF_{patch}, RF_{perlayer}} \text{Overhead}(RF_{patch}) + \lambda \times \text{Memory}(RF_{perlayer})$$

受约束于：

$$RF_{patch} \times RF_{perlayer} \geq RF_{total}$$

其中 $\lambda$ 是权衡参数。通过拉格朗日乘数法，可以得到最优的感受野分配。

D. 分块边界处理

在实际实现中，需要仔细处理分块边界。对于一个核大小为 $k$ 的卷积，相邻块之间的重叠区域为：

$$\text{Overlap} = \lfloor \frac{k-1}{2} \rfloor + \lceil \frac{k-1}{2} \rceil = k - 1$$

对于多层，重叠累积为：

$$\text{Total\_Overlap} = \sum_{i=1}^{n} (k_i - 1) \times \prod_{j=1}^{i-1} s_j^{-1}$$

其中 $k_i$ 是第 $i$ 层的核大小，$s_j$ 是第 $j$ 层的步长。

6. 结论

MCUNetV2通过提出基于分块的推理来减少微型深度学习的内存使用高达8倍，极大地扩展了设计空间并在IoT设备上解锁了强大的视觉应用。通过联合优化神经架构和推理调度，MCUNetV2显著提高了微控制器上的目标检测性能16.9%，并实现了71.8%的创纪录ImageNet准确率。对于VWW数据集，MCUNetV2可以在仅32kB SRAM下实现>90%的准确率，比现有工作小4倍。这项研究在很大程度上解决了tinyML中的内存瓶颈，为超越分类的视觉应用铺平了道路。