NeurIPS 2023 | MCUFormer: 可在内存有限微控制器上部署的Vison Transformer模型

科技

2023-11-09 05:11

©PaperWeekly 原创 · 作者 | 岳廷

研究方向 | 计算机视觉

论文标题：

MCUFormer: Deploying Vision Transformers on Microcontrollers with Limited Memory Resources

论文链接：

https://arxiv.org/pdf/2310.16898

代码链接：

https://github.com/liangyn22/MCUFormer

要解决的问题：由于 GPU 的价格高昂且能耗大，因此在物联网设备（如微控制器）上部署深度模型对生态人工智能有重大贡献。传统方法成功地在微控制器上实现了高分辨率图像的卷积神经网络推理，但在许多视觉应用中取得最先进性能的 Vision Transformer 框架仍然未被探索。

解决方案：文章提出了一种名为 MCUFormer 的硬件-算法联合优化方法，用于在内存极其有限的微控制器上部署 Vision Transformer。具体来说，作者们共同设计了 Transformer 架构，并构建了推理操作库以适应内存资源限制。他们将一次性网络架构搜索（NAS）推广到微控制器的内存预算下，以发现具有最高任务性能的最优架构。为了减少内存，他们扩大了 Vision Transformer 的现有搜索空间，考虑了低秩分解维度和 Patch 分辨率。

效果：实验结果表明，MCUFormer 在 STM32F746 微控制器上使用 320KB 内存进行图像分类时，在 ImageNet 上达到了 73.62% 的 top-1 准确率。

额外说明：该文章提出的方案不仅在物联网设备上有较大的实用意义，在任何终端设备上都有较大的实际意义。因为终端设备价格为王，更小的内存消耗意味着更低的设备价格或者更加流畅的设备体验（很多设备中，如手机电脑平板，电视机等，足够的内存空余意味更加流畅的设备体验）。因此终端设备如手机平板等上都可参考！

文章的贡献总结如下:

文章提出了第一个成功在具有有竞争力性能的 MCU 上部署 Vison Transformer 的硬件算法协同优化框架
文章提出了 Vison Transformer 内存限制下的有效架构搜索方法，并构建了具有高效内存调度的运算符库以实现实际部署
文章进行了广泛的实证研究，以证明在极端有限内存的 MCU 上部署 Vison Transformer 的可行性

具有内存约束的网络架构搜索流程：

方法

在这一节中，文章详细介绍了在微控制器上部署 Vison Transformer 的硬件算法协同优化框架。文章首先介绍具有有限内存约束的 Vison Transformer 网络架构搜索，然后演示微控制器上 Vison Transformer 推理的运算符库构建。

1.1 内存约束下的网络架构搜索

传统的架构搜索方法通常考虑计算复杂度（FLOPs）和存储成本（参数数量），而峰值内存由于微控制器极端有限的 SRAM 而隐藏了 Vison Transformer 的部署。文章扩展一次性 NAS 以搜索最优架构，其中文章首先学习一些随机包含所有可能拓扑选择的超网，并利用进化算法根据资源预算搜索最优超网。图 1 展示了具有内存约束的网络架构搜索流程。

为了充分利用 Vison TransformerNAS 方法的成就，文章在 AutoFormer 的基础上构建子网搜索空间，AutoFormer 使用训练好的超网络搜索嵌入维度、QKV 维度、MLP 比率、头数和层数。超网的搜索空间由低秩分解比率和 Token 数定义，它们可以显著影响前向传播期间的峰值内存。

在巨大的搜索空间中有效地获取有效的超网，文章在给定低秩分解比率和 Token 数组成的固定搜索空间中学习超网，并通过预测任务性能与搜索空间之间的相关性来演化超网优化的搜索空间。为低秩分解，文章推广了多层感知机（MLP）层的形式如下：

其中表示具有个 Token 和每个 Token 个元素的输入特征。和是两个低秩矩阵近似 MLP 层中的原始权重，表示偏置项。由于低秩矩阵的维度影响内存占用和任务性能之间的权衡，文章需要在搜索空间选择中找到最佳的低秩分解比率。

