边缘AI组内分享
分类: 分享 发布于: 2024-06-26

Tensorflow Lite For Micro Controller(TFLite micro):嵌入式智能的技术革命

在进行分享之前,先展示一个实物Demo。 wio

这是一块seeed出品的MCU开发板,在背面焊接了一颗模拟麦克风来捕捉音频。当说话人讲“yes”的时候,蓝灯将点亮;说“no”的时候,红灯将点亮。

wio

这个语音关键词识别 demo,可能听起来并不稀奇,毕竟我们的手机也能做到。但对于这个装置来讲,有两个数字会彻底改变你对它的认知。

第一,它的功耗不到10毫瓦,一节普通的5号电池就能让它连续工作一个月

第二,这个模型的大小仅有18KB,执行推理时RAM占用仅7KB——相当于一条短信的存储空间,却能在微控制器上实时完成语音识别

这正是TFLM的意义所在:将AI塞进指甲盖大小的芯片里,螺丝壳里做道场。通过极致的优化手段,让廉价、低功耗的设备也能拥有智能。

一、简介

1. 什么是MCU,为什么要在MCU上部署深度学习模型?

MCU(Microcontroller Unit) 是集成了CPU、内存、IO的单片微型计算机,功耗极低(毫瓦甚至微瓦级)、成本低廉(几元到几十元),广泛应用于嵌入式设备中,比如智能家居传感器、穿戴设备、工业控制器等。

此次demo使用的主控芯片 意法半导体STM32-h725AE 是ST公司的高端MCU,基于性能强劲的Arm cortex-M7架构,主要参数如下:

  • 最高主频 550Mhz (此demo降频运行在60Mhz)
  • ROM 512KB
  • RAM 564KB
  • 零售价: 约25元

airpods

在MCU上部署模型有以下优势:

  • 超低功耗
  • 超低成本
  • 实时性与低延迟
  • 隐私与安全

适用场景包括智能音箱(小度小度、小爱同学)、智能门锁的人脸/指纹识别、智能穿戴设备如手表的抬腕亮屏、计步器等,工业领域也有大量应用,例如震动故障检测等。

试想一下,如果这样的技术普及,在复杂度相对低的任务上,未来的智能家居传感器可以完全不用充电,工厂里的设备能自主预警故障,甚至医疗植入器件能本地分析数据——而这一切,无需依赖云端;

另一方面,在大模型能力日趋强大的时代,通过在微控制器上处理实时传感数据(如语音唤醒、异常检测),作为感知的第一只触手,来唤醒云端大模型进行深度决策(如自然语言理解、图像生成)—— 既保障了隐私与实时性,又解锁了复杂智能,让‘边缘轻量化+云端强脑力’优势互补,通过“终端感知过滤,云端深度思考”的模式,达到算力与能耗的最优平衡。

2. TFLM是什么?

资源极度受限的微控制器 (MCU) 设计的__轻量级机器学习推理框架__,是 TensorFlow Lite 的嵌入式子集。核心目标是在 ≤256KB RAM 的设备上运行神经网络模型,突破传统深度学习框架的硬件边界。

与传统框架对比:

特性 TensorFlow TensorFlow Lite (TFLite) TFLite Micro (TFLM)
运行设备 服务器/PC 手机/边缘设备 MCU(无OS)
内存需求 GB级 MB级 KB级
依赖项 Python/C++运行时 部分系统库 零依赖(裸机)
模型体积 100+ MB 1-10 MB 10-100 KB

TFLM是为资源极端受限的嵌入式设备(如微控制器)设计的轻量级机器学习推理框架,其优化手段涵盖模型压缩、内存管理、算子优化和硬件适配等多个层面,架构图 arc

核心组件解析:

  • Interpreter​​: 推理引擎,负责模型加载、图解析、算子调度。主要实现了.tflite 模型格式的加载与解析,仅需单块连续内存(tensor_arena)存储模型权重、激活值和中间张量,所有内存通过预分配的 Tensor Arena 静态管理,避免碎片化问题
  • Operator Kernels 算子调度与扩展​,通过 ​​OpResolver​​ 动态注册算子实现(如卷积、全连接层),仅加载模型所需的算子内核,减少代码体积。提供通用实现、硬件加速库(如 ARM CMSIS-NN)的集成,以及第三方适配API,方便用户实现自己的算子,或者针对特殊架构如risc-V等保留扩展性
  • Memory Manager TFLM的核心基础设施,负责内存管理。使用静态内存分配和内存池(Tensor Arena),借助张量内存复用、内存地址对齐等手段,完全摒弃动态内存分配,提升内存利用效率、减少碎片

