[卷积神经网络]YOLOv11实战:从零构建自定义数据集训练流程
1. 数据集准备与格式转换
第一次用YOLOv11训练自己的模型时,最头疼的就是数据集处理。我的数据是VOC格式的,但YOLO需要特定的文本标注格式。这里分享一个我验证过的完整转换方案,包含几个容易踩坑的细节。
VOC转YOLO格式的核心是坐标归一化处理。VOC使用绝对坐标(xmin,ymin,xmax,ymax),而YOLO需要转换为相对坐标(x_center,y_center,width,height)。我写了个Python脚本自动处理:
import xml.etree.ElementTree as ET def convert_box(size, box): dw = 1./size[0] dh = 1./size[1] x = (box[0] + box[2])/2.0 y = (box[1] + box[3])/2.0 w = box[2] - box[0] h = box[3] - box[1] x = x * dw w = w * dw y = y * dh h = h * dh return (x,y,w,h) def convert_annotation(xml_file, txt_file, classes): tree = ET.parse(xml_file) root = tree.getroot() size = root.find('size') w = int(size.find('width').text) h = int(size.find('height').text) with open(txt_file, 'w') as f: for obj in root.iter('object'): cls = obj.find('name').text if cls not in classes: continue cls_id = classes.index(cls) xmlbox = obj.find('bndbox') b = (float(xmlbox.find('xmin').text), float(xmlbox.find('ymin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymax').text)) bb = convert_box((w,h), b) f.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n")实际使用时要注意三点:
- 类别顺序必须固定,建议用列表不要用字典
- 图像尺寸要从XML里读取,不能假设都是统一尺寸
- 路径处理建议用os.path.join,避免跨平台问题
2. 数据集划分策略
数据集划分直接影响模型效果。我习惯用8:1:1的比例分割训练集、验证集和测试集。这里推荐使用sklearn的train_test_split分层抽样:
from sklearn.model_selection import train_test_split def split_dataset(image_dir, label_dir, test_size=0.2): images = [f for f in os.listdir(image_dir) if f.endswith('.jpg')] labels = [f.replace('.jpg', '.txt') for f in images] # 先分训练集和临时集 train_img, temp_img, train_lbl, temp_lbl = train_test_split( images, labels, test_size=test_size, random_state=42) # 再分验证集和测试集 val_img, test_img, val_lbl, test_lbl = train_test_split( temp_img, temp_lbl, test_size=0.5, random_state=42) return { 'train': (train_img, train_lbl), 'val': (val_img, val_lbl), 'test': (test_img, test_lbl) }关键点:
- 保持图像和标注文件同步分割
- 随机种子固定保证可复现
- 测试集要完全隔离,只在最终评估使用
3. 配置文件定制
YOLOv11需要两个核心配置文件:数据配置和模型配置。数据配置示例:
# VOC.yaml path: ../datasets/VOC train: images/train val: images/val test: images/test nc: 6 # 类别数 names: ['person', 'car', 'dog', 'cat', 'bicycle', 'motorcycle']模型配置我推荐从官方yolo11s.yaml开始修改:
# yolo11-custom.yaml nc: 6 # 必须与数据配置一致 scales: s: depth: 0.33 width: 0.50 max_channels: 1024 backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 # ...保持其他默认配置特别注意:
- 修改nc后要同步调整最后一层Detect的参数
- 小数据集建议用s版本,减少过拟合风险
- 输入尺寸保持640x640除非有特殊需求
4. 训练过程优化
启动训练的核心代码很简单:
from ultralytics import YOLO model = YOLO('yolo11-custom.yaml') results = model.train( data='VOC.yaml', epochs=100, imgsz=640, batch=16, device='0' )但有几个提升效果的关键技巧:
- 学习率预热:添加
warmup_epochs=3参数 - 早停机制:
patience=10当验证集指标不再提升时停止 - 数据增强:默认的augment已经很充分,不建议新手修改
- 混合精度:
amp=True能显著减少显存占用
训练过程中要重点关注三个指标:
- train/box_loss:检测框回归损失
- train/cls_loss:分类损失
- val/mAP@0.5:验证集平均精度
5. 模型评估与导出
训练完成后,用测试集进行最终评估:
model = YOLO('runs/detect/train/weights/best.pt') metrics = model.val( data='VOC.yaml', split='test', conf=0.5 # 置信度阈值 ) print(f"mAP@0.5: {metrics.box.map}")模型导出为ONNX格式便于部署:
model.export(format='onnx', dynamic=True)常见问题处理:
- 如果显存不足:减小batch_size或imgsz
- 如果过拟合:增加数据增强或减少模型规模
- 如果训练不稳定:检查数据标注质量
6. 实际应用建议
在工业项目中,我发现这些实践特别有用:
- 使用wandb或tensorboard记录训练过程
- 对困难样本进行针对性数据增强
- 尝试不同anchor box设置(尤其非常规目标)
- 测试时使用TTA(Test Time Augmentation)提升稳定性
对于边缘设备部署,建议:
- 导出时进行量化(
int8) - 使用TensorRT加速
- 对输出做后处理过滤