将 Token 数和低秩分解比率直接集成到 AutoFormer 中的原始搜索空间通常会导致超网收敛问题，因为不同低秩分解比率的最优架构之间存在很大差异。因此，文章首先优化超网络的搜索空间，它由 Token 数和低秩分解比率组成，然后在超网络搜索空间的更新版本中执行架构搜索。以上两步迭代实现，直到收敛或达到最大搜索成本。

为架构搜索枚举所有搜索空间是不可行的，因为给定搜索空间中的单个超网的训练成本非常高。启发于文献 [12] 的观察，即任务错误与连续选择之间存在强相关性，文章通过考虑性能、内存占用和搜索空间因素（包括 Token 数和低秩分解比率）之间的关系来演化搜索空间。用于评估不同超网络搜索空间优化程度的得分函数定义如下：

其中是搜索空间中的超网络准确率，是一个超参数。是内存占用在 SRAM 限制内的子网络比例。同时，枚举超网中的所有子网络以获取平均内存是不可行的，因为子空间中的候选网络数量极大，文章随机采样个子网络来近似估计给定超网络的值。使用线性拟合函数估计得分与搜索空间因素（低秩分解比率和 Token 数）之间的倾向关系，即，其中是截距，由低秩分解比率和 Token 数组成。根据以下规则更新搜索空间：

其中表示第步中的的第个元素，表示中的第个元素。优化步长在文章的方法中设置为相邻因子值的间隔，更新阈值是优化的超参数。表示小于的最大整数。

训练超网络至收敛需要不可接受的优化成本，文章在训练过程中拟合任务性能与搜索空间连续选择的相关性。对于每个演化的搜索空间，文章从恢复点或从头开始训练超网个时期。为了公平地发现任务性能与连续选择之间的倾向关系，文章使用相同的总训练时期中的数据点拟合相关性权重矩阵。

由于线性拟合函数仅在给定搜索空间的小局部区域内成立，文章利用片 wise 线性拟合函数，其中仅考虑靠近给定搜索空间的前 k 个最近邻进行拟合。数据点采样策略表示如下：

其中表示用于给定演化搜索空间的线性函数拟合的由任务准确率、内存与搜索空间连续因子组成的数据点集合。的近邻集表示为，搜索空间中的超网总训练时期表示为。当给定搜索空间的领域为空时，文章将总训练时期为的前近邻构成领域。文章在这些搜索空间中训练超网络个时期以更新准确率，以公平地估计给定搜索空间的相关矩阵。

在最优搜索空间中选择的超网用于后续训练直到完全收敛，并用于 Vison Transformer 的网络架构搜索。在此过程中，应用进化算法以最大限度地提高准确率同时遵守内存约束。

1.2 Vison Transformer推理运算符库的构建

在微控制器上部署 Vison Transformer 用于推断需要运算符库将模型转换为可执行函数。现有的深度学习运算符库可与卷积神经网络比较，而无法实现 Vison Transformer 所需的运算，如 LayerNorm 和 GeLU。同时，库的推理效率无法充分利用微控制器的内存预算，无法部署具有高任务性能的较大 Vison Transformer 模型。

尽管在 [11] 中提出了适用于所有网络架构的有效内存调度技术，但文章提出以下内存利用增强技术以启用 Vison Transformer 在微控制器上的部署。

Patch 嵌入分解：在 Vison Transformer 推理中，图像通过 Patch 嵌入层进行 Patch 化。输入样本通过卷积运算转换为 Patch，其中内核尺寸和步长都等于 Patch 大小。由于 Patch 嵌入层中的内核尺寸（大于）远大于常见的卷积神经网络，Patch 嵌入层的大部分内存消耗来自滤波器权重而不是输入和输出激活。

为了有效解决此问题，文章将大感受野卷积分解为多个小感受野卷积运算。为了将多个卷积运算的结果相加，文章使用一个小内存缓冲区来记录累加的中间值，从而有效减少峰值内存使用。图 2（a）展示了原始 Patch 嵌入运算与文章的分解之间的比较。将 Patch 嵌入层分解为过多的小感受野卷积运算可能明显增加推理延迟，因此文章选择分解 Patch 嵌入层的感受野为，以实现内存占用与推理延迟之间的满意折衷。

