4、模型训练与导出：使用Keras/TensorFlow训练图像分类模型，保存为.h5模型文件，导出为TensorFlow Lite格式。

好，到了这一步，咱们终于要开始动手训练模型了。

说实话，很多做嵌入式的朋友一听到「训练模型」就头大。觉得那是算法工程师的事。我以前也这么想，直到有一次项目里，算法团队给的模型在PC上跑得飞起，一放到STM32上直接内存溢出……嗯，从那以后我就决定，自己得会这一套流程。

说白了，嵌入式AI工程师不需要会发明新网络结构。但你要会训练、会导出、会量化。今天这一章，我就带你走一遍完整的流程。

4.1 准备数据集

先说说数据集。我习惯用一个小数据集来演示，比如识别猫和狗。你想想看，如果连二分类都搞不定，多分类就更别谈了。

数据集的结构很简单：

dataset/
├── train/
│   ├── cat/    (猫的图片)
│   └── dog/    (狗的图片)
├── validation/
│   ├── cat/
│   └── dog/
└── test/
    ├── cat/
    └── dog/

每个类别我建议至少100张图。太少的话，模型学不到特征。我在项目中遇到过客户只给了20张图，结果模型准确率只有60%，后来加了数据增强才勉强能用。

4.2 构建模型

对于STM32这种资源受限的设备，模型不能太大。我个人习惯用MobileNetV2或者自己搭一个小型CNN。

这里我写一个轻量级的例子：

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

def build_small_cnn(input_shape=(128, 128, 3), num_classes=2):
    model = models.Sequential([
        layers.Conv2D(16, (3, 3), activation='relu', input_shape=input_shape),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(32, (3, 3), activation='relu'),
        layers.MaxPooling2D(2, 2),
        layers.Conv2D(64, (3, 3), activation='relu'),
        layers.MaxPooling2D(2, 2),
        layers.Flatten(),
        layers.Dense(128, activation='relu'),
        layers.Dropout(0.5),
        layers.Dense(num_classes, activation='softmax')
    ])
    return model

model = build_small_cnn()
model.summary()

你看，这个模型只有三层卷积加一个全连接层。参数量大概在几十万级别。STM32F4系列勉强能跑，F7或H7系列就轻松多了。

我的小技巧： 输入图片尺寸不要太大。128x128已经够用。我之前试过224x224，模型大了三倍，准确率只提升了2%。不值得。

4.3 数据预处理与增强

数据增强是必须的。为什么？因为嵌入式设备采集的图片，光照、角度、遮挡都千奇百怪。你不做增强，模型在实验室里跑99%，一到现场就掉到70%。

from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest'
)

validation_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    'dataset/train',
    target_size=(128, 128),
    batch_size=32,
    class_mode='categorical'
)

validation_generator = validation_datagen.flow_from_directory(
    'dataset/validation',
    target_size=(128, 128),
    batch_size=32,
    class_mode='categorical'
)

这里要注意，验证集不要做增强。只做归一化就行。否则你没法评估模型的真实表现。

4.4 训练模型

训练过程其实很简单。但我建议你关注几个关键点：

  • 学习率： 我一般从0.001开始。如果loss震荡，就降到0.0001。
  • 早停法： 监控验证集loss，连续5个epoch不下降就停止。
  • 模型保存： 只保存验证集loss最小的那个模型。
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint

callbacks = [
    EarlyStopping(monitor='val_loss', patience=5, restore_best_weights=True),
    ModelCheckpoint('best_model.h5', monitor='val_loss', save_best_only=True)
]

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
    loss='categorical_crossentropy',
    metrics=['accuracy']
)

history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // 32,
    epochs=50,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples // 32,
    callbacks=callbacks
)

训练完成后，你会得到一个 best_model.h5 文件。这就是我们要的模型。

注意： 训练时如果显存不够，减小batch_size。我遇到过batch_size设64直接OOM，改成16就好了。另外，建议用GPU训练，CPU训练一个模型可能要几个小时。

4.5 评估模型

训练完了，别急着导出。先看看模型在测试集上的表现。

test_datagen = ImageDataGenerator(rescale=1./255)
test_generator = test_datagen.flow_from_directory(
    'dataset/test',
    target_size=(128, 128),
    batch_size=32,
    class_mode='categorical'
)

loss, accuracy = model.evaluate(test_generator)
print(f'Test accuracy: {accuracy:.4f}')

如果准确率低于90%，我建议你回去调整模型或数据。我曾经有一个项目，测试集准确率只有85%，放到STM32上直接掉到70%。后来发现是量化精度损失导致的。所以，尽量让模型在PC上就达到95%以上。

4.6 导出为TensorFlow Lite格式

这一步是关键。我们要把.h5模型转换成.tflite文件。而且，为了在STM32上跑，必须做量化。

import tensorflow as tf

# 加载训练好的模型
model = tf.keras.models.load_model('best_model.h5')

# 转换为TFLite格式，并做量化
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 这里做量化：将float32转为int8
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]  # 或者用int8

# 如果你要用int8量化，需要提供代表数据集
def representative_dataset():
    for _ in range(100):
        data = np.random.randn(1, 128, 128, 3).astype(np.float32)
        yield [data]

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_model = converter.convert()

# 保存模型
with open('model_quantized.tflite', 'wb') as f:
    f.write(tflite_model)

print('模型导出完成！')
重点： 量化后的模型大小会缩小到原来的1/4左右。比如原来10MB的模型，量化后只有2.5MB。这对于STM32的Flash来说，是巨大的优势。

4.7 验证导出的TFLite模型

导出之后，我建议你在PC上先验证一下。别直接丢到板子上才发现不对。

import numpy as np
import tensorflow as tf

# 加载TFLite模型
interpreter = tf.lite.Interpreter(model_path='model_quantized.tflite')
interpreter.allocate_tensors()

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

# 准备一张测试图片
test_image = ...  # 读取并预处理图片，尺寸128x128，归一化
test_image = np.expand_dims(test_image, axis=0).astype(np.int8)

# 推理
interpreter.set_tensor(input_details[0]['index'], test_image)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])

print('预测结果:', output)

如果输出结果和原始模型一致，那就没问题了。如果不一致，检查一下量化参数是否正确。

4.8 常见问题与避坑

问题 原因 解决方法
量化后准确率暴跌 代表数据集不够代表性 增加代表数据集样本，覆盖各种场景
TFLite模型无法加载 使用了不支持的算子 检查模型是否包含自定义层，替换为标准层
模型太大放不进Flash 没有做量化或量化不到位 使用int8量化，并考虑减小输入尺寸
推理速度太慢 模型层数太多或输入太大 减少卷积层数，使用深度可分离卷积

我曾经犯过一个低级错误：量化时忘了设置代表数据集，结果模型直接输出乱码。排查了两天才发现……嗯，从那以后我每次导出都会先验证一遍。

4.9 小结

这一章我们走完了从数据集准备到TFLite导出的全流程。你想想看，其实并不复杂。核心就三步：

  1. 用Keras训练一个轻量级CNN模型
  2. 保存为.h5文件
  3. 用TFLiteConverter量化导出为.tflite

下一章，我们会把这个.tflite模型部署到STM32上，用CubeAI生成C代码。到时候你就能在板子上跑自己的模型了。

记住，模型训练是基础。基础打牢了，后面移植才顺利。别急着跳步骤，一步一步来。

← 上一章 📖 返回目录 下一章 →