4、TensorFlow Lite模型训练基础：使用Keras构建简单分类模型、模型训练与评估、模型导出为.tflite格式

在将TensorFlow Lite部署到STM32微控制器之前，我们必须首先解决一个根本性问题：如何获得一个既满足任务精度要求、又能在资源受限的MCU上高效运行的模型。本章将聚焦于模型的生产源头——使用Keras构建并训练一个简单的图像分类模型，并完成向.tflite格式的转换。这一过程是后续所有量化、编译与硬件移植工作的前提。我们将以经典的MNIST手写数字识别为例，完整走通从数据准备、模型构建、训练评估到格式导出的全链路，并重点剖析为嵌入式部署而进行的模型设计考量。

Keras模型训练到TFLite导出工程链路 1. 数据准备 MNIST加载/归一化 (28x28灰度, 0-1) 2. Keras模型 Flatten + Dense(128) + Dense(10, softmax) 3. 训练与评估 Adam优化器 Accuracy/Val_accuracy 4. 导出TFLite Converter .tflite文件 STM32部署约束（模型设计必须考虑） RAM: 通常 64KB - 512KB Flash: 通常 256KB - 2MB 算力: 无FPU或单精度FPU 功耗: 电池供电场景 算子支持: TFLite Micro内置算子 量化: 需INT8/FP16支持 关键决策路径 ① 模型复杂度 vs 精度：Dense层数越多，参数量越大，Flash占用越高 ② 输入尺寸：28x28=784浮点运算，若为96x96则运算量增加约11.7倍 ③ 激活函数：ReLU比sigmoid/tanh计算更省，且利于量化

上图展示了从Keras模型到TFLite文件的完整工程链路，以及STM32部署对模型设计施加的硬性约束。在实际项目中，我们往往需要根据目标MCU的Flash和RAM容量，反向调整模型结构。例如，一个具有128个神经元的隐藏层，其权重参数为784×128 + 128（偏置） = 100,480个浮点数，以FP32存储需约393KB，这已经超出了许多入门级STM32的Flash容量。因此，模型小型化与量化是嵌入式AI落地的关键。

4.1 核心原理：为什么选择Keras与TFLite的组合

Keras作为TensorFlow 2.x的高级API，提供了简洁的模型构建语法，其顺序模型（Sequential）和函数式API能够快速搭建从全连接网络到卷积神经网络的各类结构。对于嵌入式部署，Keras模型需要经过TensorFlow Lite Converter转换为FlatBuffer格式的.tflite文件。这一转换过程不仅仅是格式变化，更包含了算子融合、常量折叠等优化，以及后续可选的量化操作。

理解这一流程的核心在于认识到：训练时的模型（浮点精度）与部署时的模型（可能为整数精度）在数值表现上存在差异。训练过程使用32位浮点数进行反向传播，而部署到STM32时，为了节省内存和加速推理，通常会将权重和激活值量化为8位整数。这种量化带来的精度损失需要在训练阶段就有所预估，例如通过伪量化训练（Quantization-aware Training）来缓解。

核心概念对照表

概念训练阶段（Keras）部署阶段（TFLite Micro）
数据类型FP32（float32）INT8 / FP16 / FP32
权重存储约393KB（128神经元全连接层）约98KB（INT8量化后）
激活函数ReLU, Softmax等需TFLite Micro支持
算子支持所有TF算子仅TFLite内置算子子集
数值精度高精度，适合反向传播有限精度，适合前向推理

4.2 操作流程：从数据到.tflite的完整步骤

下面我们以MNIST数据集为例，演示完整的模型训练与导出流程。MNIST包含60,000张28×28的灰度手写数字图像，共10个类别（0-9）。该数据集规模适中，非常适合验证嵌入式部署流程。

步骤一：数据加载与预处理

import tensorflow as tf
from tensorflow.keras import layers, models

# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

# 归一化：将像素值从[0,255]缩放到[0,1]
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0

# 增加通道维度：从(28,28)变为(28,28,1)
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# 将标签转换为one-hot编码（可选，sparse_categorical_crossentropy可不用）
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)

print(f"训练集形状: {x_train.shape}, 测试集形状: {x_test.shape}")

这里的关键操作是归一化和维度调整。归一化将输入范围映射到[0,1]，这有利于梯度下降的收敛。增加通道维度是为了匹配Conv2D或Flatten层的输入要求。对于嵌入式部署，我们还可以考虑将输入尺寸进一步缩小（如从28×28降至14×14），但这会损失信息，需要根据实际任务权衡。

步骤二：构建Keras模型

