基于MCU的轻量级AI推理：TFLite Micro移植与模型量化实操

MCU端跑AI推理（如STM32G4、ESP32-S3、CH32V307等）的主流方案是Google的TensorFlow Lite Micro（TFLite Micro）——它能在无OS、无动态内存分配的环境下运行，模型以C数组形式编译进固件。但浮点模型直接塞进MCU既不现实（体积大、无FPU时速度慢），所以要走"训练后全整数量化（int8）→C数组→TFLite Micro解释器"的标准路径。下面以关键词识别（KWS）模型部署到STM32G431（170MHz Cortex-M4，32KB SRAM）为例。

一、PC端：训练后全整数量化

浮点模型转int8是部署的必经之路——模型体积缩4倍，无FPU的MCU上推理速度提升3倍以上。关键是提供representative_dataset做校准，让转换器估算激活值范围：

import tensorflow as tf

import numpy as np

# 从训练集取100个样本做校准

def representative_dataset():

   for _ in range(100):

       data = np.random.randn(1, 49, 40, 1).astype(np.float32)  # 假设MFCC特征

       yield [data]

converter = tf.lite.TFLiteConverter.from_keras_model(kws_model)

converter.optimizations = [tf.lite.Optimize.DEFAULT]

converter.representative_dataset = representative_dataset

converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

converter.inference_input_type = tf.int8

converter.inference_output_type = tf.int8

tflite_quant_model = converter.convert()

# 转C数组

with open('kws_quant.tflite', 'wb') as f:

   f.write(tflite_quant_model)

xxd -i kws_quant.tflite > model_data.cc

这一步产出的model_data.cc包含了unsigned char kws_quant_tflite[]和长度变量，直接加入MCU工程编译。

⚠ 量化后精度通常会降1-3%，若降幅过大说明模型对量化敏感，可考虑量化感知训练（QAT）或简化模型结构。校准数据集必须是真实场景数据，不能用随机噪声糊弄。

二、MCU端：TFLite Micro解释器移植

TFLite Micro的核心对象MicroInterpreter，需要三样东西：模型C数组、OpResolver（算子解析器）、tensor_arena（静态内存池）。

#include "tensorflow/lite/micro/all_ops_resolver.h"

#include "tensorflow/lite/micro/micro_interpreter.h"

#include "model_data.h"  // xxd生成的C数组

// 仅注册模型实际使用的算子，减小二进制体积

static tflite::MicroMutableOpResolver<6> resolver;

resolver.AddConv2D();

resolver.AddDepthwiseConv2D();

resolver.AddFullyConnected();

resolver.AddSoftmax();

// tensor_arena：静态分配，大小从80KB起步逐步调

constexpr int kTensorArenaSize = 80 * 1024;

uint8_t tensor_arena[kTensorArenaSize] __attribute__((aligned(16)));

tflite::MicroInterpreter interpreter(

   tflite::GetModel(kws_quant_tflite),

   resolver,

   tensor_arena,

   kTensorArenaSize

);

// 分配张量

if (interpreter.AllocateTensors() != kTfLiteOk) {

   // 处理arena不足

}

// 推理循环

void run_inference(int8_t* mfcc_features) {

   // 1. 填充输入（int8域直接计算，无需反量化）

   memcpy(interpreter.input(0)->data.int8,

          mfcc_features, input_size);

   // 2. 执行推理

   interpreter.Invoke();

   // 3. 取输出：直接在int8域argmax

   int8_t* output = interpreter.output(0)->data.int8;

   int8_t max_val = -128;

   uint8_t class_id = 0;

   for (int i = 0; i < num_classes; i++) {

       if (output[i] > max_val) {

           max_val = output[i];

           class_id = i;

       }

   }

   // class_id即为识别结果

}

三个内存相关的关键点：

tensor_arena必须对齐到16字节，否则AllocateTensors()可能assert失败

arena大小：从80KB起步，运行时若返回kTfLiteError说明不足，按10KB步进增加直至稳定，再加10%余量

模型放Flash、arena放SRAM：模型权重只读且大，XIP执行；激活值频繁读写，必须在SRAM

三、CMSIS-NN加速：Cortex-M的必备优化

在Cortex-M4/M7/M55上，TFLite Micro可调用ARM的CMSIS-NN库替换默认算子实现，利用DSP/SIMD指令加速int8卷积、全连接等核心算子，推理速度通常可获数倍提升。

启用方式：

编译选项：开启-O3或-Ofast，启用-mcpu=cortex-m4 -mfloat-abi=hard -mfpu=fpv4-sp-d16 -DARM_MATH_DSP

链接CMSIS-NN库：OPTIMIZED_KERNEL_DIR=cmsis_nn

算子替换：通过自定义OpResolver，把CONV_2D、DEPTHWISE_CONV_2D、FULLY_CONNECTED指向CMSIS-NN版本

class CmsisNnOpsResolver : public tflite::AllOpsResolver {

public:

   CmsisNnOpsResolver() {

       ReplaceOp(tflite::BuiltinOperator_CONV_2D,

                 Register_CONV_2D_CMSIS_NN());

       ReplaceOp(tflite::BuiltinOperator_DEPTHWISE_CONV_2D,

                 Register_DEPTHWISE_CONV_2D_CMSIS_NN());

   }

};

未启用CMSIS-NN时，TFLite Micro使用纯C的参考实现，卷积层可能退化成嵌套for循环；启用后调用arm_convolve_s8()等汇编优化函数，一条SMLAD指令并行处理两个乘加。

四、三个翻车高发点

arena估算偏小：新手按模型大小×2预估arena，忽略了中间激活值开销。正确做法是先用80KB起步，运行时若Alloc失败再逐步加。2层1D-CNN约需48KB arena，加上业务buffers，STM32G4的32KB SRAM其实捉襟见肘，需升级到STM32F4/F7或外扩SRAM

算子未全部注册：模型用到了SPLIT、PAD等TFLite Micro未实现的算子，转换时就会报错"support only integer-only hardware"。解决办法：简化模型结构，或在训练时避免非常规算子

校准数据集与真实输入分布不符：用ImageNet预训练权重做迁移学习时，若校准数据仍是ImageNet而非目标场景，量化后精度断崖。校准集必须来自真实部署场景（如本设备的麦克风采集数据）