YOLOv5实战篇---从零构建VOC格式数据集,高效训练自定义模型
1. 环境准备与YOLOv5源码获取
在开始构建VOC格式数据集之前,我们需要先准备好开发环境。YOLOv5基于PyTorch框架,因此需要安装Python 3.7及以上版本。我推荐使用conda创建虚拟环境,避免与其他项目产生依赖冲突:
conda create -n yolov5 python=3.8
conda activate yolov5
接下来克隆Ultralytics官方仓库。这里建议使用国内镜像加速下载,避免直接从GitHub拉取时可能遇到的网络问题:
git clone https://gitee.com/mirrors/yolov5.git
cd yolov5
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
安装完成后,建议测试环境是否正常。运行以下命令会下载预训练模型并对示例图片进行检测:
python detect.py --weights yolov5s.pt --source data/images/bus.jpg
如果看到终端输出了检测结果,并且在runs/detect/exp目录下生成了标注后的图片,说明环境配置成功。这里特别提醒,如果使用GPU训练,需要确保CUDA和cuDNN已正确安装。可以通过 nvidia-smi 命令查看GPU状态,以及 torch.cuda.is_available() 验证PyTorch是否能识别到GPU。
2. VOC数据集结构解析与创建
VOC格式是目标检测领域最常用的数据集格式之一,其目录结构有严格规范。我们先来看标准VOC数据集的目录树:
VOCdevkit/
└── VOC2007/
├── Annotations/ # 存放XML标注文件
├── JPEGImages/ # 存放原始图像
├── ImageSets/
│ └── Main/ # 存放数据集划分文件
└── labels/ # YOLO格式标签(转换后生成)
实际操作中,我建议在yolov5目录下创建datasets文件夹集中管理数据。以下是创建VOC格式数据集的具体步骤:
- 创建基础目录结构:
mkdir -p datasets/VOCdevkit/VOC2007/{Annotations,JPEGImages,ImageSets/Main}
-
将收集的图片放入JPEGImages目录。这里有个实用技巧:可以使用
mogrify -format jpg *.png命令批量转换PNG为JPG格式 -
使用labelImg工具标注图片。安装labelImg只需执行:
pip install labelImg
labelImg # 启动标注工具
标注时建议开启"自动保存"模式,标注文件会默认保存为PASCAL VOC格式的XML文件。每个XML文件包含对应图片中所有目标的类别和边界框信息。我遇到过标注文件损坏的情况,建议定期备份Annotations文件夹。
3. 自动化标注与数据增强技巧
手动标注费时费力,我们可以利用现代AI技术实现半自动化标注。最新推出的YOLO-World模型支持开放词汇检测,只需提供类别文本描述就能生成初步标注:
from yoloworld import YOLOWorld
model = YOLOWorld(model_id='yolo_world/l')
classes = ["person", "car", "dog"] # 自定义类别
model.set_classes(classes)
results = model.predict("image.jpg")
results.save_txt("labels/image.txt") # 保存为YOLO格式
对于已有部分标注的数据,Roboflow平台提供智能标注建议功能。上传数据后,其Active Learning算法会优先推荐最有价值的样本进行标注,能减少30%以上的标注工作量。
数据增强方面,YOLOv5内置了Mosaic9等高级增强策略。在data/hyps/hyp.scratch-low.yaml中可以调整增强参数:
hsv_h: 0.015 # 色调增强强度
hsv_s: 0.7 # 饱和度增强范围
hsv_v: 0.4 # 明度增强范围
degrees: 10.0 # 旋转角度范围
translate: 0.1 # 平移比例
scale: 0.9 # 缩放范围
shear: 2.0 # 剪切强度
实测发现,适当增强能提升模型鲁棒性,但过度增强反而会降低精度。建议初期保持默认参数,后续再逐步调整。
4. 数据集划分与格式转换
我们需要将数据集划分为训练集、验证集和测试集。创建split_train_val.py脚本实现自动划分:
import os
import random
def split_dataset(xml_path, output_path, train_val_ratio=0.9, train_ratio=0.9):
xml_files = [f[:-4] for f in os.listdir(xml_path) if f.endswith('.xml')]
random.shuffle(xml_files)
total = len(xml_files)
train_val = int(total * train_val_ratio)
train = int(train_val * train_ratio)
with open(f"{output_path}/train.txt", 'w') as f:
f.writelines([x+'\n' for x in xml_files[:train]])
with open(f"{output_path}/val.txt", 'w') as f:
f.writelines([x+'\n' for x in xml_files[train:train_val]])
with open(f"{output_path}/test.txt", 'w') as f:
f.writelines([x+'\n' for x in xml_files[train_val:]])
运行脚本后,接下来需要将VOC格式的XML标注转换为YOLO格式的txt文件。创建voc2yolo.py转换脚本:
import xml.etree.ElementTree as ET
import os
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def convert_annotation(xml_file, classes):
in_file = open(xml_file)
out_file = open(xml_file.replace('Annotations','labels').replace('.xml','.txt'), 'w')
tree = ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
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('xmax').text),
float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert((w,h), b)
out_file.write(f"{cls_id} {' '.join([str(a) for a in bb])}\n")
转换完成后,每个图片会对应一个txt文件,内容格式为 class_id x_center y_center width height ,所有坐标都是归一化后的相对值。
5. 配置文件定制与模型训练
在data目录下创建自定义数据集的YAML配置文件,例如data/mydata.yaml:
train: ./datasets/VOCdevkit/VOC2007/train.txt
val: ./datasets/VOCdevkit/VOC2007/val.txt
test: ./datasets/VOCdevkit/VOC2007/test.txt
# 类别数量
nc: 3
# 类别名称
names: ['person', 'car', 'dog']
选择模型配置时,对于新手建议从yolov5s开始。修改models/yolov5s.yaml中的nc参数为你的类别数:
# parameters
nc: 3 # 修改为你的类别数
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
开始训练前,建议先使用预训练权重进行迁移学习:
python train.py --img 640 --batch 16 --epochs 100 --data data/mydata.yaml \
--cfg models/yolov5s.yaml --weights yolov5s.pt --device 0
训练过程中常见问题及解决方案:
- CUDA内存不足:减小batch-size(可设为8或4)
- 损失不下降:检查标注是否正确,适当减小学习率
- 过拟合:增加数据增强,使用早停策略
6. 高级训练技巧与性能优化
对于大规模数据集,可以采用以下优化策略:
- 多GPU训练:添加
--device 0,1参数使用多卡 - 混合精度训练:
--half参数可减少显存占用 - 数据缓存:
--cache ram/disk加速数据加载 - 自动批处理:
--batch-size -1让YOLOv5自动选择最佳batch
使用DDP模式进行分布式训练示例:
python -m torch.distributed.run --nproc_per_node 2 train.py \
--img 640 --batch 64 --epochs 100 --data data/mydata.yaml \
--weights yolov5s.pt --device 0,1
训练监控方面,除了默认的TensorBoard,还可以集成Weights & Biases:
pip install wandb
wandb login # 按提示输入API key
python train.py ... --bbox_interval 1 # 每epoch记录验证结果
对于边缘设备部署,建议使用模型剪枝和量化:
from torch.quantization import quantize_dynamic
model = torch.load('yolov5s.pt')
model_quantized = quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
torch.save(model_quantized, 'yolov5s_quant.pt')
7. 模型评估与部署应用
训练完成后,使用test.py评估模型性能:
python test.py --weights runs/train/exp/weights/best.pt \
--data data/mydata.yaml --task test --augment
关键指标解读:
- mAP@0.5:IoU阈值为0.5时的平均精度
- mAP@0.5:0.95:IoU阈值从0.5到0.95的平均精度
- precision:精确率(预测为正样本中真实正样本比例)
- recall:召回率(真实正样本中被正确预测的比例)
部署到生产环境时,建议导出为ONNX格式:
python export.py --weights best.pt --include onnx --img 640 --dynamic
在Python中调用训练好的模型进行推理:
import torch
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
results = model('image.jpg')
results.show() # 显示结果
results.pandas().xyxy[0] # 获取DataFrame格式的检测结果
对于需要实时处理的场景,可以使用Triton Inference Server部署模型,实现高并发推理。我在实际项目中测试,Triton可以将推理吞吐量提升3-5倍。
更多推荐


所有评论(0)