第4章：TensorFlow Lite 模型准备

好，咱们进入第四章。这一章说白了，就是为后面的STM32部署准备“弹药”——训练一个能用的模型，再把它转换成TFLite格式。

我个人习惯把这一步叫做“模型出厂前处理”。你想想看，模型在PC上跑得再好，到了MCU上水土不服，那也白搭。所以这一章，咱们把基础打扎实。

4.1 安装TensorFlow 2.x

先搞定环境。我建议你用Python 3.8-3.10版本，太新或太旧都可能遇到兼容性问题。

# 创建虚拟环境（推荐）
python -m venv tf_env
source tf_env/bin/activate  # Linux/Mac
# 或 tf_env\Scripts\activate  # Windows

# 安装TensorFlow 2.x
pip install tensorflow==2.13.0

# 验证安装
python -c "import tensorflow as tf; print(tf.__version__)"

嗯，这里要注意：如果你用的是Apple Silicon芯片，建议装tensorflow-macos版本。我在M1 Mac上踩过坑，直接装普通版会报错。

4.2 训练一个简单的CNN模型

咱们用MNIST手写数字识别，经典中的经典。这个数据集包含0-9的手写数字图片，每张28x28像素。

为什么选它？因为模型小、训练快、效果直观。我在项目中经常用它做“冒烟测试”——先跑通MNIST，再换自己的数据集。

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

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

# 归一化到[0,1]区间
x_train, x_test = x_train / 255.0, x_test / 255.0

# 增加通道维度：28x28 -> 28x28x1
x_train = x_train[..., tf.newaxis]
x_test = x_test[..., tf.newaxis]

# 构建CNN模型
model = models.Sequential([
    layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(64, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(10, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=5, validation_data=(x_test, y_test))

# 评估模型
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2)
print(f'\n测试准确率: {test_acc:.4f}')

这个网络结构很简单：两层卷积+池化，然后全连接。我刚开始做嵌入式AI时，总觉得网络越深越好，后来发现STM32根本跑不动。你想想看，MCU的Flash和RAM就那么点，模型必须精简。

4.3 模型保存为.tflite格式

训练好的模型是Keras格式（.h5或.keras），但STM32不认识。我们需要转换成TFLite格式。

# 保存Keras模型（可选）
model.save('mnist_cnn.h5')

# 转换为TFLite格式
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()

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

print("模型已转换为TFLite格式，大小：", len(tflite_model), "字节")

转换后你会得到一个二进制文件。我习惯用ls -lh看看大小——原始模型大概几百KB，对于STM32来说还是偏大。别急，下一节咱们做量化。

4.4 模型量化（int8）

量化，说白了就是把模型从float32压缩到int8。这样做的好处有两个：模型体积缩小4倍，推理速度提升2-4倍。代价是精度会略微下降，但通常可以接受。

我在项目中做过对比：MNIST模型量化后准确率从99.2%降到98.8%，几乎没差别。但模型大小从300KB降到75KB，STM32F4系列都能轻松装下。

# 准备代表性数据集（用于校准）
def representative_dataset():
    for i in range(100):
        yield [x_train[i:i+1].astype(np.float32)]

# 带量化的转换
converter = tf.lite.TFLiteConverter.from_keras_model(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()

# 保存量化后的模型
with open('mnist_cnn_quant.tflite', 'wb') as f:
    f.write(tflite_quant_model)

print("量化模型大小：", len(tflite_quant_model), "字节")

这里有个关键点：representative_dataset函数提供100张图片做校准。数量太少，量化效果差；太多，转换时间变长。我一般取100-500张，够用了。

4.5 验证量化模型

转换完别急着部署，先验证一下精度。

# 加载量化模型
interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 测试一张图片
test_image = x_test[0:1].astype(np.int8)  # 注意类型转换
interpreter.set_tensor(input_details[0]['index'], test_image)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])

print("预测结果：", np.argmax(output))
print("真实标签：", y_test[0])

如果精度下降在1%以内，那就没问题。如果下降太多，可以尝试用float16量化（体积减半，精度损失更小），或者调整校准数据集。

4.6 本章小结

咱们这一章干了四件事：装环境、训模型、转格式、做量化。最终拿到一个mnist_cnn_quant.tflite文件，大概75KB左右。

这个文件就是下一章要部署到STM32上的“弹药”。我个人习惯把量化前后的模型都留着——调试时用float32版本，部署时用int8版本。

最后提醒一句：模型文件要放在STM32项目的Model文件夹里，以数组形式嵌入。具体怎么做，下一章咱们细聊。