3. 模型训练与导出(上):用 TensorFlow 训练 CIFAR-10 图像分类模型
好,咱们进入实战环节。这一章的目标很明确:训练一个能识别 CIFAR-10 数据集的图像分类模型,然后把它量化、导出成 TFLite 格式。这是后续部署到 MCU 的第一步,也是最重要的一步。
我刚开始做 TinyML 时,总觉得训练模型是数据科学家的事,固件工程师只管跑推理就行。后来踩过几次坑才发现——模型怎么训练、用什么量化方式,直接决定了它在单片机上能不能跑、跑多快。所以这块咱们得亲手过一遍。
3.1 环境准备与数据集理解
先搭环境。我习惯用 TensorFlow 2.x,配合 Keras 高层 API,写起来最顺手。
pip install tensorflow==2.13.0 numpy matplotlib
CIFAR-10 是什么?说白了就是 10 类物体的 32x32 彩色小图。每类 6000 张,共 60000 张。训练集 50000 张,测试集 10000 张。
| 类别编号 | 类别名称 | 示例 |
|---|---|---|
| 0 | 飞机 | airplane |
| 1 | 汽车 | automobile |
| 2 | 鸟 | bird |
| 3 | 猫 | cat |
| 4 | 鹿 | deer |
| 5 | 狗 | dog |
| 6 | 青蛙 | frog |
| 7 | 马 | horse |
| 8 | 船 | ship |
| 9 | 卡车 | truck |
注意,32x32 的分辨率对 MCU 来说其实不小了。我见过很多项目直接拿 224x224 的 MobileNet 去移植,结果 Flash 和 RAM 双双爆掉。CIFAR-10 的尺寸刚好适合做教学——模型小、训练快、部署也现实。
3.2 构建模型:为 MCU 量身定制
模型结构怎么选?我的建议是:别一上来就搬 ResNet、EfficientNet。那些模型是为 GPU 设计的,参数动辄几百万,MCU 根本扛不住。
咱们自己搭一个小型 CNN,参数控制在 10 万以内。这样量化后模型体积大概 40~80KB,大部分 Cortex-M4/M7 都能跑。
import tensorflow as tf
from tensorflow.keras import layers, models
def build_tiny_cnn(input_shape=(32, 32, 3), num_classes=10):
model = models.Sequential([
# 第一层卷积:提取低级特征
layers.Conv2D(16, (3, 3), activation='relu', padding='same', input_shape=input_shape),
layers.MaxPooling2D((2, 2)),
# 第二层卷积:提取中级特征
layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
layers.MaxPooling2D((2, 2)),
# 第三层卷积:提取高级特征
layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
layers.MaxPooling2D((2, 2)),
# 分类头
layers.Flatten(),
layers.Dense(64, activation='relu'),
layers.Dense(num_classes, activation='softmax')
])
return model
model = build_tiny_cnn()
model.summary()
这个结构只有 3 层卷积 + 2 层全连接,参数量大约 8.5 万。我在一个 STM32F746 项目上试过,量化后跑一次推理只要 15ms 左右,完全够用。
3.3 训练流程与数据预处理
训练前要做数据归一化。CIFAR-10 的像素值范围是 0~255,直接喂给网络会让梯度不稳定。我习惯除以 255.0 缩放到 [0,1]。
# 加载数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
# 归一化
x_train = x_train.astype('float32') / 255.0
x_test = x_test.astype('float32') / 255.0
# 标签转 one-hot
y_train = tf.keras.utils.to_categorical(y_train, 10)
y_test = tf.keras.utils.to_categorical(y_test, 10)
编译模型时,优化器我选 Adam,学习率默认 0.001 就挺好。损失函数用 categorical_crossentropy,多分类标配。
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
history = model.fit(x_train, y_train,
batch_size=64,
epochs=20,
validation_data=(x_test, y_test))
训练 20 个 epoch,在 GPU 上大概 5~10 分钟。CPU 的话可能要半小时。最终测试精度一般在 70%~75% 之间。别嫌低——对于 8.5 万参数的小模型,这个精度已经不错了。我见过有人用全连接层硬怼,精度只有 50% 多。
3.4 模型量化:Post-training Quantization
模型训练好了,但权重是 float32 的。一个 float32 的权重占 4 字节,8.5 万个参数就是 340KB。MCU 的 Flash 通常只有 512KB 或 1MB,还得放代码和中间结果,根本不够用。
所以我们要做 后训练量化(Post-training Quantization)。说白了就是把 float32 的权重转成 int8,体积直接缩到 1/4,推理速度还能快 2~4 倍。
TensorFlow 提供了两种量化方式:
| 量化方式 | 精度损失 | 推理速度 | 适用场景 |
|---|---|---|---|
| 动态范围量化 | 较小 | 中等 | 快速验证 |
| 全整数量化 | 稍大 | 最快 | MCU 部署 |
咱们的目标是 MCU,所以直接上 全整数量化。它需要一个小型校准数据集来统计激活值的范围。
# 保存 float32 模型
model.save('cifar10_float32.h5')
# 转换为 TFLite 并量化
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.uint8
converter.inference_output_type = tf.uint8
tflite_quant_model = converter.convert()
# 保存量化模型
with open('cifar10_quantized.tflite', 'wb') as f:
f.write(tflite_quant_model)
这里 representative_dataset 是一个生成器,从训练集中取 100~200 张图片做校准。
def representative_dataset():
for i in range(200):
yield [x_train[i:i+1]]
3.5 量化效果验证
量化完一定要验证精度。我见过有人量化完直接部署,结果精度从 72% 掉到 50%,还以为是模型没训练好。
# 加载量化模型
interpreter = tf.lite.Interpreter(model_path='cifar10_quantized.tflite')
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# 验证精度
correct = 0
for i in range(len(x_test)):
# 输入需要转成 uint8
input_data = (x_test[i] * 255.0).astype(np.uint8).reshape(1, 32, 32, 3)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])
if np.argmax(output) == np.argmax(y_test[i]):
correct += 1
print(f'Quantized accuracy: {correct / len(x_test):.4f}')
一般来说,全整数量化后精度损失在 1%~3% 之间。如果损失超过 5%,说明模型本身对量化不友好,可能需要调整结构或做量化感知训练(QAT)。
3.6 本章知识体系
下面这张图把整个流程串起来了,你可以对照着看:
3.7 本章小结
咱们走完了从数据加载到模型量化的完整流程。你手里现在应该有两个文件:cifar10_float32.h5 和 cifar10_quantized.tflite。后者就是咱们下一阶段要部署到 MCU 上的模型。
量化后的模型体积大概 40~80KB,精度 70% 左右。这个精度在 MCU 上已经够用了——你想想看,一个 200MHz 的 Cortex-M4 能在几十毫秒内完成推理,代价只是损失两三个百分点的精度,这笔账怎么算都划算。