Due to the high price and heavy energy consumption of GPUs, deploying deep models on IoT devices such as microcontrollers makes significant contributions for ecological AI. Conventional methods successfully enable convolutional neural network inference of high resolution images on microcontrollers, while the framework for vision transformers that achieve the state-of-the-art performance in many vision applications still remains unexplored. In this paper, we propose a hardware-algorithm co-optimizations method called MCUFormer to deploy vision transformers on microcontrollers with extremely limited memory, where we jointly design transformer architecture and construct the inference operator library to fit the memory resource constraint. More specifically, we generalize the one-shot network architecture search (NAS) to discover the optimal architecture with highest task performance given the memory budget from the microcontrollers, where we enlarge the existing search space of vision transformers by considering the low-rank decomposition dimensions and patch resolution for memory reduction. For the construction of the inference operator library of vision transformers, we schedule the memory buffer during inference through operator integration, patch embedding decomposition, and token overwriting, allowing the memory buffer to be fully utilized to adapt to the forward pass of the vision transformer. Experimental results demonstrate that our MCUFormer achieves 73.62\% top-1 accuracy on ImageNet for image classification with 320KB memory on STM32F746 microcontroller. Code is available at https://github.com/liangyn22/MCUFormer.

翻译：由于GPU价格高昂且能耗巨大，在微控制器等物联网设备上部署深度模型对生态AI具有重要贡献。传统方法已成功实现高分辨率图像在微控制器上的卷积神经网络推理，但针对在众多视觉应用中取得最先进性能的视觉Transformer框架仍未被探索。本文提出一种名为MCUFormer的硬件-算法协同优化方法，在内存极受限的微控制器上部署视觉Transformer，通过联合设计Transformer架构与构建推理算子库来适配内存资源约束。具体而言，我们推广一次性神经架构搜索（NAS）以发现满足微控制器内存预算下任务性能最优的架构，通过引入低秩分解维度和补丁分辨率对现有视觉Transformer搜索空间进行扩展，从而降低内存消耗。在视觉Transformer推理算子库构建方面，我们通过算子融合、补丁嵌入分解和令牌覆盖机制调度推理过程中的内存缓冲区，实现内存缓冲区的高效利用以适配视觉Transformer的前向传播。实验结果表明，在STM32F746微控制器仅320KB内存条件下，我们的MCUFormer在ImageNet图像分类任务中达到73.62%的Top-1准确率。代码开源地址：https://github.com/liangyn22/MCUFormer。