可以通过一个实例,了解训练到部署的全流程:

二、模型训练&部署

Demo源自TensorFlow官方代码库中的Speech Commands项目。训练数据采用开源数据集《Speech Commands Dataset》,包含 35 个词的短音频(1秒长度,16bit@16kHz 采样率),端到端链路: mfcc flow

1. 预处理

首先通过以下手段将原始音频流做数据增强:

  • 添加背景噪声(随机 SNR)
  • 时间偏移(Time Shift)
  • 音量扰动(Volume Jitter)

然后进行特征提取。采用滑动窗口的方式进行采样,每30ms采样一次,前后两个窗口重叠10ms。这样,1s的音频就可以采样49次。针对每一个窗口内的数据,提取40维的MFCC特征 (https://docs.pytorch.org/audio/main/tutorials/audio_feature_extractions_tutorial.html) mfcc flow

[MFCC参考:https://docs.pytorch.org/audio/...]

由此,每1s的时间片段,可以生成一个40*49维的特征图。通过脚本将其转化为热力图,以下为两个不同命令的表现:

mfcc flow

2. 训练

其模型结构采用了一个经过深度优化的__深度可分离卷积神经网络__(Depthwise Separable CNN),专为边缘设备设计。

https://www.tensorflow.org/tutorials/audio/simple_audio?hl=zh-cn micro_speech_quantized

它是二维卷积层(Convolutional 2D)、全连接层和Softmax层组成的简易模型,是典型的tiny_conv架构。模型结构抽象:

全连接层 wio 在全连接层,输出为输入的加权和。全连接层可以将多维特征展平,以适应输出层的维度需求,同时将前一层提取的分散特征整合为高阶语义表示,通常应用于分类决策场景。

卷积层 single channel

各种卷积 conv 卷积层是整个模型的核心,其作用是特征提取。卷积操作的本质,是卷积核与滑动窗口内的局部输入相乘求和,值越大说明匹配度越高,反之则反。因此经过训练的卷积核,会跟与之相匹配的模式输入“强烈反应”,以此机制来提取特征。

卷积层中,输出的神经元不再与每个输入相连接,而是只与receptive field内的输入建立联系,在很多资料中被翻译为感受野,它是类似于“视野”的一个概念。

每个卷积核专注于一种局部模式,如边缘、颜色突变、语音频段特征,由多个卷积核共同捕捉输入中的有效模式,并传递到下一层或深层核中组合出高级特征甚至语义特征。 conv_cat

与全连接层相比,卷积层通过共享权重的方式,即同一组卷积核在输入上滑动复用,大大减少模型参数,减少了内存交换,有很高的内存效率;同时其并行特性适合采用硬件加速。

标准卷积在特定场景下存在一些局限,为了解决这些局限,衍生出多种卷积:

  • 步幅卷积:显式降低特征图分辨率、扩大感受野
  • 空洞卷积:扩大感受野
  • 转置卷积:扩大特征图分辨率
  • 分组卷积:降低计算量和参数
  • 1x1卷积:通道融合、升降维

深度可分离卷积(Depthwise Separable CNN)就是上述__分组卷积+1x1卷积__的组合,分组卷积把filter分成若干组,将输入通道均分给这些组处理,不同组之间没有融合。深度可分离卷积是分组卷积的极端情况:第一层每组filter只负责1个输入通道,各通道之间彻底没有融合,通过第二层1x1卷积实现通道融合和维度适配。

softmax层

通常用在多分类模型中,用来将神经网络的输出转化成概率分布: softmax 以micro speech为例:

# 假设模型输出的logits为:
[z_silence, z_unknown, z_yes, z_no] = [1.5, 0.8, 3.2, -0.5]

# 计算指数
exp_silence = e^{1.5}  4.4817
exp_unknown = e^{0.8}  2.2255
exp_yes     = e^{3.2}  24.5325
exp_no      = e^{-0.5}  0.6065

# 计算指数和
sum_exp = 4.4817 + 2.2255 + 24.5325 + 0.6065  31.8462

# 计算各类别概率
p_silence = 4.4817 / 31.8462  0.1407  # 14.07%
p_unknown = 2.2255 / 31.8462  0.0699  # 6.99%
p_yes     = 24.5325 / 31.8462  0.7703  # 77.03%
p_no      = 0.6065 / 31.8462  0.0190  # 1.90%

基于以上模型架构进行训练:

  • 官方训练脚本:tensorflow/tensorflow/examples/speech_commands/train.py
  • 参数配置:tensorflow/tensorflow/examples/speech_commands/

脚本采用动态学习率,起始轮次使用高学习率提高收敛速率,到达一定轮次后衰减为低学习率避免震荡,由于模型参数量较少(不到2万),使用纯CPU训练预计花费2小时。为了兼容MCU,训练时模拟了 8bit整数量化:

converter = tf.lite.TFLiteConverter.from_saved_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_model = converter.convert()

量化感知训练(Quantization-Aware Training, QAT)

在模型训练过程中模拟量化效果的技术,旨在使模型在后续的整数量化后保持更高的精度,核心思想是在训练阶段模拟量化噪声​​:即在训练前向传播时,向模型权重和激活值注入模拟的量化误差(如四舍五入噪声),使模型适应低精度计算。

但这样会带来一个严重问题,量化操作的不可导,使得权重无法更新。在训练中,模型需要计算loss对权重w的梯度d_loss/dw,即当权重 w 微小变化时,损失函数 loss 的变化率:

  • 正梯度​​:w 增加时 loss 上升 → 该减小 w
  • 负梯度​​:w 增加时 loss 下降 → 该增大 w

梯度通过反向传播修正机制,指导权重的更新。而四舍五入函数 y = round(x) 在非半整数点(如30.2),x微小变化时y不变 → 导数dy/dx = 0;在半整数点时,如x=30.5,从左侧逼近(30.4→30.49...)导致 y始终=30即左导数为0,从右侧逼近(30.51→30.6)造成y突然=31,即右导数几乎无穷大,左右导数不相等​即该点导数不存在。 ​​ QAT通过​​直通估计器(Straight-Through Estimator, STE)​​绕过量化操作的不可导问题,正常更新浮点权重。

QAT是训练边缘模型时常见的做法,通过在训练阶段适应量化后的数值分布,量化截断点落在信息密集区,学习到的权重天然适配低精度计算,尤其是避免低bit量化时(4bit甚至2bit量化)精度崩溃的问题。与QAT类似的技术是PTQ,训练后量化。后者仅依赖校准数据静态补偿,往往在量化后造成较大的精度损失。

3. 模型转换与量化

到训练停止的一步,已经得到了一堆ckpt文件,如下:

ckpt-0001.index
ckpt-0001.data-00000-of-00001
ckpt-0001.meta

ckpt即check point,它不仅仅作为训练过程中的恢复点,也保存了模型训练在某一刻的状态,包括:

  • 模型权重(weights)
  • 优化器状态(如 Adam 优化器的动量值)
  • 训练进度(如当前 epoch 数、全局步数)

训练结束后,可选择验证集性能最优的 ckpt 导出原始模型。至此还不能直接部署到MCU,需要经过一系列转换和优化:

softmax

其中,上图虽然指出.h5或.pb的原始模型,TFLM的支持并不仅限于此,它兼容多种格式,由转换器负责处理不同模型实现的细节,包括:

  • TensorFlow 生态格式​​
  1. .ckpt(Checkpoint)​​:.meta图结构、.index参数索引、.data权重值、checkpoint记录路径构成,训练中保存模型状态,支持断点续训和参数恢复
  2. ​.pb(Protocol Buffers)​​: 包含两种形式:①​​Frozen Graph​​(单文件):固化计算图与权重,仅支持推理;② SavedModel​​(目录):含模型结构、权重、签名(如输入/输出张量名),适合部署;支持跨语言部署(C++/Java调用)、支持TensorFlow Serving
  3. ​​.h5(HDF5)​​特点​​:单文件包含模型架构+权重+优化器状态,适合快速实验,不支持自定义层,跨框架兼容性差
  • PyTorch 生态格式​​
  1. .pt/.pth​​: 保存方式​​为​​权重字典​​(state_dict):仅参数,需代码重建模型。.pt文件仅保存模型权重(state_dict),不包含网络结构。加载时通常需​​先通过config.json重建模型架构​​。
  • 跨框架交换格式​​
  1. ​​.onnx(Open Neural Network Exchange)​​​​基于计算图(节点=算子、边=张量流),二进制存储,各框架都有完善工具链,包括转换器torch.onnx.export();TF → tf2onnx,​推理加速​​:ONNX Runtime(CPU/GPU)、TensorRT(NVIDIA GPU优化)

TFLM对其他平台的模型不提供直接支持,但可以通过工具将onnx转换为支持的SavedModel,进而导出为.tflite模型。关键API:

# SavedModel → TFLite
converter = tf.lite.TFLiteConverter.from_saved_model("saved_model")
tflite_model = converter.convert()
open("model.tflite", "wb").write(tflite_model)

# 量化
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # 默认量化
tflite_model = converter.convert()

量化关键步骤​​:

  • 模型冻结(去除训练相关节点)
  • 算子兼容性检查
  • 量化校准(需提供representative_dataset)

其中关键在于量化校准过程,通过校准数据,用于确定浮点数值到整数的映射范围(动态范围计算)​​,通过对校准数据执行推理,来记录记录每层激活值范围,进而统计最大值/最小值、计算缩放因子scale=(max-min)/255

从而将FP32的无限范围映射到INT8的[-128,127]。

5. 后处理

经过上述训练,得到了一个量化后的模型,它接收从1ms的语音片段抽取的MFCC特征图(40*49),执行一次推理后输出四分类的概率。

模型在实时推理过程中,处理的是流式数据,通常通过滑动窗口进行离散化处理。这里采用100ms滑动窗口间隔,即前后两次推理保留900ms的重叠,这样能够尽可能避免关键语素处在窗口边缘,导致有效特征丢失的问题。同样由于模型推理的抖动性,通常需要滤波手段来滤除单次推理错误和噪声干扰。

常见的手段包括:

  • 时间窗加权投票法
  • 基于统计的滤波(中值滤波、均值滤波、熵值滤波​等)
  • 基于状态机的滤波(隐马尔可夫模型 (HMM)等)

受限于MCU的资源限制,采取计算量小的时间窗加权投票法: softmax

三、TFLM核心原理

在上一章节,我们训练了一个模型,在部署全流程中与TFLM的交互以及使用的API如下: arc

1. .tflite模型结构与内存管理

.tflite模型模型定义了计算图和参数,其结构如下图所示:

tflite_struct.png

其由FlatBuffers定义,包含三层结构:

  • 操作符表(Operator Codes)​​:内置或自定义算子(如Conv2D、Add)
  • 子图(SubGraphs)​​:包含输入/输出张量、算子执行序列,支持多子图并行
  • 缓冲区(Buffers)​​:存储权重、常量等裸数据,与张量分离以支持延迟加载

在解释器的加载和执行过程中,经历以下阶段:

  • 创建解释器:使用 tflite::MicroInterpreter 类创建一个解释器实例。解释器需要模型数据、操作注册表、张量分配器和内存缓冲区作为输入。
  • 解析模型:解释器会解析传入的模型数据,确保模型格式正确,并提取模型的基本信息,如输入和输出张量的数量和类型。
  • 初始化Op解析器:解释器会使用传入的Op解析器(如 tflite::AllOpsResolver)来解析模型中使用的所有Op。Op解析器包含了所有可能的操作实现,解释器会根据模型中的Op类型选择相应的实现。
  • 创建张量分配器:解释器会使用传入的张量分配器(如 tflite::MicroAllocator)来管理内存分配。张量分配器负责在给定的内存缓冲区中分配和管理张量所需的内存。
  • 初始化张量:解释器会初始化模型中的所有张量,包括输入张量、输出张量和中间张量。初始化过程包括设置张量的类型、形状和内存位置。
  • 分配内存张量:通过AllocateTensors()触发,核心目标是​​高效规划与分配张量内存
  • 推理

受限于嵌入设备的内存限制,内存的分配是一个TFLM的重点,所有操作基于一个全局的Arena分配器​​:

__ 统一内存池​​__

tensor_arena预先分配一块​​连续物理内存​​(如SRAM)推理过程需要的内存均由TFLM进行管理,不仅不再需要动态内存分配相关的依赖,还避免动态内存分配造成的碎片问题,同时规避实时系统中,内存分配的不确定性。 分配示例:

constexpr size_t kArenaSize = 1024*8;  // 8KB
__attribute__((section(".dtcm_ram"))) alignas(16) uint8_t g_arena[kArenaSize];

其内存池划分为三个逻辑区域:

  • ​​HEAD区​​:存储​​非持久性张量​​(如中间激活值),通过内存规划器动态复用空间
  • TEMPORARY区​​:短期临时缓冲区(如算子内部计算缓存),执行后立即释放。
  • TAIL区​​:存储​​持久性数据​​(如权重、常量张量),生命周期覆盖整个模型推理过程。

memory

__ 内存规划器__

TFLM 提供了两种主要的内存规划器:

  • LinearMemoryPlanner (线性分配器)

    使用一种简单的线性分配策略。它按照内存块请求的顺序分配内存,所有内存块在内存中线性排列,因而内存不存在复用,整理利用率可能偏低。但其优点在于分配速度快,在单一路径的简单模型如线性链式结构上表现良好,实时性高,在某些资源极度受限的设备上此分配器是最优解,因为可能无法承担贪心算法的运行时开销。

  • ​​GreedyMemoryPlanner(贪心算法分配器)​​

    动态分析张量生命周期,最大化重叠非活跃张量的内存空间达成复用目的,如早期激活值与后续计算复用。其可以通过​​__离线分配(offline allocation)__的手段进一步降低运行时开销,即在模型编译或转换阶段预先确定内存块的偏移量。缺点是内存布局固定,无法在运行时调整内存使用。

    同时,GMP也提供__在线分配(online allocation)__方法,解决内存布局无法调整的问题,但其会引入分配开销,也可能会导致内存碎片从而影响内存利用率。

__ FlatBuffer __

FlatBuffer是google开发的序列化库,但其与传统的序列化库的工作方式截然不同。FlatBuffers的核心特点是“零解析”:预编译的内存布局设计​​,使序列化后的二进制数据在内存中的存储格式与程序中的数据结构完全一致,无需额外解析。其两个重要特性:

  • ​​预定义 Schema 驱动​​: 用户需预先通过 Schema 定义数据结构(如 table Person { name:string; age:int; })。FlatBuffers 编译器(flatc)根据 Schema 生成目标语言(如 C++/Java)的​​强类型访问接口​​,这些接口直接映射二进制数据的偏移量。
  • ​​直接内存访问​​: 序列化后的数据保存在​​连续的字节缓冲区​​中。程序通过生成的访问接口(如 person->name())直接根据​​偏移量定位字段位置​​,跳过了解析、拷贝过程

零解析的关键依赖于两项核心技术: 一是__vTable(虚拟表)机制__​​,即每个 table 类型对象关联一个 vTable(存储各字段的偏移地址),读取字段时:字段值地址 = 对象基址 + vTable[字段偏移]。二是__​​分层寻址设计​​__,针对不同的数据类型设计不同的寻址模式,如:①​​标量/Struct​​:直接内联在父对象中,无间接访问;②​​String/Vector​​:存储指向实际数据的指针(偏移量),二次跳转;③​​Union​​:先读取类型标签,再按类型跳转到对应数据。

除了上述特性,FlatBuffer还严格采用的小端序(Little-Endian)存储 + 字节对齐规则,确保相同二进制数据在不同硬件平台一致解析,达成跨平台一致性。FlatBuffer 在强化内存利用率的同时,也引入一些弊端,比如数据不可变,如需修改某字段则需要完整重建整个buffer,是读写平衡的极端应用场景。

2. 模型与计算图优化

量化​

将FP32权重和激活值转换为INT8格式,减少75%的模型大小和内存占用。量化通过缩放因子(scale)和零点(zero_point)映射浮点范围到整数范围,反量化则通过逆操作恢复近似浮点值。 ​​ 执行int8量化理论上相比FP32模型降低3/4的内存空间。针对ARM Cortex-M设备,INT8乘加指令比FP32快4倍,由此理论上也会加速4倍推理速度。

部分高端MCU内置了NPU(如stm32N6),中低端型号通常也内置了DSP或通过指令集支持SIMD计算,这些硬件通常只支持int8,量化后模型在硬件加速下可以获得大幅性能提升。

*混合精度量化

旨在通过动态分配量化位宽(如FP32、FP16、INT8等)来平衡模型的计算效率、内存占用与推理精度。其核心思想是根据模型各层对精度的敏感度差异,自适应地选择最优的量化策略,从而在硬件资源受限的场景下实现高效部署。

算子融合(Op Fusion)

将相邻算子合并为单一内核(如 Conv2D + ReLU => FusedConv2DWithReLU),减少中间张量存储和调度开销。避免中间结果的存储和加载,减少内核调用次数,更好地利用缓存和并行指令。融合可以在模型转换阶段,也可以在算子注册时由Interpreter选择Op解析器时选择相应的实现,其包括以下重要步骤:

  • 计算图重写(Graph Rewriting):
  • 模式匹配:转换器识别计算图中的特定算子序列(如 Conv2D → ReLU)
  • 节点替换:将匹配的序列替换为融合后的算子节点(如 FusedConv2DWithReLU)。
  • 内存访问优化:
  • 中间结果跳过存储:融合后的算子直接传递数据到下一计算阶段,无需写入全局内存。
  • 硬件指令级优化: 在 ARM Cortex-M 上,融合算子可结合 SMLAL(乘加)和 SMAX(最大值)指令,避免冗余数据加载(MatMul+Add => FusedMatMulWithAdd)
; ARM Cortex-M 汇编示例(伪代码)
SMLAL r0, r1, r2, r3  ; Conv2D 计算
SMAX  r0, r0, #0       ; 立即执行 ReLU

剪枝(Pruning)

剪枝(Pruning)是深度学习中一种重要的模型压缩技术,其通过移除神经网络中冗余或重要程度相对低的参数(如权重、神经元、通道等),在尽量保持模型性能的前提下,显著降低模型的复杂度和资源消耗。TFLite模型转换器在启用optimizations=tf.lite.Optimize.EXPERIMENTAL_SPARSITY可保留稀疏结构,确保.tflite文件包含稀疏张量元数据,并在推理过程中结合硬件支持,自动调用稀疏计算内核。

3. 硬件加速

CMSIS-NN​​(Cortex Microcontroller Software Interface Standard - Neural Network)

https://arm-software.github.io/CMSIS-NN/latest/

ARM专为Cortex-M系列微控制器设计的神经网络加速库。其核心目标是通过优化计算内核,​​最大化推理性能​​并​​最小化内存占用​​,使AI模型能在资源受限的嵌入式设备(如STM32、瑞萨RA系列)高效运行。

TFLM可以充分利用CMSIS-NN的加速能力,在创建工程时可以选择相关支持:

python tensorflow/lite/micro/tools/project_generation/create_tflm_tree.py  --makefile_options="TARGET=cortex_m_generic TARGET_ARCH=cortex-m7 OPTIMIZED_KERNEL_DIR=cmsis_nn" --example="micro_speech" ../gen_micro_speech_nn
处理器架构 核心特性 CMSIS-NN 支持程度 CMSIS-NN依赖的硬件支持
Cortex-M0/M0+/M3 基础 RISC 内核,无 DSP/FPU 基础支持:纯 C 实现,无 SIMD 优化,依赖软件模拟运算 无专用硬件加速单元;仅依赖 CPU 标量指令
Cortex-M4 支持 DSP 扩展指令,可选 FPU 高度优化:启用 DSP 指令加速(如 SMLAD),并行处理 4 组 INT8 乘加运算 DSP 扩展指令:单周期 MAC、双 16 位乘加;SIMD 并行(需硬件支持)
Cortex-M7 双发射流水线,高主频,支持 Cache 深度优化:DSP 指令 + Cache 局部性优化,减少内存延迟 DSP 扩展指令 + 超标量架构(双发射);指令/数据 Cache(减少内存访问延迟)
Cortex-M55/M85 支持 MVE(Helium 向量指令集) 极致优化:MVE 指令并行处理 16 个 INT8 操作,支持 INT4 量化及稀疏计算 MVE 向量单元(128 位寄存器);硬件循环加速;可选 Ethos-U NPU 协同(芯片级集成)

hd

CMSIS-NN的算法库直接提供卷积、全连接等算子,这些算子通过编译优化层对接底层硬件。

SIMD指令扩展

​​SIMD(单指令多数据)通过单指令同时操作多个数据(如128位寄存器拆分为4个32位浮点数),将计算吞吐量提升​​2-16倍​​(取决于数据类型和指令集)。例如在Arm NEON指令集中,Q0 ~ Q15 为16 个 128 位寄存器可拆分为:

  • ​​4 个 32 位通道​​ → 同时操作 4 个 int32
  • ​​8 个 16 位通道​​ → 同时操作 8 个 int16
  • 16 个 8 位通道​​ → 同时操作 ​​16 个 int8​​

示例:

VADD.F32 Q0, Q1, Q2   ; Q0 = Q1 + Q2(4个float并行加)

; 输入: 
; Q1 = [c0, c1, c2, c3]
; Q2 = [d0, d1, d2, d3]

; 输出
; Q0 = [c0+d0, c1+d1, c2+d2, c3+d3]

其他指令不逐一列举。

DSP指令

DSP指令集是专为高效执行​​滤波、FFT、卷积、编解码​​等信号处理任务设计的硬件加速指令。其核心优势在于​​单周期乘加(MAC)、并行计算(SIMD/VLIW)、专用寻址模式(循环/位反转)和饱和运算​​,可大幅提升实时性敏感任务的性能。

以FIR(有限冲激响应)为例: y[n] = sum (h[k] * x[n - k]), k = 0 ~ N - 1 C语言实现:

// 传统C实现低效
for (int i = 0; i < output_len; i++) {
    y[i] = 0;
    for (int j = 0; j < N; j++) {
        y[i] += h[j] * x[i + j];  // 双层循环+多次内存访问
    }
}

DSP加速:

MOV R0, #x_base      ; 输入信号基地址
MOV R1, #h_base      ; 系数基地址
MOV R2, #N           ; 滤波器阶数
MOV R3, #0           ; 累加器清零

FIR_LOOP:
    LDR R4, [R0], #4 ; 加载x[i+j]并自增地址(自动偏移)
    LDR R5, [R1], #4 ; 加载h[j]并自增地址
    MAC R3, R4, R5    ; 单周期完成乘加:R3 += R4 * R5
    SUBS R2, R2, #1   ; 计数器减1
    BNE FIR_LOOP      ; 循环直到N次完成

L1-Cache

在高端的Arm架构中,CPU因为IPC设计和工艺的水平提高可以达到非常高的主频,而传统的RAM要花费更高的代价才能够匹配CPU的速度。因此Cache用于填补CPU与RAM之间的性能差距​​,利用空间局部性​​(连续地址访问)和​​时间局部性​​(重复访问相同数据),将高频数据缓存在Cache中,规避RAM速率瓶颈。 cache

但在Cache的利用中,有许多关注点,其核心在于充分利用缓存局部性,下面两种循环写法对性能产生较大影响:

int arr[1024][1024];


for (int i = 0; i < 1024; i++) {
    for (int j = 0; j < 1024; j++) {
        arr[i][j] = i + j;
    }
}

for (int j = 0; j < 1024; j++) {
    for (int i = 0; i < 1024; i++) {
        arr[i][j] = i + j;
    }
}

CMSIS-NN的Cache优化策略​:

  • ​​连续内存访问​​:

    采用​​NHWC布局​​(通道维度连续存储),替代传统NCHW。例如,卷积计算时,单点空间位置的多通道数据连续存储,单次Cache加载可处理多个通道,减少跳跃访问导致的Cache Miss

  • 计算分块(Tiling)与Partial Im2Col​

  • ​​Partial Im2Col​​:传统Im2Col将卷积展开为大矩阵,消耗大量内存(如1024×1024输入需生成GB级矩阵)。CMSIS-NN采用​​分块Im2Col​​,仅生成部分列(如4×4小块),确保临时数据适配L1 Cache。

  • ​​循环分块​​:将大卷积拆分为小块计算(例如8×8输入块),使权重和输入数据均能驻留Cache,减少主存访问

  • 内存预取与数据对齐​​

  • ​​预取指令​​:在计算当前块时,预加载下一块数据至Cache,隐藏内存延迟

  • ​​强制对齐​​:数据地址按32字节(Cache Line大小)对齐,确保单次加载完整利用Cache行

总结:作为专为微控制器设计的轻量级推理框架,TFLM通过模型量化、剪枝等压缩手段,配合硬件协同优化,使AI模型能在资源受限的嵌入式设备上高效本地运行。

四、扩展阅读