📚 课程目录第4章 TensorFlow模型基础

4、TensorFlow模型基础：构建简单Keras模型、训练并保存为H5或SavedModel格式

在进入TFLite Micro的模型转换流程之前，首先需要理解上游模型是如何产生和保存的。TFLite Micro本质上是一个推理引擎，它不负责训练，而是消费已经训练好的模型文件。因此，掌握TensorFlow中Keras API的建模、训练与导出方法，是后续所有转换工作的前提。本章将围绕一个具体的图像分类任务，完整演示从数据准备、模型构建、训练到保存为H5与SavedModel两种格式的全过程，并深入剖析两种格式的差异及其对后续TFLite转换的影响。

4.1 核心原理：Keras模型的生命周期与导出格式

Keras是TensorFlow的高层API，它提供了简洁的接口来定义、训练和保存深度学习模型。一个完整的Keras模型生命周期包含五个阶段：数据准备、模型构建、模型编译、模型训练和模型保存。其中，模型保存是连接训练与部署的桥梁，其格式选择直接影响后续TFLite转换的路径。

在TFLite Micro的上下文中，我们最关心的是模型的计算图结构和权重参数。Keras模型在内存中是一个包含层（Layer）、权重（Weights）、优化器状态（Optimizer State）和编译配置的复合对象。保存模型时，我们需要决定保留哪些信息：是仅保留推理所需的架构和权重，还是连同训练状态一起保存以便恢复训练。

TensorFlow提供了两种主要的保存格式：H5（Hierarchical Data Format 5）和SavedModel。H5是Keras的传统格式，它将模型架构、权重和优化器状态打包进一个二进制文件中，适合单文件分发。SavedModel是TensorFlow的推荐格式，它是一个包含变量、图和资产文件的目录结构，更适合生产环境中的版本管理和服务部署。对于TFLite转换而言，两种格式都可以作为输入，但转换器的API调用方式略有不同。

关键理解：H5格式本质上是Keras的“快照”，它依赖于Keras的层定义来反序列化模型；而SavedModel是TensorFlow的“标准交换格式”，它包含了完整的计算图定义，不依赖Keras的Python类。因此，SavedModel在跨版本兼容性和部署灵活性上更优，而H5在简单场景下更便捷。

4.2 操作流程：从零构建并训练一个图像分类模型

下面我们以经典的Fashion-MNIST数据集为例，演示一个完整的模型构建与训练流程。Fashion-MNIST包含10类灰度服饰图像，每张图像尺寸为28×28像素，非常适合作为TFLite Micro的入门示例。

4.2.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,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]

# 划分验证集（从训练集中取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部署，输入数据的格式必须与转换时保持一致，否则推理结果会出错。

4.2.2 构建模型架构

model = keras.Sequential([
    keras.layers.Conv2D(8, (3,3), activation='relu', input_shape=(28,28,1)),
    keras.layers.MaxPooling2D((2,2)),
    keras.layers.Conv2D(16, (3,3), activation='relu'),
    keras.layers.MaxPooling2D((2,2)),
    keras.layers.Flatten(),
    keras.layers.Dense(10, activation='softmax')
])

model.summary()

这个模型包含两个卷积层和两个池化层，最后通过全连接层输出10个类别的概率分布。选择较小的卷积核数量（8和16）是为了控制模型大小，使其适合微控制器有限的存储资源。对于TFLite Micro，模型参数总量通常需要控制在几百KB以内，因此设计网络时需要在精度和体积之间权衡。

4.2.3 编译与训练

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

history = model.fit(
    x_train, y_train,
    batch_size=32,
    epochs=5,
    validation_data=(x_val, y_val)
)

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

编译阶段需要指定优化器、损失函数和评估指标。这里使用Adam优化器和稀疏交叉熵损失（因为标签是整数而非one-hot编码）。训练过程中，fit()方法会返回一个History对象，记录每个epoch的损失和准确率，便于后续可视化分析。

4.3 案例拆解：保存模型并验证两种格式的差异

训练完成后，我们需要将模型保存为H5和SavedModel两种格式，并验证它们的结构差异。这是一个关键的实操环节，因为不同的保存方式会影响后续TFLite转换的代码路径。

对比维度	H5格式	SavedModel格式
文件形态	单个二进制文件（.h5）	目录（包含saved_model.pb、variables/、assets/）
内容完整性	包含架构、权重、优化器状态、编译配置	包含完整计算图、变量、签名（Signature）
依赖Keras层定义	是，反序列化时需要Keras层类	否，图定义独立于Keras
TFLite转换API	`tf.lite.TFLiteConverter.from_keras_model(model)`	`tf.lite.TFLiteConverter.from_saved_model(dir)`
适用场景	快速实验、模型交换	生产部署、版本管理、跨框架使用

4.3.1 保存为H5格式

model.save('fashion_mnist_model.h5')
print("H5模型已保存")

