写在前面:本项目是基于 ultralytics 版 YOLOv5 源码改成 OneFlow 后端的结果,本工程的目的是做一个拥有更快训练速度的 YOLOv5 ,同时提供丰富的中文教程和源码细节解读,使得读者可以更加深入的了解 YOLOv5 。本 README 的其中一些部分也是直接用的 ultralytics 版 YOLOv5 README 的翻译,我们将相关链接替换为了 OneFlow 后端 YOLOv5 对应的链接。
- 🌟 v1.2.0 正式开源。v1.2.0同步了ultralytics yolov5的v7.0版本,同时支持分类,目标检测,实例分割等任务 详情请看:Release Note
- 🎉 代码仓库地址
- 🎉 文档网站地址
- 🎉 OneFlow 安装方法 (注意
⚠️ : 目前OneFlow 需要安装 nightly 版本,等OneFlow后续新版本发布后可以安装稳定版。此外 OneFlow 目前仅对Linux平台有完整支持,请 Windows 和 Macos 用户注意)
不过即使你对 OneFlow 带来的性能提升不感兴趣,我们相信OneYOLOv5文档网站中对 ultralytics/yolov5 相关教程的汉化以及源码剖析也会是从零开始深入学习 YOLOv5 一份不错的资料。欢迎在仓库给我们提出宝贵的意见。🌟🌟🌟
请查看 文档网站 获取关于训练,测试和推理的完整文档。
注意
-
oneflow目前不支持windows平台
-
--batch 必须是GPU数量的倍数。
-
GPU 0 将比其他GPU占用略多的内存,因为它维护EMA并负责检查点等。
安装
在Python>=3.7.0 的环境中克隆版本仓并安装 requirements.txt,包括 OneFlow nightly 。
git clone https://github.com/Oneflow-Inc/one-yolov5 # 克隆
cd one-yolov5
pip install -r requirements.txt # 安装
YOLOv5实例分割模型支持使用 --data coco128-seg.yaml
参数自动下载 COCO128-seg
测试数据集(测试数据集表示能测试项目正常运行的小数据集), 以及使用 bash data/scripts/get_coco.sh --train --val --segments
或者使用 python train.py --data coco.yaml
下载 COCO-segments
数据集
# Single-GPU
python segment/train.py --weights yolov5s-seg.of --data coco128-seg.yaml --epochs 5 --img 640
# Multi-GPU DDP
python -m oneflow.distributed.launch --nproc_per_node 4 segment/train.py --weights yolov5s-seg.of --data coco128-seg.yaml --epochs 5 --img 640 --device 0,1,2,3
注意 :
- {
.of
: 代表OneFlow预训练权重 ,.pt
: 代表 PyTorch 预训练权重 } --weights yolov5s-seg.of
表示使用OneFlow预训练权重 , 也是支持使用 PyTorch 预训练权重 如--weights yolov5s-seg.pt
- 模型权重将自动从 github 下载(建议如果没有设置代理,可以提前将模型下载到电脑本地 使用
--weights 本地路径/yolov5s-seg.of
)
在ImageNet-1k数据集上验证YOLOv5m-seg 模型的精度
bash data/scripts/get_coco.sh --val --segments # download COCO val segments split (780MB, 5000 images)
python segment/val.py --weights yolov5s-seg.of --data coco.yaml --img 640 # validate
使用预训练模型(YOLOv5m-seg) 预测图片
python segment/predict.py --weights yolov5m-seg.of --data data/images/
将 yolov5s-seg
模型导出为 ONNX 格式 示例
python export.py --weights yolov5s-seg.of --include onnx --img 640 --device 0
消费级显卡的春天,GTX 3090 YOLOv5s单卡完整训练COCO数据集缩短11.35个小时
注意: 这里的模型权重为v1.1.0版本的权重,最新权重下载 releases/tag/v1.2.0
模型 | ONNX版本模型 | 规模 (像素) |
mAP验证 0.5:0.95 |
mAP验证 0.5 |
速度 CPU b1 (ms) |
速度 V100 b1 (ms) |
速度 V100 b32 (ms) |
参数 (M) |
浮点运算 @640 (B) |
---|---|---|---|---|---|---|---|---|---|
YOLOv5n | YOLOv5n.onnx | 640 | 28.0 | 45.7 | 45 | 6.3 | 0.6 | 1.9 | 4.5 |
YOLOv5s | YOLOv5s.onnx | 640 | 37.4 | 56.8 | 98 | 6.4 | 0.9 | 7.2 | 16.5 |
YOLOv5m | YOLOv5m.onnx | 640 | 45.4 | 64.1 | 224 | 8.2 | 1.7 | 21.2 | 49.0 |
YOLOv5l | YOLOv5l.onnx | 640 | 49.0 | 67.3 | 430 | 10.1 | 2.7 | 46.5 | 109.1 |
YOLOv5x | YOLOv5x.onnx | 640 | 50.7 | 68.9 | 766 | 12.1 | 4.8 | 86.7 | 205.7 |
YOLOv5n6 | YOLOv5n6.onnx | 1280 | 36.0 | 54.4 | 153 | 8.1 | 2.1 | 3.2 | 4.6 |
YOLOv5s6 | YOLOv5s6.onnx | 1280 | 44.8 | 63.7 | 385 | 8.2 | 3.6 | 12.6 | 16.8 |
YOLOv5m6 | YOLOv5m6.onnx | 1280 | 51.3 | 69.3 | 887 | 11.1 | 6.8 | 35.7 | 50.0 |
YOLOv5l6 | YOLOv5l6.onnx | 1280 | 53.7 | 71.3 | 1784 | 15.8 | 10.5 | 76.8 | 111.4 |
YOLOv5x6 + [TTA][TTA] |
YOLOv5x6.onnx | 1280 1536 |
55.0 55.8 |
72.7 72.7 |
3136 - |
26.2 - |
19.4 - |
140.7 - |
209.8 - |
表格注释 (点击扩展)
- 所有检查点都以默认设置训练到300个时期. Nano和Small模型用 hyp.scratch-low.yaml hyps, 其他模型使用 hyp.scratch-high.yaml.
- mAPval 值是 COCO val2017 数据集上的单模型单尺度的值。
复现方法:python val.py --data coco.yaml --img 640 --conf 0.001 --iou 0.65
- 使用 AWS p3.2xlarge 实例对COCO val图像的平均速度。不包括NMS时间(~1 ms/img)
复现方法:python val.py --data coco.yaml --img 640 --task speed --batch 1
- TTA 测试时数据增强 包括反射和比例增强.
复现方法:python val.py --data coco.yaml --img 1536 --iou 0.7 --augment