model = models.Sequential([
    layers.Input(shape=(28, 28, 1)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

model.summary()

该模型包含一个Flatten层将784个像素展平，一个128神经元的全连接隐藏层，以及一个10神经元的输出层。总参数量为784×128 + 128 + 128×10 + 10 = 101,770个参数。以FP32存储约需398KB，INT8量化后约100KB。对于具有512KB Flash的STM32F4系列，这个模型是可行的，但若MCU Flash仅有128KB，则需要考虑减少隐藏层神经元数量或使用深度可分离卷积。

步骤三：模型编译与训练

model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

history = model.fit(x_train, y_train,
                    batch_size=32,
                    epochs=5,
                    validation_split=0.2)

训练过程中，我们使用Adam优化器（学习率默认0.001），批量大小32，训练5个epoch。验证集占20%。训练完成后，我们应评估模型在测试集上的表现：

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print(f"测试集准确率: {test_acc:.4f}")

通常，这个简单模型在MNIST上能达到约97%-98%的测试准确率。如果准确率低于95%，可能需要增加训练轮次或调整模型结构。

步骤四：导出为TensorFlow Lite格式

# 转换模型
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

# 保存为文件
with open('mnist_model.tflite', 'wb') as f:
    f.write(tflite_model)

print(f"模型大小: {len(tflite_model)} bytes")

此时生成的.tflite模型仍为FP32精度。为了适配STM32，我们通常需要进一步进行量化。最简单的量化方式是训练后动态范围量化（Post-training Dynamic Range Quantization），它仅将权重量化为INT8，激活值仍为浮点，推理时动态计算缩放因子。

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

with open('mnist_model_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print(f"量化模型大小: {len(tflite_quant_model)} bytes")

经过动态范围量化后，模型大小通常能减少约4倍。但要注意，这种量化方式在MCU上运行时，激活值仍需要浮点运算单元（FPU），如果STM32没有FPU，推理速度会显著下降。因此，对于Cortex-M0/M0+这类无FPU的内核，我们需要使用全整数量化（Full Integer Quantization），这需要提供代表性数据集来校准激活值的范围。

重要提醒：全整数量化需要额外的校准步骤。如果直接使用动态范围量化，在无FPU的MCU上可能无法运行或速度极慢。务必根据目标芯片的硬件特性选择合适的量化方案。

4.3 案例拆解：针对STM32F401RE的模型优化实战

假设我们使用STM32F401RE开发板（Cortex-M4F，具有FPU，Flash 512KB，RAM 96KB）来部署手写数字识别。初始模型（128隐藏层，FP32）大小为398KB，可以存入Flash，但RAM占用需要计算：输入784×4B=3.1KB，隐藏层激活128×4B=0.5KB，输出10×4B=40B，加上中间变量，总RAM占用约5KB，远低于96KB，因此RAM不是瓶颈。

然而，如果我们希望进一步压缩模型以留出更多Flash空间给其他功能，可以考虑以下优化路径：

优化方案模型大小（FP32）模型大小（INT8）测试准确率Flash节省
基线：128隐藏层398KB100KB97.8%-
方案A：64隐藏层200KB50KB96.9%50%
方案B：32隐藏层101KB26KB95.2%74%
方案C：64+64双隐藏层260KB65KB97.5%35%

从表中可以看出，将隐藏层从128减至64，准确率仅下降0.9%，但模型大小减半。对于手写数字识别这种相对简单的任务，64神经元的隐藏层已经足够。如果任务更复杂（如识别10种不同物体），则可能需要保留更多神经元或改用卷积网络。

在实际部署中，我们还应该考虑使用模型剪枝技术，将接近零的权重直接置零，然后进行稀疏存储。但TFLite Micro对稀疏模型的支持有限，因此更推荐直接调整网络结构。

经验之谈：在嵌入式AI项目中，模型设计应遵循"最小充分"原则。先从一个较小的模型开始，逐步增加复杂度，直到满足精度要求。这比先训练大模型再压缩更高效。

4.4 常见误区与总结复盘

在模型训练与导出过程中，初学者常犯以下错误，这些错误可能导致后续STM32移植失败：

  • 忽略输入数据的归一化方式：训练时使用/255.0归一化，但部署时若忘记在C代码中执行相同操作，会导致推理结果完全错误。必须将归一化操作固化在模型内部（如添加Rescaling层），或确保在MCU端实现相同的预处理。
  • 使用不支持的算子：某些Keras层（如BatchNormalization在训练模式下）在TFLite转换时可能产生不兼容的算子。解决方法是确保模型处于推理模式（training=False），或使用TFLite支持的替代层。
  • 混淆动态范围量化与全整数量化：动态范围量化仅权重为INT8，激活仍为FP32，需要FPU支持。全整数量化则完全使用INT8运算，但需要校准数据集。选择哪种量化方式取决于MCU的算力。
  • 忽视输出后处理：Softmax输出的是概率分布，在MCU上通常不需要计算完整的Softmax（涉及指数运算），而是直接比较logits（未经过Softmax的原始输出）的大小。因此，在转换模型时可以考虑去掉Softmax层，将输出改为线性层，从而减少计算量。

总结本章，我们完成了从Keras模型构建到.tflite导出的完整流程。关键要点包括：模型结构直接影响Flash/RAM占用；量化是嵌入式部署的必经之路；训练与部署环境的数据预处理必须保持一致。在下一阶段，我们将深入探讨如何将生成的.tflite模型转换为C++数组，并集成到STM32工程中，实现真正的边缘推理。

← 上一章 📖 返回目录 下一章 →