1. 工业小数据集下的YOLOv6迁移学习实战
在工业质检领域，我们常常面临样本不足的困境。最近在一个齿轮缺陷检测项目中，客户只提供了1200张标注图像，但要求检测精度mAP50必须达到0.85以上。经过多次实验验证，我总结出一套基于YOLOv6的30轮快速收敛方案，最终在测试集上达到了0.88的mAP50。下面分享完整实现细节。
2. 核心方案设计思路
2.1 为什么选择迁移学习？
当训练数据少于5000张时，从头训练(target-from-scratch)会导致严重的过拟合。我们做过对比实验：
- 从头训练：30轮后训练集mAP50=0.95，验证集仅0.62
- 迁移学习：30轮后验证集稳定在0.88+
关键差异在于预训练模型已经具备：
- 通用物体特征提取能力（边缘、纹理等）
- 多尺度感知能力（COCO数据集中包含各种尺寸物体）
2.2 数据增强策略设计
针对1200张的小数据集，我们采用三级增强策略：
基础增强（YOLOv6内置）
python复制# configs/data/yolov6s_finetune.py
img_aug = {
'hsv_h': 0.015,
'hsv_s': 0.7,
'hsv_v': 0.4,
'degrees': 5.0,
'translate': 0.1,
'scale': 0.9,
'shear': 0.0,
'flip_lr': 0.5,
'mosaic': 1.0,
'mixup': 0.1
}
自定义增强（Albumentations扩展）
python复制import albumentations as A
extra_aug = A.Compose([
A.RandomShadow(p=0.3),
A.GridDropout(ratio=0.2, p=0.5),
A.RandomSunFlare(p=0.1)
], bbox_params=A.BboxParams(format='yolo'))
样本重平衡
对缺陷样本进行oversampling：
- 正常齿轮：800张 → 保持800
- 裂纹缺陷：200张 → 复制到400
- 缺齿缺陷：200张 → 复制到400
注意：增强后的样本需要重新计算anchor尺寸，我们使用k-means++对增强后数据重新聚类得到：
anchors = [[12,16], [19,36], [40,28], [36,75], [76,55]]
3. 模型训练关键配置
3.1 网络结构微调
保留YOLOv6s的主干网络（backbone），仅调整检测头：
yaml复制# yolov6s_finetune.yaml
head:
num_layers: 2 # 原为3，减少复杂度
in_channels: [128, 256, 512]
feat_channels: 128 # 原为256
3.2 分阶段训练策略
阶段一：冻结训练（第1-10轮）
python复制# 冻结backbone参数
for param in model.backbone.parameters():
param.requires_grad = False
optimizer = SGD([
{'params': model.head.parameters(), 'lr': 0.001}
], momentum=0.9)
阶段二：全网络微调（第11-30轮）
python复制# 解冻所有参数
for param in model.parameters():
param.requires_grad = True
optimizer = SGD([
{'params': model.backbone.parameters(), 'lr': 0.0001},
{'params': model.head.parameters(), 'lr': 0.001}
], momentum=0.9)
3.3 学习率调度
采用余弦退火+热启动：
python复制lr_scheduler = CosineAnnealingWarmRestarts(
optimizer,
T_0=5, # 5轮为一个周期
T_mult=2,
eta_min=1e-5
)
4. 实测效果与调优记录
4.1 训练过程监控
| 训练轮次 | 学习率 | 训练mAP50 | 验证mAP50 |
|---|---|---|---|
| 1-5 | 0.001 | 0.65 | 0.63 |
| 6-10 | 0.001 | 0.82 | 0.79 |
| 11-15 | 0.0005 | 0.88 | 0.85 |
| 16-20 | 0.0002 | 0.91 | 0.87 |
| 21-25 | 0.0001 | 0.93 | 0.88 |
| 26-30 | 0.00005 | 0.94 | 0.88 |
4.2 关键调参经验
-
批次大小选择：
- 显存8G：batch_size=16
- 显存16G：batch_size=32
- 显存24G：batch_size=48
-
早停策略配置：
python复制early_stopping = EarlyStopping( patience=5, # 连续5轮验证集mAP不提升则停止 delta=0.001 # 最小改善阈值 )
5. 常见问题解决方案
5.1 过拟合现象处理
症状：训练集指标持续上升，验证集指标波动或下降
解决方案：
- 增加MixUp比例（0.1→0.3）
- 添加CutOut增强
- 在检测头添加Dropout层（rate=0.2）
5.2 显存溢出排查
当出现CUDA out of memory时：
- 减小batch_size（32→16）
- 关闭mosaic增强（最后一轮训练时）
- 使用梯度累积：
python复制# 每4个batch更新一次参数 optimizer.step_every = 4
5.3 小目标检测优化
对于小于32x32像素的缺陷：
- 修改anchor尺寸：
yaml复制anchors: [[8,10], [13,16], [19,24]] - 增加P2特征层（高分辨率检测头）
- 使用BiFPN替换原FPN
6. 工程部署建议
6.1 模型量化方案
python复制from pytorch_quantization import quant_modules
quant_modules.initialize()
model_fp32 = load_model("yolov6s_finetune.pt")
model_int8 = torch.quantization.convert(model_fp32)
量化后模型大小从45MB降至11MB，推理速度提升2.3倍。
6.2 TensorRT加速
bash复制trtexec --onnx=yolov6s.onnx \
--saveEngine=yolov6s.engine \
--fp16 \
--workspace=4096
在Jetson Xavier NX上实测：
- FP32: 45 FPS
- FP16: 78 FPS
- INT8: 115 FPS
7. 项目完整代码结构
code复制yolov6_finetune/
├── configs/
│ ├── data/
│ │ └── yolov6s_finetune.py
│ └── model/
│ └── yolov6s_finetune.yaml
├── data/
│ ├── images/ # 原始图像
│ ├── labels/ # YOLO格式标注
│ └── augmented/ # 增强后数据
├── tools/
│ ├── train.py # 训练脚本
│ └── augment.py # 数据增强脚本
└── README.md # 完整参数说明
这套方案已在三个工业场景验证：
- 齿轮缺陷检测（1200张，mAP50=0.88）
- PCB板检测（800张，mAP50=0.85）
- 纺织品瑕疵检测（1500张，mAP50=0.89）
关键是要根据具体场景调整数据增强策略，比如PCB检测需要增加仿射变换来模拟不同拍摄角度，纺织品检测则需要强化纹理相关的增强。