运算符集成：网络参数量化和用整数算术替代 MAC 运算已被证明在不造成性能下降的情况下有助于减少卷积神经网络的内存。但是，具有量化运算符的常规运算符库由于特定运算符包括 GeLU 激活和层规范化而无法用于 Vison Transformer 推理。由于 GeLU 需要高斯分布的累积分布函数（CDF），文章用量化的乘法和 softmax 函数逼近激活函数，两者都可以用 int8 整数实现。类似地，运算符库中的整数算术不支持层规范化中的平方根运算符。

文章构建了用定点迭代方法计算平方根运算符输出的替代方程式。考虑层规范化运算符的推理延迟，文章仅对替代方程式进行3次迭代以在推理延迟与预测精度之间达到平衡。文章在附录中详细介绍了仅用整数的平方根运算伪代码。

Token 覆盖：在 Vison Transformer 推理中造成高内存消耗的矩阵乘法仅在 Token 内部施加元素交互，而 Token 之间的交互在内存消耗远低于 SRAM 限制的自注意力计算中启用。因此，文章可以在全连接层的推理期间覆盖中间特征，因为每个 Token 的输入激活在获取对应的 Token 输出特征后将不再使用。

图 2（b）说明了现有推理运算符库中的矩阵乘法内存调度和文章的 Token 覆盖技术，其中文章的方法在相同的内存空间写入矩阵乘法的每个 Token 的输出，与输入激活相同。由于峰值内存使用发生在矩阵乘法中，文章根据网络架构为 Vison Transformer 推理分配最大内存大小，根据以下标准：

其中，和分别表示第层全连接层中权重、输入特征和输出特征的高度，，和表示上述张量的宽度。由于输入激活消耗的内存可以被输出特征覆盖，文章仅分配与输入和输出中较大者相同大小的内存缓冲区。

实验

在这一节中，文章首先介绍硬件配置、数据集和硬件算法协同优化框架的实现细节，然后文章进行全面的性能分析，包括与网络架构搜索和运算符库构建中的技术相关的任务性能和内存消耗。最后，文章与现有的微控制器深度模型部署的网络架构搜索方法和运算符库进行比较。

2.1 设置

硬件配置：文章在不同的微控制器上部署 Vison Transformer，这些微控制器具有不同的资源约束，包括 STM32F427（Cortex-M4/256KB 内存/ 1MB 闪存）、STM32F746（Cortex-M7/ 320KB 内存/1MB 闪存）和 STM32H743（Cortex-M7/512KB 内存/ 2MB 闪存）。在性能分析中，文章在 STM32F746 上评估文章的框架，以获得准确性和内存消耗。

数据集：文章在 ImageNet 上进行实验进行图像分类，其包含 120 万幅训练图像和 5 万幅验证图像，涵盖 1000 个类别。所有图像均缩放和偏置到 [-1,1] 范围进行归一化。在训练过程中，文章将图像的较短边调整为 256，并随机裁剪 240×240 的区域。推理时，文章使用 240×240 的中心裁剪。

实现细节：对于 Vison Transformer 的网络架构搜索，文章的低秩分解比率和 Token 大小的选择可以从和中选取。超网设计和进化搜索与 [11] 中的保持一致，以充分利用具有高自由度的架构搜索的潜力。

在搜索空间演化过程中，文章仅选择具有相同总训练时期的 5 个最近邻来拟合分段线性函数，以获取斜率。每次演化，文章在存在的恢复点或从头开始在选定的搜索空间中训练超网 30 个时期。如果给定超网络的邻域集为空，文章选择总训练时期比给定超网络少 30 的前 5 个最近邻。

