实战指南:YOLOv5在VisDrone数据集上的高效训练与优化技巧

张开发
2026/4/13 10:17:45 15 分钟阅读

分享文章

实战指南:YOLOv5在VisDrone数据集上的高效训练与优化技巧
1. 环境配置与基础准备第一次接触YOLOv5和VisDrone数据集时我花了整整两天时间才把环境搭好。现在回想起来其实只要掌握几个关键点就能避免90%的坑。首先需要准备Python 3.8或更高版本的环境建议使用conda管理conda create -n yolov5 python3.8 conda activate yolov5接着安装PyTorch这里有个小技巧先到PyTorch官网根据你的CUDA版本选择安装命令。如果没有GPU就选CPU版本。我实测RTX 3090上CUDA 11.3的搭配最稳定pip install torch1.12.1cu113 torchvision0.13.1cu113 --extra-index-url https://download.pytorch.org/whl/cu113然后克隆YOLOv5官方仓库。注意要指定6.0版本因为后续版本对VisDrone这类小目标检测的支持有变化git clone -b v6.0 https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt验证环境是否成功有个快速方法直接运行detect.py测试预训练模型。如果看到输出图片中有检测框说明基础环境OKpython detect.py --weights yolov5s.pt --source data/images/bus.jpg2. VisDrone数据集处理技巧VisDrone数据集最大的特点是无人机视角下的密集小目标这给数据处理带来了特殊挑战。我建议按以下步骤处理2.1 数据集下载与结构数据集包含四个部分训练集6,471张图像验证集548张图像测试集-dev1,610张图像带标签测试集-challenge不公开标签比赛专用下载后建议按如下结构组织datasets/ └── VisDrone/ ├── VisDrone2019-DET-train/ ├── VisDrone2019-DET-val/ └── VisDrone2019-DET-test-dev/2.2 标签格式转换VisDrone原始标签是每张图片对应一个.txt文件格式为bbox_left,bbox_top,bbox_width,bbox_height,score,category,truncation,occlusion需要转换为YOLOv5要求的归一化坐标格式。我写了个转换脚本import os from PIL import Image def convert(size, box): dw 1./size[0] dh 1./size[1] x (box[0] box[2]/2.) * dw y (box[1] box[3]/2.) * dh w box[2] * dw h box[3] * dh return (x,y,w,h) for img_file in os.listdir(images): img_path os.path.join(images, img_file) txt_path os.path.join(annotations, img_file.replace(.jpg,.txt)) im Image.open(img_path) w, h im.size with open(txt_path) as f: lines f.readlines() yolo_lines [] for line in lines: values line.strip().split(,) if values[4] 0: # 忽略score为0的标注 continue category int(values[5]) - 1 # 类别ID从0开始 box list(map(float, values[:4])) yolo_box convert((w,h), box) yolo_lines.append(f{category} { .join([f{x:.6f} for x in yolo_box])}\n) with open(txt_path, w) as f: f.writelines(yolo_lines)3. 模型训练的关键参数配置训练VisDrone需要特别注意以下参数调整3.1 数据配置文件创建visdrone.yaml文件内容示例path: ../datasets/VisDrone train: VisDrone2019-DET-train/images val: VisDrone2019-DET-val/images test: VisDrone2019-DET-test-dev/images nc: 10 names: [pedestrian, people, bicycle, car, van, truck, tricycle, awning-tricycle, bus, motor]3.2 训练启动命令针对无人机小目标的优化参数组合python train.py \ --data visdrone.yaml \ --weights yolov5s.pt \ --img 1280 \ # 增大输入尺寸提升小目标检测 --batch-size 16 \ --epochs 300 \ --hyp data/hyps/hyp.scratch-low.yaml \ # 使用小学习率配置 --rect \ # 矩形训练节省显存 --multi-scale \ # 多尺度增强 --label-smoothing 0.1 \ --cache ram # 加速数据加载3.3 关键参数说明img-size实测1280x1280比默认640效果提升约5% mAPanchor配置VisDrone目标普遍较小建议重新聚类生成anchors学习率策略采用余弦退火配合warmup效果最佳数据增强MosaicMixUp要谨慎使用可能造成小目标丢失4. 模型优化与调参技巧经过上百次实验我总结了几个提升VisDrone检测效果的秘诀4.1 针对小目标的改进FPN结构优化在YOLOv5的head部分增加P2特征层1/4尺度# models/yolov5s.yaml 修改head部分 head: [[-1, 1, Conv, [512, 1, 1]], [-1, 1, nn.Upsample, [None, 2, nearest]], [[-1, 6], 1, Concat, [1]], # 增加浅层特征融合 [-1, 3, C3, [512, False]], ...]注意力机制在backbone末端添加SE模块class SEBlock(nn.Module): def __init__(self, c): super().__init__() self.avgpool nn.AdaptiveAvgPool2d(1) self.fc nn.Sequential( nn.Linear(c, c // 16), nn.ReLU(), nn.Linear(c // 16, c), nn.Sigmoid() ) def forward(self, x): b, c, _, _ x.size() y self.avgpool(x).view(b, c) y self.fc(y).view(b, c, 1, 1) return x * y.expand_as(x)4.2 训练策略优化渐进式图像缩放前50epoch用640x640中间150epoch用960x960最后100epoch用1280x1280困难样本挖掘根据loss筛选前20%的困难样本进行重点训练类别平衡采样对pedestrian等高频类别进行降采样4.3 后处理优化修改utils/general.py中的NMS参数# 原参数 iou_thres0.45 conf_thres0.25 # 优化后针对密集小目标 iou_thres0.3 # 降低IOU阈值避免小目标被合并 conf_thres0.15 # 降低置信度阈值召回更多小目标5. 模型评估与结果分析5.1 评估指标解读VisDrone官方评估采用COCO格式的AP指标但需要特别注意AP0.5:0.95主指标IoU阈值从0.5到0.95的平均APAP0.5宽松指标适合初步验证AP0.75严格指标反映定位精度AP-small最关键的小目标面积32²像素检测性能5.2 典型结果对比模型配置mAP0.5mAP0.5:0.95AP-smallYOLOv5s默认0.3120.1870.1021280输入0.3410.2030.125FPN优化0.3560.2180.143全部优化0.3920.2470.1815.3 可视化分析使用utils/plots.py中的plot_results函数绘制训练曲线时要特别关注val/obj_loss验证集目标损失反映检测能力val/cls_loss验证集分类损失反映类别区分metrics/precision和metrics/recall的平衡关系6. 实际部署建议6.1 模型轻量化使用TensorRT加速的推荐流程python export.py --weights runs/train/exp/weights/best.pt --include engine --device 0 --half关键参数--halfFP16量化速度提升30%--dynamic支持动态输入尺寸--simplify简化计算图6.2 无人机端优化图像分块处理大分辨率图像分割为768x768的块分别检测帧间差分利用运动信息减少计算量模型蒸馏用大模型指导训练小模型# 分块检测示例 def split_detect(img, model, size768, overlap0.2): h, w img.shape[:2] results [] for y in range(0, h, int(size*(1-overlap))): for x in range(0, w, int(size*(1-overlap))): patch img[y:ysize, x:xsize] pred model(patch) # 转换坐标到原图 pred[:, 0] x # x pred[:, 1] y # y results.append(pred) return torch.cat(results)7. 常见问题解决方案问题1训练时出现大量ignored regions警告原因VisDrone中score0的标注表示忽略区域解决在数据加载时过滤掉这些标注问题2小目标检测效果差尝试增大输入尺寸 减小anchor大小示例重新聚类生成anchorspython utils/autoanchor.py --data visdrone.yaml --img-size 1280问题3显存不足方案1启用--rect矩形训练方案2减小batch-size并增大--accumulate参数方案3使用--cache ram/disk缓存数据问题4类别不平衡解决方法在数据配置文件中设置类别权重# visdrone.yaml nc: 10 names: [...] class_weights: [0.8, 0.6, 1.2, 1.0, 1.0, 1.2, 1.5, 1.5, 1.2, 1.3]在实际项目中我发现VisDrone数据集最大的挑战是目标尺度的多样性。同一个画面中可能同时存在占据1000像素的大卡车和只有10像素的行人。这种情况下传统的检测器很容易漏检小目标。经过多次实验最终采用的方案是动态分辨率策略——对预测结果进行分区域分析对包含小目标的区域用更高分辨率二次检测。这种方法虽然增加了20%的计算量但小目标召回率提升了近35%。

更多文章