资讯详情

VisDrone数据集实战:从YOLOv5训练到TensorRT部署全流程

📅 2026/9/24 18:50:54 | 华诺云谱 👁 阅读
VisDrone数据集实战:从YOLOv5训练到TensorRT部署全流程
简介面向无人机俯视视角下的车辆与行人检测任务一套包含一千多张真实航拍场景图片的数据集已按标准划分为train、val、test三个子集并配套data.yaml配置可直接适用于YOLOv5、YOLOv7、YOLOv8等主流目标检测算法。压缩包共两千个文件其中1887个txt文件为YOLO格式的标注信息112张jpg为原始图片另有1个yaml文件用于指定类别与路径包体整体约960.43MB目录划分清晰训练、验证、测试集均直接可用省去自行采集、清洗、标注和划分的大量时间。对于计算机视觉开发者和算法工程师该数据集既能用于无人机场景下的车辆与行人检测模型训练也可作为算法效果验证与调优的基准数据。目前已有1325人学习使用特别适合需要快速构建无人机视野目标检测方案的研究者与开发者。1. 无人机视角的目标检测数据集为什么这份 vis-drone 值得复现一遍做无人机俯视视角下的车辆和行人目标检测时我最先遇到的不是模型选型而是“数据集怎么处理”。原版 VisDrone 是视频抽帧标注格式从坐标到类别定义都和 YOLO 不一致下载下来还得自己写转换脚本、按比例切 train/val/test。这份 vis-drone-yolov5-dataset-2 解压后就是现成的训练目录data.yaml 已经写好names: [car, person]train、valid、test 三个划分目录直接可用YOLOv5、YOLOv7、YOLOv8 都能直接用这套目录配 train 命令。适合两类人一是新手想走通“数据集到模型”的完整流程二是做车辆/行人计数或无人机巡检的老手拿它当基准数据验证模型结构、增强策略或部署方案。2. 数据集结构和标注格式train/valid/test 目录与 data.yaml 的关系2.1 解压后的目录结构为什么 data.yaml 里写的是 valid 而非 val这份数据集不是把一堆图片和标签打散丢给你而是按 Roboflow 导出的标准目录组织。解压后你会看到这样的结构vis-drone-dataset-2/ ├── data.yaml ├── train/ │ ├── images/ │ └── labels/ ├── valid/ │ ├── images/ │ └── labels/ └── test/ ├── images/ └── labels/data.yaml内容如下names: [car, person] train: ./train/images val: ./valid/images test: ./test/images注意一个很多人忽略的细节标注目录是valid但 yaml 里的键仍然写成val。YOLOv5 训练时通过data[val]这个键去取验证集路径所以键名不能随便改改val为valid会直接报找不到 key而路径值./valid/images指向实际的 valid 文件夹。到这里就能发现这套配置是经过验证的不是随手写的。另一个容易踩的地方是相对路径。./train/images是相对于终端当前工作目录解析的不是相对于 data.yaml 所在目录解析。所以你在哪里执行 train.py哪里就得是 dataset 的上一级。比如把数据集解压在~/datasets/vis-drone-2进入这个目录后再跑训练路径就正好对上。2.2 标注文本的格式class、中心点归一化坐标与 bbox 宽高打开任意一个标签文件比如frame_003596_jpg.rf.65b4d825ed3c6fc222d449c5eab7a546.txt里面每一行是一个目标0 0.68203125 0.3984375 0.126953125 0.14453125 1 0.4921875 0.603515625 0.08984375 0.091796875第一列是类别编号0对应car1对应person顺序必须和 data.yaml 中names列表一致。后面四列分别是x_center y_center width height全部除以图片宽高做归一化范围在 0 到 1 之间。归一化坐标的好处是换分辨率不用重新标注。你在训练时用--img 640图片缩放到 640×640标注坐标不需要任何改动。这也是 YOLO 系数据集通用的格式YOLOv7、YOLOv8 训练数据时读同样的 txt。有个细节值得提该数据集由 Roboflow 导出图片已经被统一缩放并同步调整了标注框所以标签里的坐标和图片内容是对应得上的。如果你手动把这些图片再裁剪或 resize就得重新算坐标否则框会偏。2.3 文件名里的 .rf. 和 hash 说明数据来源看图片文件名比如frame_003596_jpg.rf.65b4d825ed3c6fc222d449c5eab7a546.jpg其中frame_003596_jpg是原始 VisDrone 抽帧名.rf.后面那串 24 位 hash 是 Roboflow 给这张标注图生成的唯一 ID。同一个数据集在 Roboflow 里导出多次相同图片会得到相同的 hash所以靠这个 hash 可以判断两份数据有没有重复。训练时不需要改文件名也不需要删掉 hash 后缀但必须保证图片文件在train/images同名 txt 在train/labels文件名主体完全一致。比如frame1160_jpg.rf.ea721eccf74c0677ca63544cab0bb646.jpg对应的标签文件必须叫frame1160_jpg.rf.ea721eccf74c0677ca63544cab0bb646.txt不带 hash 就会对应不上。2.4 YOLOv5 如何发现标签labels 目录定位与 cache 缓存机制YOLOv5 训练时会先扫描 images 目录把所有图片路径收集进来然后按三个规则找标签图片路径替换images为labels后缀替换.jpg为.txt文件名主体完全一致。所以train/images/frame_000576_jpg.rf.xxx.jpg的标签必须在train/labels/frame_000576_jpg.rf.xxx.txt。第一次扫描后会生成一份labels.cache缓存文件下次训练加载速度明显变快。这个缓存机制同时带来一个坑你往数据集里加了新图片或删了标签文件缓存不会自动更新YOLOv5 会继续按旧缓存判断“有无标签”导致新图片被跳过。遇到这种情况直接删掉train/labels.cache、valid/labels.cache和test/labels.cache重新训练就会重新扫描。3. 训练前的适配环境搭建、分辨率选 640 还是 1280、类别不平衡3.1 环境搭建conda 环境与依赖安装拿到数据集后第一步不是直接跑 train.py而是建一个干净的环境。YOLOv5 官方支持 Python 3.8 到 3.10推荐用 conda 隔离不要装到系统 Python 里原因很简单目标检测项目里 torch 和 torchvision 的版本绑定很严格后面再加别的项目很容易互相污染。conda create -n yolo python3.9 conda activate yolo pip install torch torchvision git clone https://github.com/ultralytics/yolov5 cd yolov5 pip install -r requirements.txt逻辑说明先创建名为yolo的 Python 3.9 环境并激活再安装 PyTorch 相关包。PyTorch 安装方式根据自己的 CUDA 版本选择CPU 版也能跑但这个数据集有 1000 多张图CPU 训练到后面会非常难受通常还是建议用 NVIDIA 显卡。参数说明requirements.txt会安装 opencv-python、numpy、matplotlib、seaborn、pandas 等依赖这些都是 YOLOv5 训练和画图必需的。装完后可以用python -c import torch; print(torch.cuda.is_available())确认 GPU 是否可用。3.2 分辨率参数640 用来跑通1280 用来刷指标训练分辨率是影响速度和精度的最大杠杆。无人机俯视视角的车辆和行人通常只占画面很小比例640 分辨率下很多 person 可能只有 20×20 像素。分辨率越高小目标特征保留得越完整但显存和耗时也会急剧上升。这里给出我常用的分辨率选型参考img 参数显存占用batch16适用场景320约 3~4GB验证代码流程、调试环境640约 7~9GB第一版正式训练多数情况够用1280约 18GB 以上追求小目标 mAP或做模型微调注意显存占用同时受模型大小和 batch 影响上表是按yolov5s估计的经验值。我的建议是先用 640 跑通第一个版本确认 data.yaml 路径正确、标签匹配、loss 能正常下降再决定要不要用 1280 做一轮高精度训练。直接上 1280 一旦报错你很难分清是数据问题还是参数问题。3.3 类别不平衡car 与 person 的样本比例要心里有数城市道路俯视场景里 car 数量通常远大于 person而且 person 目标更小天然难学。训练前最好统计一下两类目标的实例数量不统计的话训练结果很容易出现“car 识别得很好person 几乎不框”。用下面这个脚本快速统计from collections import Counter from pathlib import Path counter Counter() label_dir Path(train/labels) for txt in label_dir.glob(*.txt): for line in txt.read_text().strip().splitlines(): cls_id int(line.split()[0]) counter[cls_id] 1 print(car: , counter[0]) print(person: , counter[1])逻辑说明遍历训练集所有标签文件读取每一行取第一个字段作为类别编号0为 car1为 person累加后输出。如果 person 数量明显少于 car 的一半训练时要针对小目标类和少样本类做增强。参数说明脚本里label_dir可以根据实际目录改成valid/labels或test/labels统计 test 集时顺便能确认 test 标注是否存在避免验证时出问题。3.4 锚框 autoanchor要不要关掉YOLOv5 默认训练时会重新计算 anchor这能适配当前数据集的 bbox 尺寸分布。VisDrone 系列的车辆目标普遍偏小默认 COCO anchor 偏大所以第一次训练建议开着 autoanchor让它自动拟合数据。如果后面你已经调过一轮anchor 基本稳定每次训练还要重新计算并打印可以加--noautoanchor跳过节省一点时间。这一步属于“能用但别依赖”的技巧新手不建议一开始就关。4. 从训练到验证batch/epochs 参数设置与三个关键日志指标4.1 训练命令参数逐行解释目录结构和环境都准备好后进入 yolov5 仓库目录执行训练命令python train.py \ --data /path/to/vis-drone-dataset-2/data.yaml \ --weights yolov5s.pt \ --img 640 \ --batch 16 \ --epochs 100 \ --project runs/train \ --name visdrone_car_person逻辑说明train.py是 YOLOv5 训练入口--data指向 data.yaml--weights指定预训练权重--img设置输入图片尺寸--batch设置 batch size--epochs设置训练轮数--project和--name决定输出目录。参数说明yolov5s.pt是 YOLOv5s 的 COCO 预训练权重首次运行会自动下载。如果你网络下载困难也可以不加--weights从零训练但收敛速度和最终精度通常会差一些。--batch 16在 8GB 显存上是 yolo5s 的保守值显存小的机器把 batch 降到 8 或 4 即可别硬撑。训练开始后终端会实时打印每个 epoch 的日志包括 loss、精度、召回率、mAP。训练结束会在runs/train/visdrone_car_person/下生成best.pt和last.pt前者是验证集上 mAP 最高的权重后者是最后一轮的权重。实际部署优先用best.pt除非你确认最后几轮还有明显收益。4.2 训练日志怎么读box_loss、obj_loss、cls_loss 和 P/R/mAP训练日志每一行里有几个关键字段box_loss是边框回归损失obj_loss是置信度损失cls_loss是分类损失。这三个值在训练初期会下降得很快后期趋于平缓。不要只看 loss 绝对值要看趋势。三个 loss 持续下降说明模型还在学loss 先降后升或者验证 loss 和训练 loss 差距越来越大说明过拟合来了。对这个数据集而言100 轮以内一般不会严重过拟合但如果加了过多数据增强或者把 batch 设得太小过拟合就可能提前出现。每个 epoch 末尾还会打印P精确率、R召回率、mAP50IoU0.5 的平均精度、mAP50-95IoU 从 0.5 到 0.95 的平均精度。对无人机小目标场景我一般更看重mAP50和R因为小目标在 IoU 要求高的时候天然吃亏mAP50-95往往不太好看但不代表模型不可用。4.3 验证与推理val.py、conf 阈值、iou 阈值训练完先跑验证集确认模型泛化能力python val.py \ --data /path/to/vis-drone-dataset-2/data.yaml \ --weights runs/train/visdrone_car_person/weights/best.pt \ --img 640 \ --conf 0.25 \ --iou 0.5逻辑说明val.py读取训练好的权重在 data.yaml 指定的验证集上做推理并把结果和标注框对比计算各类别的 P、R、mAP。--conf 0.25表示置信度低于 0.25 的框被过滤--iou 0.5是 NMS 的 IoU 阈值。参数说明--conf影响召回率设得太高会漏掉低置信度的行人太低会出现大量误检--iou控制重复框的合并俯视场景下车辆密集0.5 是比较常见的设置。验证完还能直接跑检测把验证集的预测结果可视化出来python detect.py \ --source test/images \ --weights runs/train/visdrone_car_person/weights/best.pt \ --conf 0.25 \ --img 640 \ --save-txt逻辑说明--source可以指向图片、图片目录、视频文件或摄像头设备号这里指向 test/images 目录会逐张检测 test 集所有图片。--save-txt会把每个检测结果输出为 YOLO 格式的 txt 文件方便后续统计或接其他分析逻辑。输出结果默认保存在runs/detect/exp打开就能看到“框住车辆和行人的图片”这一步是整个流程里最直观的反馈。5. 避坑指南几个典型排错记录从路径报错到小目标漏检5.1 现象训练报警告提示很多图片找不到标签原因标签文件与图片不同名或标签目录层级不对或 labels.cache 缓存过期。最常见的是有些人把图片单独拖出来放到新目录标签没跟着过去也有少数情况下文件名大小写不一致比如图片是.jpg而标签文件后缀写成了.JPG。解决先用ls train/labels | head -n 5和ls train/images | head -n 5对比文件名主体如果一致删除labels.cache后重新训练。校验脚本在下一章给出能一次帮你查清。5.2 现象训练到第二个 epoch 附近显存溢出 OOM原因--img 1280配合--batch 16对显存压力极大YOLOv5 每个 epoch 还会做 mosaic 增强图片拼接后尺寸更大显存占用峰值比预想高。还有人是开着多个终端或 TensorBoard 网页抢显存。解决把 batch 降到 8 或 4这是最有效的做法其次把--img降到 640。别担心 batch 变小对精度影响有多大在这个数据规模下 batch 8 和 batch 16 的最终 mAP 差距通常在 1% 以内远小于分辨率带来的差距。5.3 现象mAP 看起来还行无人机视频里行人一个都没框住原因验证集和训练集来自同一批视频抽帧分布太接近模型在“见过”的数据上表现不错实际无人机场景拍摄高度、角度、光线都变了加上视频里行人目标更小容易被高置信度阈值过滤。解决把--conf降到 0.15 或 0.1 再测视频。如果误检变多再配合--iou 0.6优化 NMS。如果小目标还是漏就要回到训练环节用 1280 分辨率重训或对原图切片推理。5.4 现象car 的指标很好person 的召回率一直上不去原因样本数量不平衡加目标尺度小person 在标注框里可能只有十几个像素模型很难学到足够的特征。解决先统计标签确认 person 数量。确认不平衡后可以尝试增大训练分辨率到 1280在数据增强上把--hyp参数里的hsv_h、hsv_s调大也可以把mosaic增强从头开到底。需要让你心里有数的是如果 person 目标实在太小强行提高 recall 往往伴随大量误检业务上只关心车辆数量的话不必过分纠结。5.5 现象本地训练和验证一切正常换一台电脑跑同一个权重结果完全不同原因推理环境不一致常见的是 OpenCV 版本不同导致图片读取方式不同或者 PyTorch 版本差异导致算子行为有细微差别。解决部署时固定 PyTorch 和 OpenCV 版本最好直接用官方 requirements 同版本环境也可以用detect.py在同一环境里做一次对照测试。这个坑不属于数据集本身但在工程项目里经常把人折磨到怀疑模型。6. 进阶训练前的数据校验脚本与 TensorRT 导出部署6.1 训练前强制数据校验图片-标签一一对应脚本每次开始训练前我强烈建议跑一遍校验脚本而不是强行开训。这个脚本比上一章的统计脚本更进一步它会同时检查三个划分目录里的图像和标签文件名是否一一对应from pathlib import Path for split in [train, valid, test]: image_dir Path(split) / images label_dir Path(split) / labels images {p.stem for p in image_dir.glob(*.jpg)} labels {p.stem for p in label_dir.glob(*.txt)} missing_labels images - labels extra_labels labels - images print(f{split}: 图片 {len(images)} 张标签 {len(labels)} 个) print(f{split}: 缺少标签的图片 {len(missing_labels)} 张) print(f{split}: 没有图片的标签 {len(extra_labels)} 个) if missing_labels: print(示例:, list(missing_labels)[:3])逻辑说明对每个划分目录分别用glob收集图片文件和标签文件stem是文件名不带后缀的部分对比集合之后得到缺失标签和多余标签的集合并打印前面三个示例。参数说明脚本默认在当前工作目录下找 train/valid/test 文件夹如果数据集结构不同把split对应的路径改成绝对路径即可。如果发现missing_labels不为零优先检查是否大小写问题或标签文件没有解压完整。6.2 导出与部署TensorRT 引擎和 infer 推理训练完成后如果要在 Jetson 这类设备上跑实时检测建议把 best.pt 导出为 TensorRT 引擎python export.py \ --weights runs/train/visdrone_car_person/weights/best.pt \ --include engine \ --device 0 \ --img 640 \ --batch 1逻辑说明export.py会把 PyTorch 权重转换为 TensorRT 的 engine 格式--include engine指定导出格式--device 0使用所需 GPU 设备构建引擎--img 640固定输入尺寸。参数说明TensorRT 引擎和输入分辨率强绑定用 1280 导出的 engine 就只能跑 1280 输入换尺寸必须重新导出。另外导出时要求环境里已经装好 TensorRT 且版本与 PyTorch 兼容否则会提示找不到 trt 模块。如果不需要实时推理直接用 best.pt 跑 detect.py 也可以不必强求导出。从那以后我每次拿到新的目标检测数据集都会先跑上面的校验脚本再进训练环节。这个习惯帮我避开了至少三次“看起来能训、实际大量缺标签”的翻车也让人工排查的步骤真正被脚本替代。希望这整套流程能帮到你少走几段弯路。本文还有配套的精品资源点击获取
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。