然后，文章在这些搜索空间中训练超网络以获取准确率，以公平地估计给定搜索空间的相关性矩阵。初始化的数据点用于拟合趋势的连续选择是随机采样的。对于运算符库构建，文章在推理期间使用 int8 量化对 Vison Transformer 中的所有张量进行量化。分解 Patch 嵌入层的滤波器大小设置为，具有多次前向传递以减少峰值内存，文章迭代了从定点迭代方法获得的替代分配4次以计算层规范化中的平方根。超参数分配和近似方法在附录中列出。

2.2 性能分析

搜索空间演化的效果：低秩分解比率和 Token 数显著影响内存，这对于 Vison Transformer 的部署很重要。为了发现在任务性能和内存之间取得满意折衷的最优搜索空间，文章有效地演化搜索空间。文章比较了不同的搜索空间选择标准，包括随机采样搜索空间、最大搜索空间和将搜索空间选择组合到超网学习中。

表 1 显示了不同方法的 top-1 准确率、内存占用和闪存，其中文章的方法在准确率和效率方面都优于其他方法。随机选择搜索空间无法发现任务性能与存储和内存效率之间的最佳折衷，而选择内存约束下子网络比例最低的最大复杂度搜索空间则优于演化方法，因为内存限制内的子网络比例很低。将搜索空间选择直接组合到超网训练中会将搜索复杂度增加 15-16 个数量级，搜索缺陷阻碍了获取最佳网络架构。

运算符库中的内存调度效果：为了在微控制器上执行 Vison Transformer，文章构建了一个运算符库将模型转换为可执行函数。此外，文章还通过 Patch 嵌入分解、运算符集成和 Token 覆盖重新调度推理期间的内存，以充分利用微控制器的内存约束。

图 3 展示了具有不同内存调度技术的运算符库的峰值内存比较。运算符集成对内存减少贡献最大（23.8%），因为 int8 张量表示的内存仅为 float 的 25%。Patch 嵌入分解也是必要的，因为大的卷积核需要大的内存空间进行 MAC 运算。Token 覆盖也通过在每次矩阵乘法中重用缓冲区来减少内存消耗，因为只需要用于表示较大一个的输入和输出的内存即可。

总之，将所有这些技术集成到运算符库中可以减少 9% 的内存消耗，并允许在微控制器上部署大规模 Vison Transformer 而具有满意的任务性能。

拟合趋势的影响：搜索空间根据连续选择与任务性能之间拟合的趋势演化，有效的趋势估计对获取最佳网络架构很重要。考虑到简化和拟合精度，文章利用分段线性函数来拟合具有邻近给定搜索空间的数据点的趋势。大的邻域对异常值更具鲁棒性，而小的邻域为选定的搜索空间生成更准确的拟合函数。

文章研究邻域大小对拟合趋势的影响，表 2 描绘了不同邻域规模的数据点采样结果的推理效率和 top-1 准确率。在准确率和内存约束内选择邻域中的适度数据点数量导致最高准确率，这是由于有效的搜索空间演化。

文章还提供了一个搜索空间演化示例的可视化，如图 4 所示，其中演化过程中被选择更频繁的搜索空间被验证为更优。不同颜色表示各候选的演化路径，箭头表示演化的方向。双箭头表示演化后的候选人在当前总训练时期之后再训练次。

搜索空间演化中超参数的影响：搜索空间演化中的关键超参数包含更新阈值和，文章在敏感性分析中改变其值以研究其影响。图 5 展示了不同阈值分配下的 top-1 准确率，其中中等阈值导致最高性能。大阈值导致没有足够置信的频繁搜索空间变化，而小阈值强制演化过程停留在局部最优点。同时，性能对低秩分解比率的更新阈值更敏感，它明显改变了最优架构。

2.3 大规模图像分类评估

在本节中，文章将 MCUFormer 中的网络架构搜索技术与最先进的架构设计方法进行比较，包括基于 Vison Transformer 的 NAS 方法 AutoFormer。文章还将运算符库与针对卷积神经网络推理设计的现有运算符库进行了比较。由于常规运算符库如 CMSIS-NN 和 TinyEngine 无法直接用于 Vison Transformer 推理，因为缺少可执行的层规范化和 GeLU 等特定函数，文章补充了这些缺失的函数以确保公平比较 Vison Transformer 推理中的内存调度。

