YOLO26 训练实战:自定义数据集全流程

1. 环境准备

# 安装 Ultralytics
pip install ultralytics==8.3.0

# 验证
python -c "from ultralytics import YOLO; print('OK')"

2. 数据集准备

2.1 数据标注

# 使用 Roboflow 在线标注
# 或使用 LabelImg 本地标注
pip install labelImg
labelImg  # 启动 GUI

# 标注格式(YOLO):
# class_id x_center y_center width height
# 0 0.5 0.5 0.3 0.4

2.2 数据集结构

dataset/
├── train/
│   ├── images/
│   │   ├── 001.jpg
│   │   └── 002.jpg
│   └── labels/
│       ├── 001.txt
│       └── 002.txt
├── valid/
│   ├── images/
│   └── labels/
└── data.yaml
# data.yaml
path: ./dataset
train: train/images
val: valid/images
test: test/images

names:
  0: helmet
  1: person
  2: no-helmet

nc: 3

3. 模型训练

3.1 基础训练

#!/usr/bin/env python3
"""train_yolo26.py"""
from ultralytics import YOLO

# 加载预训练模型
model = YOLO("yolo26s.pt")

# 训练
results = model.train(
    data="data.yaml",
    epochs=100,
    imgsz=640,
    batch=16,
    device="0",
    workers=8,
    patience=20,
    save=True,
    save_period=10,
    project="runs/detect",
    name="custom_yolo26s",
    exist_ok=True,
    pretrained=True,
    optimizer="auto",
    lr0=0.01,
    lrf=0.01,
    momentum=0.937,
    weight_decay=0.0005,
    warmup_epochs=3,
    warmup_momentum=0.8,
    warmup_bias_lr=0.1,
    box=7.5,
    cls=0.5,
    dfl=1.5,
    plots=True,
)

print(f"最佳 mAP50: {results.best_map50:.4f}")
print(f"最佳 mAP50-95: {results.best_map50_95:.4f}")

3.2 选择模型尺寸

模型选择建议:
┌─────────────┬──────────┬──────────┬──────────┐
│ 场景         │ 推荐模型  │ 训练时间  │ 精度      │
├─────────────┼──────────┼──────────┼──────────┤
│ 快速验证     │ YOLO26n  │ 2h       │ mAP 38+  │
│ 通用检测     │ YOLO26s  │ 6h       │ mAP 45+  │
│ 高精度      │ YOLO26m  │ 12h      │ mAP 51+  │
│ 极致精度     │ YOLO26l  │ 24h      │ mAP 53+  │
└─────────────┴──────────┴──────────┴──────────┘
* 训练时间基于 1000 张图片、RTX 4090

4. 超参数调优

# 超参数搜索
from ultralytics import YOLO

model = YOLO("yolo26s.pt")

# 使用 Ray Tune 超参数搜索
model.tune(
    data="data.yaml",
    epochs=30,
    iterations=100,
    optimizer="AdamW",
    plots=True,
    save=True,
    val=True,
)

4.1 关键超参数

# 超参数说明
lr0: 0.01          # 初始学习率
lrf: 0.01          # 最终学习率 = lr0 * lrf
momentum: 0.937    # SGD 动量
weight_decay: 0.0005  # 权重衰减
warmup_epochs: 3   # 预热轮数
box: 7.5           # 边界框损失权重
cls: 0.5           # 分类损失权重
dfl: 1.5           # DFL 损失权重
hsv_h: 0.015       # 色调增强
hsv_s: 0.7         # 饱和度增强
hsv_v: 0.4         # 亮度增强
degrees: 5.0       # 旋转角度
translate: 0.1     # 平移比例
scale: 0.5         # 缩放比例
mosaic: 1.0        # Mosaic 概率
mixup: 0.15        # MixUp 概率

5. 模型评估

#!/usr/bin/env python3
"""evaluate.py"""
from ultralytics import YOLO

model = YOLO("runs/detect/custom_yolo26s/weights/best.pt")

# 在测试集上评估
metrics = model.val(
    data="data.yaml",
    split="test",
    imgsz=640,
    batch=1,
    conf=0.25,
    iou=0.6,
)

print(f"mAP50: {metrics.box.map50:.4f}")
print(f"mAP50-95: {metrics.box.map:.4f}")
print(f"Precision: {metrics.box.mp:.4f}")
print(f"Recall: {metrics.box.mr:.4f}")

# 每类 AP
for i, name in metrics.names.items():
    print(f"  {name}: AP50={metrics.box.ap50[i]:.4f}")

6. 模型导出

#!/usr/bin/env python3
"""export.py"""
from ultralytics import YOLO

model = YOLO("runs/detect/custom_yolo26s/weights/best.pt")

# 导出 ONNX
model.export(format="onnx", imgsz=640, opset=11, simplify=True)

# 导出 TensorRT
model.export(format="engine", imgsz=640, half=True, batch=1)

# 导出 TFLite
model.export(format="tflite", imgsz=640, int8=True)

7. 推理测试

#!/usr/bin/env python3
"""predict.py"""
from ultralytics import YOLO

model = YOLO("runs/detect/custom_yolo26s/weights/best.pt")

# 单张图片
results = model.predict("test.jpg", conf=0.3, save=True)

# 视频
results = model.predict("video.mp4", conf=0.3, save=True, stream=True)
for r in results:
    print(f"检测到 {len(r.boxes)} 个目标")

# 摄像头
results = model.predict(source=0, conf=0.3, show=True)

8. 训练技巧

训练优化清单:
├── 数据
│   ├── 数据量:每个类别 500+ 张
│   ├── 标注质量:边界框紧贴目标
│   ├── 类别平衡:各类别数量相近
│   └── 多样性:覆盖各种场景/光照/角度
├── 模型
│   ├── 先小后大:先用 YOLO26n 验证数据质量
│   ├── 预训练:使用 COCO 预训练权重
│   └── 冻结层:冻结前 10 层 backbone
├── 训练
│   ├── 早停:patience=20
│   ├── 学习率:cosine 调度
│   ├── 数据增强:Mosaic + MixUp
│   └── 混合精度:AMP 开启
└── 评估
    ├── mAP50:粗评估
    ├── mAP50-95:精评估
    ├── PR 曲线:分析阈值
    └── 混淆矩阵:分析误检

总结

核心步骤:

  1. 数据准备:标注 + 增强 + 格式化
  2. 模型选择:YOLO26n 验证 → YOLO26s/m 正式训练
  3. 超参数调优:学习率、损失权重、增强策略
  4. 评估验证:mAP50-95 > 目标阈值
  5. 导出部署:ONNX → TensorRT/TFLite/NCNN
Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