4、模型训练准备：数据集采集与标注、数据增强策略、训练集/验证集/测试集划分，以及使用TensorFlow/Keras或PyTorch训练模型
在STM32CubeAI的模型移植流程中，模型训练准备阶段直接决定了最终部署到MCU上的模型性能上限。很多工程师在完成环境搭建后，急于开始训练网络，却忽略了数据层面的系统性工程。实际上，对于嵌入式视觉任务而言，数据质量与数据分布的管理远比网络结构的选择更为关键。本章将围绕数据集采集与标注、数据增强策略、数据集划分以及基于TensorFlow/Keras与PyTorch的训练实现展开，帮助您建立一套可复现、可评估、可部署的模型训练工作流。
4.1 数据集采集：场景覆盖与样本多样性
核心原理：嵌入式视觉模型部署后所面对的真实环境是复杂多变的。STM32CubeAI生成的模型通常运行在资源受限的MCU上，其泛化能力完全依赖于训练数据的分布覆盖度。如果训练数据仅包含单一光照、单一背景或单一角度的样本，模型在实机运行时会频繁出现误判。数据采集的核心目标是让训练集尽可能逼近目标部署场景的真实分布，包括光照强度、色温、视角偏移、遮挡程度、运动模糊以及目标物自身的形态变化。
操作流程：在采集阶段，建议按照以下步骤执行：
- 场景清单制定：列出目标应用可能出现的所有场景，例如室内/室外、白天/夜晚、强光/逆光、近距离/远距离等，并为每个场景分配采集比例。
- 设备参数固定：使用与最终部署硬件相同或相近的摄像头模组（如OV2640、OV5640），固定分辨率、帧率、白平衡模式，避免训练与推理时的域偏移。
- 样本量规划：对于分类任务，每类至少采集500-1000张原始图片；对于检测任务，每类目标实例不少于2000个标注框。若条件允许，原始数据量建议达到最终训练样本量的1.5倍，为清洗和筛选留出余量。
- 数据清洗：剔除模糊、过曝、失焦以及标注错误的图片，使用直方图或感知哈希算法辅助去重。
4.2 数据标注：格式选择与质量控制
核心原理：标注质量直接决定了模型学习的上限。对于STM32CubeAI支持的视觉任务，常见标注格式包括：图像分类的文件夹标签、目标检测的Pascal VOC XML或YOLO TXT格式、语义分割的PNG掩码。标注的核心矛盾在于人工成本与标注精度之间的平衡。对于嵌入式场景，建议优先选择边界框标注，因为其成本低于像素级分割，且对MCU端的目标检测任务足够有效。
标注流程与工具：
- 使用LabelImg或Labelme工具进行标注，导出为VOC格式或COCO格式。
- 设置标注规范：明确边界框包含完整目标物，遮挡超过50%的目标不标注，模糊目标不标注。
- 执行交叉验证：由第二人随机抽取10%的标注结果进行复核，计算IoU（Intersection over Union）一致性，要求IoU > 0.85。
- 对于分类任务，检查类别分布是否均衡，若存在长尾分布，需补充少数类样本或采用类别权重。
| 标注类型 | 适用任务 | STM32CubeAI支持格式 | 标注成本 | 推荐工具 |
|---|---|---|---|---|
| 分类标签 | 图像分类 | 文件夹结构 | 低 | 自写脚本 |
| 边界框 | 目标检测 | VOC XML / YOLO TXT | 中 | LabelImg |
| 实例分割 | 像素级分割 | COCO JSON | 高 | Labelme |
4.3 数据增强策略：提升泛化能力的关键手段
核心原理：数据增强通过在训练过程中对原始样本施加随机变换，生成语义不变但外观多样的新样本，从而扩大训练数据的有效分布范围。对于MCU部署场景，增强策略需要特别关注光照鲁棒性和几何鲁棒性，因为嵌入式设备往往面临视角变化和光照波动。但需注意，增强强度过大会导致训练分布与真实分布偏离，反而降低精度。
推荐增强组合：以下策略经过实际项目验证，适用于STM32CubeAI的典型视觉任务：
- 几何变换：随机水平翻转（概率0.5）、随机旋转（±15°）、随机缩放（0.8~1.2倍）、随机平移（±10%）。
- 色彩扰动：亮度调整（±20%）、对比度调整（±20%）、饱和度调整（±20%）、色相微调（±5°）。
- 噪声注入：高斯噪声（σ=0.01~0.03）、椒盐噪声（概率0.01），模拟传感器噪声。
- 高级策略：MixUp（按比例混合两张图及其标签）、CutMix（将一张图的区域粘贴到另一张图上）。
# TensorFlow/Keras 数据增强示例
from tensorflow.keras.preprocessing.image import ImageDataGenerator
datagen = ImageDataGenerator(
rotation_range=15,
width_shift_range=0.1,
height_shift_range=0.1,
brightness_range=[0.8, 1.2],
zoom_range=[0.8, 1.2],
horizontal_flip=True,
fill_mode='nearest'
)
# PyTorch 数据增强示例
from torchvision import transforms
transform = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(15),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.RandomResizedCrop(size=(224, 224), scale=(0.8, 1.0)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])4.4 训练集/验证集/测试集划分：评估可信度的基石
核心原理：三集划分的核心目的是模拟模型在未知数据上的表现。训练集用于参数更新，验证集用于超参数调优和早停判断，测试集仅用于最终评估。对于嵌入式模型，测试集必须与训练集完全隔离，且分布应代表真实部署场景。常见的错误是使用随机划分，导致同一场景的连续帧同时出现在训练集和测试集中，造成评估结果虚高。
划分策略与操作：
- 分层采样：按类别比例进行分层划分，确保每个类别在三集中的比例与原始数据集一致。
- 场景隔离：如果数据来自不同时间或不同地点，应确保同一场景的数据只出现在一个集合中。例如，按视频片段划分，而非按帧划分。
- 比例建议：对于中小规模数据集（<10万张），建议训练:验证:测试 = 7:2:1；对于大规模数据集，可调整为 8:1:1。
- 验证集与测试集区分：验证集用于比较不同超参数或不同网络结构的效果，测试集仅在最终模型确定后使用一次，避免对测试集产生过拟合。
| 数据集 | 用途 | 使用频率 | 典型比例 | 注意事项 |
|---|---|---|---|---|
| 训练集 | 参数更新 | 每个epoch | 70%-80% | 参与数据增强 |
| 验证集 | 超参数调优、早停 | 每个epoch后 | 10%-15% | 不参与训练，可重复使用 |
| 测试集 | 最终性能评估 | 仅一次 | 10%-15% | 完全隔离，模拟真实场景 |
4.5 使用TensorFlow/Keras与PyTorch训练模型
核心原理：TensorFlow/Keras和PyTorch是当前主流的深度学习框架。对于STM32CubeAI工作流，TensorFlow/Keras可以直接导出TFLite模型，而PyTorch需要先导出ONNX再转换为TFLite。两者在训练流程上本质相同，但API风格和调试体验有所差异。选择框架时，应考虑团队熟悉度以及后续量化工具链的兼容性。
训练流程与关键参数：
- 输入尺寸：STM32CubeAI对输入尺寸有严格限制，通常为32×32、48×48、96×96或224×224。训练时应使用与部署一致的尺寸，避免resize带来的精度损失。
- 优化器选择：推荐使用Adam（初始学习率1e-3）或SGD（初始学习率1e-2，动量0.9）。对于小数据集，Adam收敛更快；对于大数据集，SGD配合余弦退火可获得更好泛化。
- 损失函数：分类任务使用交叉熵，检测任务使用Smooth L1 + 交叉熵组合。
- 批次大小：受限于MCU内存，训练批次不宜过大，建议16-64之间。若显存不足，可使用梯度累积。
- 早停策略：监控验证集损失，若连续10个epoch无改善，则停止训练并恢复最佳权重。
# TensorFlow/Keras 训练示例
import tensorflow as tf
from tensorflow.keras import layers, models
model = models.Sequential([
layers.Conv2D(16, 3, activation='relu', input_shape=(96, 96, 3)),
layers.MaxPooling2D(2),
layers.Conv2D(32, 3, activation='relu'),
layers.MaxPooling2D(2),
layers.Flatten(),
layers.Dense(64, activation='relu'),
layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam',
loss='categorical_crossentropy',
metrics=['accuracy'])
history = model.fit(
train_generator,
validation_data=val_generator,
epochs=50,
callbacks=[tf.keras.callbacks.EarlyStopping(patience=10, restore_best_weights=True)]
)
# PyTorch 训练示例
import torch
import torch.nn as nn
import torch.optim as optim
class SimpleCNN(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
self.pool = nn.MaxPool2d(2)
self.fc1 = nn.Linear(32*24*24, 64)
self.fc2 = nn.Linear(64, 10)
def forward(self, x):
x = self.pool(torch.relu(self.conv1(x)))
x = self.pool(torch.relu(self.conv2(x)))
x = x.view(x.size(0), -1)
x = torch.relu(self.fc1(x))
return self.fc2(x)
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(50):
for images, labels in train_loader:
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()训练后评估与复盘：训练完成后，应记录以下关键指标：训练集准确率、验证集准确率、测试集准确率、损失曲线收敛情况、混淆矩阵以及每类别的精确率/召回率。若训练集准确率远高于验证集，说明存在过拟合，应增加数据增强强度或引入Dropout；若两者均低，则说明模型容量不足或数据质量有问题。对于STM32CubeAI部署，还需额外关注模型参数量和FLOPs，确保满足MCU的Flash和RAM限制。