📚 课程目录第3章模型训练与导出

3、模型训练与导出：使用Keras训练一个简单分类模型，并导出为SavedModel格式

在TFLite Micro的完整工具链中，模型训练与导出是决定最终固件推理效果的第一道关口。很多开发者习惯直接下载预训练模型或使用PyTorch导出ONNX，但在嵌入式AI领域，Keras + TensorFlow SavedModel的组合仍然是与TFLite Converter衔接最平滑、量化支持最成熟的路径。本章将聚焦于如何使用Keras训练一个适用于MCU场景的简单分类模型，并正确导出为SavedModel格式，为后续的.tflite转换与微控制器部署打下坚实基础。

3.1 核心原理：为什么选择Keras + SavedModel作为TFLite Micro的前置链路

TFLite Micro的转换工具链虽然支持多种输入格式，但SavedModel是官方推荐且经过充分测试的中间格式。其核心原因在于：SavedModel不仅保存了模型的网络结构（GraphDef）和权重（Variables），还保存了完整的函数签名（SignatureDef），这使得TFLite Converter能够精确地识别模型的输入输出张量名称、数据类型和形状。相比之下，仅保存权重（.h5）或使用ONNX格式，在转换时往往需要额外指定输入输出节点，增加了出错概率。

对于MCU场景，模型必须满足几个硬性约束：参数量通常在几十KB到几百KB之间（Flash限制），推理时RAM占用不超过几十KB，且算子必须是TFLite Micro内置算子集（如FullyConnected、Conv2D、DepthwiseConv2D、Relu、Softmax等）。Keras的Sequential API天然适合构建这种前馈网络，且其高层抽象让开发者无需关注底层张量操作，可以更专注于网络结构设计。

  关键认知：SavedModel不是简单的“模型文件”，而是一个包含模型图、权重、训练配置和推理签名的目录。TFLite Converter在转换时会解析这个目录中的saved_model.pb文件，提取推理所需的子图，并执行算子映射和量化优化。因此，导出时的任何不规范操作（如未固定输入shape、使用了自定义层）都会在转换阶段暴露出来。

3.2 操作流程：从数据准备到SavedModel导出的完整步骤

下面以一个经典的图像分类任务（Fashion-MNIST，28x28灰度图，10类）为例，演示完整的训练与导出流程。该任务非常适合TFLite Micro教学，因为模型规模小、数据易获取、分类效果直观。

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

import tensorflow as tf
from tensorflow import keras

# 加载Fashion-MNIST数据集
(x_train, y_train), (x_test, y_test) = keras.datasets.fashion_mnist.load_data()

# 归一化到[0,1]区间，并增加通道维度（28,28,1）
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
x_train = x_train[..., tf.newaxis]  # (60000, 28, 28, 1)
x_test = x_test[..., tf.newaxis]    # (10000, 28, 28, 1)

# 划分验证集（从训练集中取5000张）
x_val = x_train[-5000:]
y_val = y_train[-5000:]
x_train = x_train[:-5000]
y_train = y_train[:-5000]

print(f"训练集: {x_train.shape}, 验证集: {x_val.shape}, 测试集: {x_test.shape}")

这里的关键操作是增加通道维度。TFLite Micro的卷积算子要求输入为4D张量（batch, height, width, channels），即使灰度图也需要显式声明通道数为1。如果省略这一步，后续转换时会报“输入维度不匹配”的错误。

步骤2：构建轻量级模型

model = keras.Sequential([
    keras.layers.Input(shape=(28, 28, 1), name='input_layer'),
    keras.layers.Conv2D(filters=8, kernel_size=(3,3), activation='relu', padding='same'),
    keras.layers.MaxPooling2D(pool_size=(2,2)),
    keras.layers.Conv2D(filters=16, kernel_size=(3,3), activation='relu', padding='same'),
    keras.layers.MaxPooling2D(pool_size=(2,2)),
    keras.layers.Flatten(),
    keras.layers.Dense(units=32, activation='relu'),
    keras.layers.Dense(units=10, activation='softmax', name='output_layer')
])

model.summary()

