英文标题：AtomNet: Designing Tiny Models from Operators Under Extreme MCU Constraints
中文标题：AtomNet：在极端MCU约束下基于算子设计微型模型
作者信息
Zhiwei Dong<sup>1,3</sup>, Mingzhu Shen<sup>3</sup>, Shihao Bai<sup>3</sup>, Xiuying Wei<sup>3</sup>, Jinyang Guo<sup>2</sup>, Ruihao Gong<sup>2,3</sup>, Song-Lu Chen<sup>1</sup>, Xianglong Liu<sup>2</sup>, Xu-Cheng Yin<sup>1</sup>
<sup>1</sup>University of Science and Technology Beijing
<sup>2</sup>Beihang University
<sup>3</sup>SenseTime Research
*These authors contributed equally to this work.
Correspondence to: dongz.cn@outlook.com, m.shen23@imperial.ac.uk, baishihao@sensetime.com, w2295297125@163.com, jinyangguo@buaa.edu.cn, xhplus@163.com, songluchen@ustb.edu.cn, xlliu@nlsde.buaa.edu.cn, xuchengyin@ustb.edu.cn
论文出处
The Thirty-Ninth AAAI Conference on Artificial Intelligence (AAAI-25)
摘要
TinyML因其能够在边缘设备上提供低成本且即时的性能而备受关注。尤其是常用的微控制器单元（MCU）对峰值内存（SRAM）和存储（Flash）施加了极端限制。现有的TinyML方法通常依赖于定制且难以获取的推理库，并且需要使用高级神经架构搜索（NAS）算法来寻找可部署的架构，这耗时较长。为了解决这些问题，本文充分利用MCU中的资源，并推导出在极端MCU约束下设计模型的硬件导向指南。具体而言，通过收集Flash、SRAM和延迟的运行时数据，构建了一个名为AtomDB的数据集，并基于此建立了若干关键算子指南，以充分利用有限的Flash和SRAM，同时最小化延迟。通过将这些指南应用于分析模块，本文提出了一种混合模式，将适当的模块组织在网络的不同阶段，形成AtomNet，这是一种更硬件导向的架构，能够处理前阶段的SRAM瓶颈和后阶段的Flash瓶颈。大量实验验证了对硬件特性利用的有效性。值得注意的是，AtomNet在320KB的MCU上使用现成的官方推理库进行ImageNet任务时，首次实现了3.5%的准确率提升和超过15%的延迟降低，超越了当前最先进的方法。
引言
物联网（IoT）设备在个人健康护理和智能家居等多种应用中被广泛使用，其核心组件之一是微控制器单元（MCU），这些设备的计算资源（包括SRAM和Flash）非常有限。例如，STM32F412的最大支持为256KB SRAM和1MB Flash。因此，将流行的高效模型部署到IoT设备上面临挑战，尤其是在MCU上实现高准确率的模型仍然是一个难题。近期的研究开始关注在MCU的特殊约束下进行模型设计。例如，MCUNet提出了一个系统-算法协同设计框架和优化的分辨率-宽度搜索空间来搜索网络架构，而MCUNetV2则提出了一个基于补丁的推理框架来解决MCU的峰值内存限制。然而，这些方法仍然依赖于高成本的NAS方法，并且无法充分利用MCU上的SRAM或Flash资源，限制了更高效和有效架构设计的可能性。因此，本文提出了一种充分利用硬件潜力的硬件导向模型设计框架，通过深入分析每个单独的算子（如卷积和ReLU），收集大量运行时数据并进行深入分析，从而构建AtomDB数据库。基于AtomDB提供的算子级见解，本文提出了三条优化指南，用于在Flash、SRAM和延迟方面优化MCU友好的神经网络，并设计出AtomNet架构。
相关工作
本文回顾了近年来在模型压缩技术（包括剪枝、量化和网络架构搜索）方面的进展，这些技术旨在优化高效的模型。例如，MobileNets通过使用深度可分离卷积（SEP）减少了普通卷积的计算量。然而，这些方法要么超过了SRAM/Flash的约束，要么由于忽略了MCU的硬件特性而在MCU上未能实现实际加速。与之相对的是，本文的AtomDB和AtomNet通过分析每个算子在MCU上的特性来实现实际加速。
方法
本文提出了一种基于细粒度算子分析的网络设计框架，通过构建系统化的数据库AtomDB来研究Flash、SRAM和延迟问题。该框架的核心是通过测量大量算子在MCU上的表现，提供关于延迟、Flash和SRAM的白盒见解，而不是依赖复杂的NAS方法进行搜索。具体而言，本文首先对单个算子进行深入分析，收集其在真实硬件上的延迟、SRAM和Flash数据。然后，基于这些数据，提出了三条硬件友好的指南，用于指导混合模块模式的设计，并最终启发AtomNet的设计。
问题定义
本文的目标是在给定的延迟、SRAM和Flash上限下，设计出具有最大准确率的模型。为了实现这一目标，本文强调了在MCU上进行实际测量的重要性，因为理论估计可能会与实际测量存在偏差，从而阻碍所选架构的实际部署。
框架概述
与大多数现有工作不同，本文不是将硬件信息作为黑盒奖励来解决问题，而是通过在MCU上测量大量算子，提供关于延迟、Flash和SRAM的白盒见解。通过这种方式，可以更好地构建整个网络，以最大化利用MCU的资源，从而获得更高的准确率。
算子分析与AtomDB构建
AtomDB通过测量大量算子在MCU上的表现来构建，这些算子是推理引擎在融合和优化后实际执行的（例如ConvBNRelu）。通过收集这些算子的延迟、SRAM和Flash数据，AtomDB为架构设计提供了基础支持。本文提出了三条基于AtomDB的指南，分别从Flash、SRAM和延迟的角度优化MCU友好的神经网络。
混合模块模式
通过对模块的分析，本文发现没有任何一种模块在延迟、SRAM和Flash方面表现最佳，因此提出了一种混合模块模式，以充分利用MCU的资源。具体而言，在网络的早期阶段，由于分辨率较大，建议使用SEP模块（不带SE模块），因为它对SRAM的需求最少；在网络的后期阶段，由于Flash限制更为关键，建议使用ConvNext模块，因为它对Flash最为友好。通过这种方式，混合模块模式能够更好地处理不同阶段的资源瓶颈。
AtomNet设计
基于上述分析，本文提出了一个硬件导向的搜索空间，并采用MACs（乘积累加运算）作为准确率的指标，搜索一系列能够最大化利用硬件资源的AtomNet模型。与现有的NAS方法不同，本文采用零成本NAS方法进行搜索，以简化模型搜索过程。
实验
实验设置
本文在ImageNet数据集上使用SGD优化器进行训练，并在VWW数据集上对预训练的网络进行微调。对于PASCAL VOC数据集，使用AdamW优化器进行训练。在部署时，将AtomNet量化为INT8格式，并使用STM32 X-CUBE-AI推理库在不同的硬件上进行部署，包括STM32F412（256KB SRAM/1MB Flash）、STM32F746（320KB SRAM/1MB Flash）和STM32H743（512KB SRAM/2MB Flash）。
图像分类任务
AtomNet在ImageNet任务上取得了显著的性能提升。例如，在STM32F412（256KB SRAM/1MB Flash）上，AtomNet实现了66.7%的Top-1准确率，比当前最先进的方法MCUNetV2高出1.5%。在STM32H743（512KB SRAM/2MB Flash）上，AtomNet实现了73.0%的Top-1准确率，比MCUNetV2高出1.2%。此外，AtomNet在延迟-准确率权衡方面也表现出色，与MCUNet系列和其他现有方法相比，在所有硬件平台上均实现了显著的准确率提升。
VWW和目标检测任务
在VWW任务中，AtomNet在10FPS、5FPS和320KB SRAM设置下均优于MCUNet，最大模型的准确率达到了95.0%，超过了MCUNetV2的94.0%。在PASCAL VOC数据集上的目标检测任务中，AtomNet在256KB和512KB SRAM限制下分别比MCUNetV2提高了0.7%和1.6%的mAP。
消融研究
指南的有效性
本文通过对比实验验证了三条指南的有效性。例如，在Flash方面，基于深度可分离卷积（DW）的网络显著减少了Flash使用量，并提高了准确率。在SRAM方面，通过在SRAM主导阶段替换模块，成功减少了SRAM的使用。在延迟方面，通过更多地使用ConvNext模块，AtomNet在满足SRAM和Flash约束的同时，显著降低了延迟。
混合模块模式的有效性
通过对比实验，本文证明了混合模块模式在资源受限的MCU上的优势。与仅使用单一模块模式的模型相比，AtomNet在相同的资源约束下实现了更高的准确率。此外，通过可视化SRAM使用情况，本文展示了混合模块模式在平衡跨层SRAM使用方面的优势。
讨论
本文的方法专注于在有限的SRAM和Flash的MCU上进行硬件导向的模型设计，这与常见的NAS算法是正交的。即使使用简单的零成本指标，本文的方法也能显著优于现有的MCUNet系列。此外，如果结合先进的NAS方法（例如Real et al. 2019），还可以进一步降低延迟并提升性能。然而，由于其高成本，本文并未将其作为默认设置。实验结果表明，通过对硬件的深入分析和MCU友好的设计模式，可以在不需要复杂NAS方法的情况下搜索到性能更优的模型。
结论
本文通过构建AtomDB和基于算子级指南分析模块，提出了一种在极端MCU资源约束下的混合模块模式。通过简单的硬件导向搜索策略，AtomNet进一步提升了硬件利用率和计算效率。该方法在TinyML设置下为分类任务和目标检测任务的准确率性能设定了新的标杆。