H5格式保存后，可以使用keras.models.load_model()重新加载。需要注意的是，H5文件包含了优化器状态，因此加载后可以直接继续训练。但对于TFLite转换，我们只关心推理部分，优化器状态会被忽略。

4.3.2 保存为SavedModel格式

model.save('saved_model_fashion', save_format='tf')
print("SavedModel目录已保存")

SavedModel格式会生成一个目录，其中saved_model.pb是计算图的序列化定义，variables/目录存放权重变量。这种格式不依赖Keras的Python环境，可以被TensorFlow Serving、TFLite等不同组件直接消费。

经验建议：在TFLite Micro项目中，推荐使用SavedModel格式作为中间产物。原因有二：一是SavedModel包含明确的输入/输出签名，转换时不容易出错；二是当模型需要迭代更新时，SavedModel的目录结构便于版本管理。

4.3.3 验证两种格式的加载

# 加载H5模型
h5_model = keras.models.load_model('fashion_mnist_model.h5')

# 加载SavedModel
saved_model = tf.saved_model.load('saved_model_fashion')

# 验证推理结果一致性
import numpy as np
sample = x_test[0:1]
h5_pred = h5_model.predict(sample)
saved_pred = saved_model(sample)
print(f"H5预测类别: {np.argmax(h5_pred)}, SavedModel预测类别: {np.argmax(saved_pred.numpy())}")

通过对比预测结果，可以确认两种格式保存的模型在推理行为上完全一致。这个验证步骤非常重要，它确保我们在后续TFLite转换时不会因为格式问题引入误差。

4.4 常见误区与总结复盘

在模型构建和保存过程中，初学者经常遇到以下几类问题，需要特别注意：

输入形状不匹配：Keras模型在构建时定义的input_shape必须与训练数据的形状严格一致。如果训练时使用了(28,28,1)，但转换TFLite时传入的输入张量形状是(28,28)，会导致转换失败或推理错误。
忽略归一化层：很多人在训练时手动对数据做归一化，但模型本身不包含归一化操作。这意味着部署到TFLite Micro后，微控制器上的输入数据也必须手动归一化，否则模型性能会大幅下降。更稳妥的做法是将归一化作为模型的第一层（如tf.keras.layers.Rescaling），这样转换后的TFLite模型就自带预处理能力。
保存格式选择不当：如果使用model.save('model.h5')，Keras会根据文件扩展名自动选择H5格式；如果使用model.save('model_dir')，则默认保存为SavedModel。建议显式指定save_format参数，避免歧义。
优化器状态带来的体积膨胀：H5格式会保存优化器的动量等信息，导致文件体积比纯权重文件大。对于TFLite转换，这些信息是多余的，但不会影响转换结果。

风险提示：在训练过程中，如果使用model.save_weights()只保存权重，而不保存架构，那么后续加载时必须重新构建完全相同的模型架构。这种方式不适合TFLite转换，因为TFLite需要完整的计算图定义。务必使用model.save()保存完整模型。

4.4.1 案例复盘：一个典型的转换失败场景

假设某工程师训练了一个模型，输入形状为(32,32,3)，但在保存时使用了model.save('model.h5')。随后他尝试用TFLiteConverter转换，却遇到以下错误：

ValueError: None is only supported in the 1st dimension. Tensor 'input_1' has invalid shape 'None, 32, 32, 3'

这个错误的根源在于Keras模型的输入维度中，batch size被表示为None（动态维度）。TFLite转换器要求输入形状完全固定，不能有动态维度。解决方案是在转换时指定固定的batch size：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.experimental_new_converter = True
# 通过指定输入形状来固定batch size
converter._experimental_default_to_single_batch_in_tensor_list_ops = True
tflite_model = converter.convert()

更推荐的做法是在模型构建时就使用固定的batch size，或者在转换前用model.build(input_shape)重新指定形状。

4.4.2 总结：模型保存的最佳实践

为了确保后续TFLite Micro转换的顺利进行，建议遵循以下流程：

在模型构建时，明确输入形状，并考虑将归一化等预处理操作纳入模型内部。
训练完成后，同时保存H5和SavedModel两种格式，H5用于快速验证，SavedModel用于正式转换。
在转换前，使用model.summary()确认模型结构，并检查参数量是否适合目标微控制器的存储限制。
验证模型在测试集上的性能，确保精度满足需求后再进行转换。

核心要点回顾：本章通过一个完整的Fashion-MNIST分类示例，演示了Keras模型的构建、训练和保存流程。重点在于理解H5与SavedModel两种格式的本质差异：H5是Keras的私有格式，依赖Python环境；SavedModel是TensorFlow的标准格式，独立于Keras。对于TFLite Micro部署，推荐使用SavedModel作为转换输入。同时，要注意输入形状的固定性、归一化操作的位置以及模型体积的控制，这些因素直接影响后续转换的成败和部署效果。