该模型参数量约为1.2万，模型文件大小约50KB，非常适合MCU部署。注意使用了Input(shape=(28,28,1))显式声明输入层，并给输入输出层命名（name='input_layer'和name='output_layer'），这有助于后续在TFLite转换时明确指定输入输出张量。

步骤3：编译与训练

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

history = model.fit(
    x_train, y_train,
    batch_size=32,
    epochs=10,
    validation_data=(x_val, y_val),
    callbacks=[
        keras.callbacks.EarlyStopping(patience=2, restore_best_weights=True),
        keras.callbacks.ReduceLROnPlateau(factor=0.5, patience=1)
    ]
)

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

训练过程中，EarlyStopping回调会在验证集准确率不再提升时自动停止训练并恢复最佳权重，避免过拟合。ReduceLROnPlateau则在验证损失停滞时降低学习率，帮助模型收敛到更优解。对于MCU场景，我们不需要追求极致精度（90%以上即可），更重要的是模型规模与精度的平衡。

步骤4：导出为SavedModel格式

# 导出为SavedModel（推荐）
model.save('fashion_mnist_model', save_format='tf')

# 同时保存一份.h5备份（便于快速加载和检查）
model.save('fashion_mnist_model.h5')

# 验证导出结果
loaded_model = tf.keras.models.load_model('fashion_mnist_model')
print(f"加载模型输入: {loaded_model.inputs}")
print(f"加载模型输出: {loaded_model.outputs}")

导出后，fashion_mnist_model目录下应包含saved_model.pb、variables/和assets/三个子目录。使用load_model验证可以确保模型完整可恢复，这是转换前的重要自检步骤。

3.3 案例拆解：一个真实部署中的模型导出问题排查

在某智能家居项目中，我们需要在ESP32-S3上部署一个关键词唤醒模型（识别“你好”和“关闭”两个词）。团队最初使用PyTorch训练了一个GRU模型，并转换为ONNX再转TFLite，结果在量化时遇到严重精度损失（准确率从95%降至82%）。后来改用Keras重新训练，并严格遵循SavedModel导出规范，问题得到解决。

对比维度	PyTorch + ONNX方案	Keras + SavedModel方案
量化支持	需手动指定输入输出节点，易出错	Converter自动解析签名，稳定可靠
算子兼容性	GRU算子需映射为TFLite自定义算子	改用Conv1D + Dense，全部为内置算子
模型大小	ONNX文件包含冗余元数据，约200KB	SavedModel + 量化后仅35KB
调试便利性	需要Netron等工具手动检查节点	Keras的summary()直接输出参数量

这个案例说明，选择Keras + SavedModel不仅是技术偏好，更是工程效率的保障。在嵌入式AI领域，时间成本和调试成本往往比模型精度更重要。

3.4 常见误区与总结复盘

误区1：忽略输入shape的固定。Keras默认允许动态batch维度，但TFLite Micro要求完全固定的shape。如果使用Input(shape=(None, 28, 28, 1))，转换时会报错。务必使用具体数值。

误区2：使用不支持的激活函数。例如LeakyReLU在TFLite Micro中需要额外算子支持，建议使用ReLU或ReLU6。如果必须使用LeakyReLU，需在转换时指定allow_custom_ops=True，但这会增加部署复杂度。

误区3：导出时未指定save_format='tf'。在TensorFlow 2.x中，model.save('path')默认会根据文件扩展名推断格式。如果路径没有.h5后缀，默认保存为SavedModel；但为了明确性，建议显式指定save_format='tf'。

经验总结：训练阶段就要为部署考虑。建议在模型设计时遵循以下原则：优先使用Conv2D + DepthwiseConv2D + Dense的组合；激活函数只用ReLU/ReLU6/Softmax；避免使用BatchNormalization（量化时可能引入额外计算）；输入输出张量命名清晰。这些习惯能大幅减少后续转换和部署的排障时间。

最后，务必在导出后执行一次完整的加载验证，并检查saved_model.pb的文件大小。如果模型文件超过1MB，说明网络结构过于复杂，需要重新设计。对于TFLite Micro，理想的目标是模型文件小于200KB，这样在常见的MCU（如STM32F4、ESP32）上才能流畅运行。