Yolov5自定义图片训练测试及模型调优（详细过程）

一、数据准备与标注规范

1.1 数据集结构构建

自定义训练需遵循Yolov5标准目录结构：

custom_dataset/
├── images/
│   ├── train/      # 训练集图片
│   ├── val/        # 验证集图片
│   └── test/       # 测试集图片（可选）
└── labels/
    ├── train/      # 训练集标签
    └── val/        # 验证集标签

建议按71比例划分数据集，图片与标签文件需保持同名（如img1.jpg对应img1.txt）。

1.2 标注文件格式

Yolov5采用YOLO格式的txt标注文件，每行包含：
<class_id> <x_center> <y_center> <width> <height>
所有坐标需归一化到[0,1]区间。例如：

0 0.542 0.678 0.215 0.332  # 类0，中心点(0.542,0.678)，宽高占图像比例

推荐使用LabelImg、CVAT等工具进行标注，完成后需通过脚本验证标注质量：

import os
def check_annotations(label_dir):
    for file in os.listdir(label_dir):
        if file.endswith('.txt'):
            with open(os.path.join(label_dir, file), 'r') as f:
                lines = f.readlines()
                for line in lines:
                    parts = line.strip().split()
                    if len(parts) != 5:
                        print(f"Error in {file}: Invalid format")
                    # 可添加更多校验逻辑

二、模型配置与训练参数

2.1 配置文件修改

在data/目录下创建自定义数据集配置文件（如custom.yaml）：

# custom.yaml示例
train: ../custom_dataset/images/train
val: ../custom_dataset/images/val
nc: 3  # 类别数量
names: ['class1', 'class2', 'class3']  # 类别名称

2.2 模型选择策略

Yolov5提供五种规模模型：

• Yolov5s：最快（3.7ms/img），适合嵌入式设备
• Yolov5m：平衡型（5.6ms/img）
• Yolov5l：高精度（7.8ms/img）
• Yolov5x：最高精度（11.7ms/img）
• Yolov5n：纳米版（1.4ms/img）

根据硬件条件选择：

# GPU训练示例（使用Yolov5s）
python train.py --img 640 --batch 16 --epochs 100 \
               --data custom.yaml --weights yolov5s.pt \
               --name custom_train --cache ram

2.3 关键训练参数

参数	说明	推荐值
--img	输入分辨率	640（平衡速度精度）
--batch	批处理大小	根据GPU显存调整（16-64）
--epochs	训练轮次	100-300（小数据集可减少）
--lr0	初始学习率	0.01（默认）
--lrf	学习率衰减系数	0.01（默认）
--weight-decay	权重衰减	0.0005（防止过拟合）

三、训练过程监控与优化

3.1 实时监控指标

训练日志包含关键指标：

• box_loss：边界框回归损失
• obj_loss：目标存在性损失
• cls_loss：分类损失
• mAP@0.5：IoU=0.5时的平均精度

建议每10个epoch保存一次权重：

# 在train.py中添加检查点保存逻辑
if epoch % 10 == 0:
    torch.save(model.state_dict(), f'runs/train/exp/weights/best_{epoch}.pt')

3.2 常见问题处理

问题1：训练初期loss波动大

• 解决方案：降低初始学习率至0.001，或使用学习率预热

问题2：验证集mAP停滞

• 解决方案：
  • 增加数据增强（HSV调整、随机缩放）
  • 调整 --patience 参数（默认50，可减至20）
  • 尝试混合精度训练（--half）

问题3：GPU显存不足

• 解决方案：
  • 减小--batch-size（如从16降至8）
  • 使用梯度累积（--accumulate）
  • 启用--evolve进行超参数搜索

四、模型测试与评估

4.1 测试命令示例

python val.py --data custom.yaml --weights runs/train/exp/weights/best.pt \
             --img 640 --task val --half --save-txt --save-conf

4.2 评估指标解析

• P/R曲线：精确率-召回率曲线
• F1分数：最佳阈值下的综合指标
• 混淆矩阵：分析类别误判情况

生成可视化报告：

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import seaborn as sns
# 假设已有真实标签和预测标签
cm = confusion_matrix(true_labels, pred_labels)
plt.figure(figsize=(10,8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.xlabel('Predicted')
plt.ylabel('True')
plt.show()

五、模型调优实战技巧

5.1 超参数优化

使用遗传算法进行自动调参：

python train.py --evolve 300 --img 640 --batch 16 \
               --data custom.yaml --weights yolov5s.pt

生成的hyp_evolved.yaml包含优化后的参数组合。

5.2 模型剪枝与量化

剪枝操作（减少30%-50%参数）：

from models.experimental import attempt_load
model = attempt_load('best.pt', map_location='cpu')
pruned_model = torch.nn.utils.prune.global_unstructured(
    model, pruning_method=torch.nn.utils.prune.L1Unstructured, 
    amount=0.3)

量化操作（INT8精度）：

python export.py --weights best.pt --include torchscript,onnx,coreml \
                --img 640 --optimize

5.3 部署优化建议

1. TensorRT加速：

   trtexec --onnx=model.onnx --saveEngine=model.engine --fp16

2. OpenVINO优化：

   mo --framework onnx --input_model model.onnx \
   --output_dir optimized --data_type FP16

六、完整工作流程示例

1. 数据准备：

   # 示例：使用脚本自动划分数据集
   python -c "
   import os, shutil, random
   src = 'raw_data'
   dst = 'custom_dataset'
   files = [f for f in os.listdir(src) if f.endswith('.jpg')]
   random.shuffle(files)
   train, val = files[:int(0.7*len(files))], files[int(0.7*len(files)):]
   for f in train: shutil.copy(os.path.join(src,f), os.path.join(dst,'images','train',f))
   for f in val: shutil.copy(os.path.join(src,f), os.path.join(dst,'images','val',f))
   "

2. 启动训练：

   python train.py --img 640 --batch 16 --epochs 200 \
               --data custom.yaml --weights yolov5s.pt \
               --name industrial_defect --cache ram \
               --hyp hyp.scratch-low.yaml

3. 结果分析：

   import pandas as pd
   results = pd.read_csv('runs/train/industrial_defect/results.csv')
   print(results[['epoch', 'metrics/precision', 'metrics/mAP_0.5']].tail())

七、进阶优化方向

1. 多尺度训练：在配置文件中添加：

   # 在data/custom.yaml中添加
   multi_scale: True  # 启用随机尺度训练

2. 注意力机制集成：
   修改models/yolov5s.yaml，在backbone中添加：

   # 在C3模块后添加注意力
   [[-1, 1, CBAM, [256]]]  # CBAM注意力模块

3. 知识蒸馏：
   使用大模型指导小模型训练：

   # 在loss计算中添加蒸馏损失
   teacher_output = teacher_model(images)
   distill_loss = F.mse_loss(student_output, teacher_output)
   total_loss = box_loss + obj_loss + cls_loss + 0.5*distill_loss

通过系统化的训练流程和针对性的优化策略，Yolov5在自定义数据集上的表现可显著提升。实际项目中，建议从Yolov5s模型开始，逐步优化数据质量、调整超参数，最后进行模型压缩，以实现精度与速度的最佳平衡。