作为计算机视觉领域的基础性技术,目标检测在业界具有广泛应用,而YOLO系列因其综合性能较好,成为广受欢迎的首选框架。
这次,为了让用户深入了解 OneFlow 训练目标检测模型的可行性以及性能的优越性,我们将 Ultralytics 版 YOLOv5(https://github.com/ultralytics/yolov5)通过 import oneflow as torch 的方式迁移为 OneFlow 后端(对应 YOLOv5 的 commit 为:48a85314bc80d8023c99bfb114cea98d71dd0591)。
相比 PyTorch,目前 YOLOv5 在 OneFlow 上进行小 batch 训练时有 5%-10% 的性能优势,而训练大 batch 的性能与 PyTorch 持平。
需要说明的是,本次发布的 OneFlow 后端的 One-YOLOv5 只是一个基础版本,还没有用任何优化技巧,相信在后续的一些定制化的性能优化技巧下(比如 nn.Graph 加持,算子的优化),OneFlow 可以继续提升 YOLOv5 在 COCO 等数据集的训练速度,更有效缩短目标检测模型的训练时间。
此外,我们还推出了《YOLOv5 全面解析教程》,其中对 YOLOv5 的相关教程进行了汉化,并添加了一系列详细的代码解读、原理讲解以及部署教程,并将逐步配备相应的视频讲解。相信这是一份很不错地从零开始深入学习 YOLOv5 的资料,希望帮助你更好地理解和使用 YOLOv5。
- 代码仓库地址:https://github.com/Oneflow-Inc/one-yolov5
- 文档网站地址:https://start.oneflow.org/oneflow-yolo-doc/index.html
- OneFlow 安装方法:https://github.com/Oneflow-Inc/oneflow#install-oneflow
欢迎在 GitHub 上 Star One-YOLOv5 项目获取最新的动态,并在仓库提Issue、PR。以下将详细介绍 One-YOLOv5 在 COCO 上的精度以及性能表现。
1
One-YOLOv5与Ultralytics/YOLOv5精度一致
以YOLOv5n网络为例, result.csv(https://oneflow-static.oss-cn-beijing.aliyuncs.com/one-yolo/YOLOv5n_results.csv) 这个日志展示了基于 One-YOLOv5 在 COCO 上从零开始训练 YOLOv5n 网络的日志。下图展示了 box_loss , obj_loss, cls_loss ,map_0.5, map_0.5:0.95 等指标在训练过程中的变化情况:
可以看到,最终在第 300 个 epoch 时,map_0.5 达到了 0.45174,map_0.5:0.95 达到了0.27726。这与 Ultralytics/YOLOv5n 给出的精度数据一致。(https://github.com/ultralytics/yolov5#pretrained-checkpoints,注意官网给出的精度指定 iou 为 0.65 的精度,而上述csv文件中是在 iou 为 0.60下的精度,使用我们训练的权重并把 iou 指定为 0.65 可以完全对齐官方给出的精度数据)。
关于这一点,我们可以使用 Ultralytics/YOLOv5n 来验证:
python val.py --weights yolov5n.pt --data data/coco.yaml --img 640 --iou 0.60
输出:
val: data=data/coco.yaml, weights=['yolov5n.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.6, max_det=300, task=val, device=, workers=8, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=False, dnn=False
YOLOv5 v6.1-384-g7fd9867 Python-3.8.13 torch-1.10.0+cu113 CUDA:0 (NVIDIA GeForce RTX 3080 Ti, 12054MiB)
cuda:0
Fusing layers...
YOLOv5n summary: 213 layers, 1867405 parameters, 0 gradients, 4.5 GFLOPs
val: Scanning '/data/dataset/fengwen/coco/val2017.cache' images and labels... 4952 found, 48 missing, 0 empty, 0 corrupt: 100%|█████
Class Images Instances P R mAP50 mAP50-95: 100%|██████████| 157/157 [00:40<00:00, 3.
all 5000 36335 0.573 0.432 0.456 0.277
上述输出可以说明,One-YOLOv5n 与 Ultralytics/YOLOv5n 的精度完全对齐。
在 One-YOLOv5n从零开始训练 YOLOv5n 进行精度复现的命令为 (2卡DDP 模式):
python -m oneflow.distributed.launch --nproc_per_node 2 train.py --data data/coco.yaml --weights ' ' --cfg models/yolov5n.yaml --batch 64
2
小batch训练,One-YOLOv5性能更优
以下性能结果都是直接将 PyTorch 切换为 OneFlow 之后测试的,还没有做针对性优化,后续会在此基础上继续提升 OneFlow 后端 YOLOv5 的训练速度。
在 3080Ti 的性能测试结果
单卡测试结果
- 以下为GTX 3080ti(12GB) 的YOLOv5测试结果(OneFlow后端 vs PyTorch后端)
- 以下测试结果的数据配置均为coco.yaml,模型配置也完全一样,并记录训练完COCO数据集的1个epoch所需时间
- 由于OneFlow Eager目前AMP的支持还不完善,所以我们提供的结果均为FP32模式下进行训练的性能结果
- PyTorch版本 yolov5 code base:https://github.com/ultralytics/yolov5
- OneFlow版本 yolov5 code base:https://github.com/Oneflow-Inc/one-yolov5
- CUDA 版本 11.7, cuDNN 版本为 8.5.0
- 测试命令(其中 batch 参数是动态变化的)为:
python train.py --batch 16 --cfg models/yolov5n.yaml --weights '' --data coco.yaml --img 640 --device 0
可以看到,在 batch 比较小时,OneFlow 后端的 YOLOv5 相比 PyTorch 有 5%-10% 的性能优势,这得益于 OneFlow 的 Eager 运行时系统可以更快的做 CUDA Kernel Launch。而 batch 比较大的时候 OneFlow 后端的 YOLOv5 相比于 PyTorch 的性能基本持平,这是因为当 batch 比较大时 CUDA Kernel Launch 的开销相比计算的开销会比较小。
两卡DDP测试结果
- 配置和单卡均一致
- 测试命令(其中 batch 参数是动态变化的)为:
python -m oneflow.distributed.launch --nproc_per_node 2 train.py --batch 16 --data coco.yaml --weights '' --device 0,1
在 2 卡 DDP 模式下,得益于单卡的性能优势,在 batch 比较小时,OneFlow 后端的 YOLOv5 相比 PyTorch 的性能稍微领先 ,而对于大 batch 来说,OneFlow的性能相比 PyTorch 基本持平。
3
总结
基于 OneFlow 移植 Ultralytics 版的 YOLOv5 ,在精度训练达标的情况下可以在 batch 比较小时取得一些性能优势。此外,对想深入了解 YOLOv5 的用户,我们发布了《YOLOv5全面解析教程》,希望它可以成为帮助你学习 YOLOv5 的绝佳资料。
欢迎下载体验 OneFlow v0.8.0 最新版本:
https://github.com/Oneflow-Inc/oneflow/
本文暂时没有评论,来添加一个吧(●'◡'●)