译：《MCUNetV2：用于微深度学习的内存高效的基于补丁的推理》

论文地址：MCUNetV2: Memory-Efficient Patch-based Inference for Tiny Deep Learning

Motivation

如今基于微硬件的如微处理器单元(microcontroller units, MCUs)的IoT设备远处不在。而在这样的微硬件上部署深度学习模型对于人工智能的普及非常有用。但是由于内存紧张，微深度学习和移动深度学习有本质区别：一个普通的MCU通常SRAM小于512kB，这对于部署大多数现有深度学习网络而言太小了。这就提出了一个新的挑战以使用更低的内存峰值高效地进行AI推理。

已有的高效深度学习工作包括剪枝、量化和NAS，但这些方法集中在降低参数和FLOPs的数量，而不是内存瓶颈。内存预算的紧张限制了特征图/激活的大小，也限制我们使用小的模型容量或小的输入图像大小。实际上，已有的tinyML工作中的输入分辨率通常很小(<$224^2$)，这对于图片分类通常是可接受的，但对于像目标检测这样的密度预测任务却不太可行，如图：

基于这些，作者提出MCUNetV2：用于微深度学习的内存高效的基于补丁的推理。

作者贡献：

系统分析高效的CNN设置的内存使用模式，发现它们的内存分布不平衡，为优化留下巨大的空间。
提出一种基于补丁的推理调度方法，来显著降低CNN模型运行所需的峰值内存，同时通过重新分配感受野来最小化计算开销。
通过网络结构和推理调度联合设计，在MCUs上实现了对微小图像的分类和目标检测。

方法简介

系统分析微深度学习的内存瓶颈

在深度层面对每层的内存使用进行分析，结果如图，可以发现其存在不平衡激活内存分布，以MobileNetV2为例，其前5块有较高的峰值内存(>450kB)，这就是其对应的整个网络的内存瓶颈，而剩下的13个块对应的内存使用则较少。

作者还检查了其它高效网络设计，发现这种现象在不同的CNN主干中相当普遍。

此外，作者还发现这种情况适用于大多数单分支或残差CNN设计，这是由于分层结构造成的：在每个阶段之后，图像分辨率被降采样一半，导致像素减少4$\times$，而通道只增加了2$\times$或更小的比例，这就导致激活大小下降。

因此，内存瓶颈往往出现在网络的早期阶段，之后峰值内存使用会小很多。

基于这些，作者认为：如果能找到一种”绕过“初始内存密集阶段的方法，就可以大大降低两个网络的峰值内存，为优化留下更大空间。

MCUNetV2

使用基于补丁的推理打破内存瓶颈

作者提出使用一种基于补丁的推理方法(patch-based inference)来打破刚开始几层的内存瓶颈，如图

已有的深度学习推理框架(如： TensorFlow Lite Micro、TinyEngine、microTVM等)都是使用逐层(layer-by-layer)执行的方式，对于每个卷积层，推理库首先在SRAM中分配输入和输出激活缓冲，然后在整个层计算完成后释放输入缓冲。这种实现使得推理优化更容易，但SRAM必须为每层保存整个输入和输出激活，这对于CNN的初始阶段来说太大了。

而本文提出的基于补丁的推理则是以逐补丁形式在初始的内存密集阶段运行。每次，只在一个小的空间区域(比整个区域小10$\times$以上)运行模型，这可以大大降低峰值内存使用。在这个阶段完成之后，剩下的具有小峰值内存的网络则以正常的逐层方式执行。

但是，这样显著的内存节省是以计算开销为代价的，为了保持与每层推理相同的输出结果，不重叠的输出补丁对应输入中的重叠补丁(图中阴影区域)，这是因为内核大小>1的卷积滤波器有助于增加感受野，输出补丁的边界像素依赖于相邻补丁的输入，如下图，即使在最佳超参数选择下，这样的重复计算也会使得整个网络的计算量增加10~17%。

通过感受野重新分配降低计算开销

计算开销与基于补丁的初始阶段的感受野相关，考虑基于补丁阶段的输出，输入图像的感受野越大，每个补丁的分辨率越大，导致重叠面积和重复计算越多。对于MobileNetV2，如果只考虑下采样，每个输入补丁的边长为224/4=56，但考虑到感受野的增加，每个输入补丁都必须使用75$$75的形状，导致很大的重叠区域。

因此，本文提出重新分配CNN的感受野(receptive field，RF)来降低计算开销。基本思想是：

降低基于补丁的初始阶段的感受野

这有助于降低每个输入补丁的大小和重复计算量
提升后面阶段的感受野

这有助于补偿性能损失

以MobileNetV2为例，如下图，作者在每个补丁推理阶段使用更小的核和更少的块，并且增加后而推理阶段每层的块的数量，这个过程需要手动调整，并且不同情况调整也不同。

在这样重新分配之后，计算开销仅有3%，这对于减少内存的好处而言是可以忽略的。

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

由于之前MobileNetV2重新设计中所说的各种需要手动调整的过程，作者提出以自动化的方式联合优化神经结构和推理调度。给定确定的数据集和硬件限制，目标是实现满足所有限制的最高准确率。

骨架优化

使用MnasNet-alike搜索空间，对于每个倒残差块包括不同的内核大小$k_{[]}(3/5/7)$，不同的扩展比$e_{[]}(3/4/6)$，每个阶段不同的块数量$d_{[]}(2/3/4)$。作者认为：最好的搜索空间配置不仅是硬件感知的，而且是任务感知的：例如一些任务可能更喜欢更高分辨率而不是更大的模型大小，反之亦然。因此，作者在搜索空间中加入$r$和$w$，进一步扩展$w$来支持每块宽度缩放$w_{[]}$，包括$w_{[]}(0.5/0.75/1.0)$和$r(96-256)$，这增加了搜索空间的可扩展性，使得能适应不同的MCU模型和紧张的资源预算

推理调度优化

推理引擎基于TinyEngine开发，除了TinyEngine的优化选项，还需要补丁数量$p$和块数量$n$来执行基于补丁的推理。

联合搜索

需要协同设计骨架优化和推理调度，例如，给定一样的限制，可以选择使用更小的模型来满足逐层执行或者更大的模型和逐补丁推理。

Zephon Blog

《MCUNetV2:Memory-Efficient Patch-based Inference for Tiny Deep Learning》