作为计算机视觉领域的基础性技术,目标检测在业界具有广泛应用,而YOLO系列因其综合性能较好,成为广受欢迎的首选框架。

这次,为了让用户深入了解 OneFlow 训练目标检测模型的可行性以及性能的优越性,我们将 Ultralytics 版 YOLOv5(https://github.com/ultralytic...)通过 import oneflow as torch 的方式迁移为 OneFlow 后端(对应 YOLOv5 的 commit 为:48a85314bc80d8023c99bfb114cea98d71dd0591)。

相比 PyTorch,目前 YOLOv5 在 OneFlow 上进行小 batch 训练时有 5%-10% 的性能优势,而训练大 batch 的性能与 PyTorch 持平。

需要说明的是,本次发布的 OneFlow 后端的 One-YOLOv5 只是一个基础版本,还没有用任何优化技巧,相信在后续的一些定制化的性能优化技巧下(比如 nn.Graph 加持,算子的优化),OneFlow 可以继续提升 YOLOv5 在 COCO 等数据集的训练速度,更有效缩短目标检测模型的训练时间。

image.png

此外,我们还推出了《YOLOv5 全面解析教程》,其中对 YOLOv5 的相关教程进行了汉化,并添加了一系列详细的代码解读、原理讲解以及部署教程,并将逐步配备相应的视频讲解。相信这是一份很不错地从零开始深入学习 YOLOv5 的资料,希望帮助你更好地理解和使用 YOLOv5。

欢迎在 GitHub 上 Star One-YOLOv5 项目获取最新的动态,并在仓库提Issue、PR。以下将详细介绍 One-YOLOv5 在 COCO 上的精度以及性能表现。

1、One-YOLOv5与Ultralytics/YOLOv5精度一致

以YOLOv5n网络为例, result.csv
(https://oneflow-static.oss-cn...) 这个日志展示了基于 One-YOLOv5 在 COCO 上从零开始训练 YOLOv5n 网络的日志。下图展示了 box_loss , obj_loss, cls_loss ,map_0.5, map_0.5:0.95 等指标在训练过程中的变化情况:

image.png

可以看到,最终在第 300 个 epoch 时,map_0.5 达到了 0.45174,map_0.5:0.95 达到了0.27726。这与 Ultralytics/YOLOv5n 给出的精度数据一致。(https://github.com/ultralytic...,注意官网给出的精度指定 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/ultralytic...
  • OneFlow版本 yolov5 code base:https://github.com/Oneflow-In...
  • 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 

image.png

可以看到,在 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

image.png

在 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-In...


OneFlow
10 声望22 粉丝

★ OneFlow深度学习框架:github.com/Oneflow-Inc/oneflow ★ OF云平台:oneflow.cloud