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 左右,完全够用。

关键点:padding='same' 保证特征图尺寸不缩太快,否则 32x32 的输入经过三层池化就只剩 4x4 了,信息损失太大。

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% 多。

小技巧:如果精度一直上不去,试试加一点数据增强(随机翻转、平移)。但注意:MCU 端部署时不能做在线增强,所以训练时的增强要轻量。

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]]
我曾经踩过的坑:校准数据集一定要从训练集取,不要用测试集。用测试集会高估量化后的精度,部署到 MCU 上才发现实际效果差很多。另外,校准样本不要太多,200 张足够,多了反而可能过拟合到校准集上。

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 本章知识体系

下面这张图把整个流程串起来了,你可以对照着看:

CIFAR-10 模型训练与量化导出流程 数据准备 CIFAR-10 加载+归一化 模型构建 3层CNN + 2层FC 模型训练 Adam, 20 epochs 保存 Float32 cifar10_float32.h5 全整数量化 int8 权重+激活 校准数据集 200张训练样本 导出 TFLite cifar10_quantized.tflite 精度验证 ⚠ 注意 校准集用训练集 精度损失 1%~3%

3.7 本章小结

咱们走完了从数据加载到模型量化的完整流程。你手里现在应该有两个文件:cifar10_float32.h5 和 cifar10_quantized.tflite。后者就是咱们下一阶段要部署到 MCU 上的模型。

量化后的模型体积大概 40~80KB,精度 70% 左右。这个精度在 MCU 上已经够用了——你想想看,一个 200MHz 的 Cortex-M4 能在几十毫秒内完成推理,代价只是损失两三个百分点的精度,这笔账怎么算都划算。

个人建议:把量化前后的模型都保存好。float32 版本可以用来做精度基准,int8 版本用来部署。如果后续发现部署效果不理想,可以回头对比两个版本的输出差异,定位问题。