图像分类是一项具有挑战性的任务，主要是因为其大规模和高多样性，特别是当使用轻量级模型部署在微控制器上时。表 3 展示了不同的深度神经网络包括基于 CNN 和 Transformer 的架构以及各种微控制器深度学习库的 top-1 分类准确率、内存占用和闪存。表中的灰色数字表示所需的内存或闪存超过微控制器的内存或闪存约束。对于超过微控制器内存约束的架构与运算符库的组合，文章报告了使用 GPU 进行推理时达到的准确率。

对于基于 CNN 的架构，文章采用 TinyEngine 进行编译。MCUFormer-A 和 MCUFormer-T 分别表示文章的方法，其中网络架构搜索分别与 AutoFormer 替换，并使用补充了 Vison Transformer 缺失层的 TinyEngine 进行编译。除 MCUFormer-T 外，文章使用文章提出的引擎编译所有基于 Transformer 的架构。

由于 Vison Transformer 在大规模视觉任务中明显优于 CNN，因此在微控制器上部署 Transformer 是很可取的。与微控制器上的最先进 CNN 模型 EMQ 相比，文章的方法在 320KB 内存限制下提高了 5.4% 的 top-1 准确率（68.2% vs. 73.6%）。搜索空间演化中的架构搜索和运算符库中的内存调度都有助于在 MCU 上部署具有高准确率的 Vison Transformer，通过比较 MCUFormer-T、MCUFormer-A 和 MCUFormer 可以看出，因为去除了冗余表示消耗 MCU 中的内存。

尽管最先进的 NAS 方法 AutoFormer 可以有效发现 Vison Transformer 的拓扑，以实现最佳的准确率复杂性折衷，但它们未考虑微控制器的极端内存约束，这对 GPU 不是问题。直接在原始搜索空间中补充低秩分解比率和 Token 数通常会导致搜索缺陷，因为超网训练中的收敛问题。同时，TinyEngine 优化了通用深度网络部署的内存调度，但仍忽略了 Vison Transformer 特定架构的进一步内存利用率提高。

总的来说，架构搜索中的搜索空间演化和库构建中的内存调度分别提高了 top-1 准确率 4.3%（69.3% vs. 73.6%）和减少了峰值内存 63%（319kB vs. 872kB），使用 STM32F746 获得了更好的结果。

文章还在具有不同资源约束的不同控制器上部署了硬件算法协同优化框架，表 3 还演示了各种方法的 top-1 准确率。在 256KB 内存和 1MB 闪存的约束下，文章的 Vison Transformer 仍然在 ImageNet 上达到了 71.1% 的 top-1 准确率，比 EMQ 方法高出 6.9%。

随着内存和闪存约束的增加，任务性能有所提升，STM32H743 上的 512KB 内存和 2MB 闪存获得了 74.0% 的 top-1 准确率，甚至比在 GPU上推理的 Deit 高出 1.8%。因此，文章的 Vison Transformer 可以应用于对准确率有较高要求的实际应用中。

结论

文章提出了 MCUFormer 进行硬件算法协同优化，它成功地在内存极端有限的微控制器上部署了 Vison Transformer，同时保持了满意的任务性能。文章推广了网络架构搜索，扩大了搜索空间，考虑了权重矩阵的低秩分解比率和 Token 数。搜索空间演化用于超网训练，以在资源约束下发现具有最高准确率的模型拓扑，其中演化利用了任务性能与连续选择之间的相关性，提高了搜索效率。

同时，文章构建了 Vison Transformer 推理的运算符库，将模型转换为可执行函数，并通过 Patch 嵌入分解、运算符集成和 Token 覆盖全面利用内存。广泛的实验结果表明，文章的 MCUFormer 在具有竞争力性能的微控制器上实现了图像分类、人员存在识别和关键词识别等任务。文章的工作目前包含以下限制：考虑在更多不同视觉任务（如 DETR 进行对象检测）中在 MCU 上部署 Transformer 对实际应用如自动驾驶和机器人导航很重要。

更多阅读