4. TensorFlow模型训练与导出：使用Keras训练一个简单分类模型，保存为.h5文件，准备转换
好，到了真正动手的环节了。前面几章我们聊了模型迭代的整体思路和STM32Cube.AI的工具链，现在该把模型跑起来了。
这一章的目标很明确：用Keras训练一个简单的分类模型，然后导出成.h5文件。别小看这一步，我见过太多人在导出环节翻车——模型在电脑上跑得飞起，一转换就报错。说白了，就是训练时埋的坑。
4.1 选一个适合MCU的模型结构
先说说模型选型。在STM32上跑模型，跟你在服务器上跑完全是两码事。MCU的资源就那么点——Flash几MB、RAM几百KB，你还想塞个ResNet-50进去？不现实。
我个人习惯，给MCU用的模型，参数控制在100K以内，层数不超过10层。你想想看，STM32F4系列的主频才168MHz，复杂模型推理一次要好几秒，谁受得了？
这里我选一个简单的全连接网络，做MNIST手写数字识别。经典中的经典，但足够说明整个流程。
- 输入层：28×28=784个神经元（展平后的像素值）
- 隐藏层1：128个神经元，ReLU激活
- 隐藏层2：64个神经元，ReLU激活
- 输出层：10个神经元，Softmax激活
总参数量：约107K，Flash占用约430KB，RAM占用约4KB（推理时）。
4.2 数据准备与预处理
MNIST数据集Keras自带，直接加载就行。但这里有个细节——数据归一化。
我曾经在项目里吃过亏：训练时没做归一化，模型收敛慢不说，转换到STM32后推理结果全乱套。为什么？因为MCU的浮点精度有限，大数值输入会让激活值溢出。
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np
# 加载MNIST数据集
(x_train, y_train), (x_test, y_test) = 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
# 展平：从28x28变成784
x_train = x_train.reshape(-1, 784)
x_test = x_test.reshape(-1, 784)
# 标签转one-hot编码
y_train = keras.utils.to_categorical(y_train, 10)
y_test = keras.utils.to_categorical(y_test, 10)
print(f"训练集形状: {x_train.shape}, 测试集形状: {x_test.shape}")4.3 构建模型
模型构建用Keras Sequential API，简单直观。注意，这里我刻意不用BatchNormalization和Dropout——这两个层在STM32Cube.AI转换时偶尔会出问题，尤其是老版本的工具链。
model = keras.Sequential([
layers.Dense(128, activation='relu', input_shape=(784,), name='hidden1'),
layers.Dense(64, activation='relu', name='hidden2'),
layers.Dense(10, activation='softmax', name='output')
], name='mnist_classifier')
model.summary()模型摘要输出：
| 层名 | 类型 | 输出形状 | 参数量 |
|---|---|---|---|
| hidden1 | Dense | (None, 128) | 100,480 |
| hidden2 | Dense | (None, 64) | 8,256 |
| output | Dense | (None, 10) | 650 |
| 总计 | 109,386 | ||
4.4 编译与训练
编译时选Adam优化器，学习率用默认的0.001。损失函数用categorical_crossentropy，因为是多分类问题。
model.compile(
optimizer=keras.optimizers.Adam(learning_rate=0.001),
loss='categorical_crossentropy',
metrics=['accuracy']
)
history = model.fit(
x_train, y_train,
batch_size=128,
epochs=10,
validation_data=(x_test, y_test),
verbose=1
)训练过程大概30秒就结束了。最终测试精度在97%左右——够用，但不是重点。重点是后面导出和转换。
4.5 模型评估与验证
训练完先评估一下，确保模型没跑偏。
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"测试精度: {test_acc:.4f}")
# 随机抽几个样本看看预测结果
import matplotlib.pyplot as plt
predictions = model.predict(x_test[:5])
for i in range(5):
pred_label = np.argmax(predictions[i])
true_label = np.argmax(y_test[i])
print(f"样本{i}: 预测={pred_label}, 真实={true_label}")输出结果：
测试精度: 0.9734
样本0: 预测=7, 真实=7
样本1: 预测=2, 真实=2
样本2: 预测=1, 真实=1
样本3: 预测=0, 真实=0
样本4: 预测=4, 真实=4嗯，没问题。模型在电脑上跑通了。
4.6 导出为.h5文件
关键步骤来了。导出.h5文件，STM32Cube.AI才能识别。
# 保存整个模型（结构+权重+优化器状态）
model.save('mnist_model.h5')
print("模型已保存为 mnist_model.h5")
# 也可以只保存权重（更轻量，但需要配合模型结构代码）
# model.save_weights('mnist_weights.h5')- 使用
model.save()保存完整模型，包含网络结构、权重和训练配置。 - 文件后缀必须是
.h5，STM32Cube.AI只认这个格式。 - 保存路径不要有中文和空格，否则工具链可能报错。
- 保存前确认模型是
training=False状态（即推理模式）。
导出后，用h5py库验证一下文件完整性：
import h5py
with h5py.File('mnist_model.h5', 'r') as f:
print("文件中的顶层组:", list(f.keys()))
# 应该看到 'model_weights', 'optimizer_weights', 'training_config' 等输出：
文件中的顶层组: ['model_weights', 'optimizer_weights', 'training_config']好，文件没问题。接下来就可以交给STM32Cube.AI了。
4.7 本章知识体系
下面这张图总结了从数据到.h5文件的完整流程：
4.8 常见问题与避坑
最后，分享几个我踩过的坑：
- 模型保存后加载报错：检查Keras版本。STM32Cube.AI对Keras 2.x支持最好，2.3.x到2.9.x都行。2.10以上可能有问题。
- 转换时提示不支持层：我遇到过Lambda层和自定义层无法转换。解决办法：用标准层替代，或者用
tf.keras.layers里的现成层。 - 推理结果全一样：大概率是预处理没对齐。训练时归一化到[0,1]，部署时也要一样。我曾经在STM32上忘了做归一化，结果模型输出全是0.1。
- 模型太大放不进Flash：参数量超过200K就要考虑剪枝或量化了。后面章节会详细讲。