一块RTX 3090加速训练YOLOv5s，时间减少11个小时，速度提升20%

作者｜BBuf

很高兴为大家带来One-YOLOv5的最新进展，在《一个更快的YOLOv5问世，附送全面中文解析教程》发布后收到了很多算法工程师朋友的关注，十分感谢。

不过，可能你也在思考一个问题：虽然OneFlow的兼容性做得很好，可以很方便地移植YOLOv5并使用OneFlow后端来进行训练，但为什么要用OneFlow？能缩短模型开发周期吗？解决了任何痛点吗？本篇文章将尝试回答这几个问题。

我曾经也是一名算法工程师，开发机器也只有两张RTX 3090消费级显卡而已，但实际上大多数由我上线的检测产品也就是靠这1张或者2张RTX 3090完成的。

由于成本问题，很多中小公司没有组一个A100集群或者直接上数十张卡来训练检测模型的实力，所以这个时候在单卡或者2卡上将目标检测模型做快显得尤为重要。模型训练速度提升之后可以降本增效，提高模型生产率。

所以，近期我和实习生小伙伴一起凭借对YOLOv5的性能分析以及几个简单的优化，将单RTX 3090 FP32 YOLOv5s的训练速度提升了近20%。对于需要迭代300个Epoch的COCO数据集来说，One-YOLOv5相比Ultralytics/YOLOv5缩短了11.35个小时的训练时间。

本文将分享我们的所有优化技术，如果你是一名PyTorch和OneFlow的使用者，尤其日常和检测模型打交道但资源相对受限，那么本文的优化方法将对你有所帮助。

One-YOLOv5链接：
https://github.com/Oneflow-Inc/one-yolov5

欢迎你给我们在GitHub上点个Star，我们会用更多高质量技术分享来回馈社区。对 One-YOLOv5 感兴趣的小伙伴可以添加bbuf23333进入One-YOLOv5微信交流群。

1

结果展示

我们展示一下分别使用One-YOLOv5以及Ultralytics/YOLOv5在RTX 3090单卡上使用YOLOv5s FP32模型训练COCO数据集的一个Epoch所需的耗时：

可以看到，在单卡模式下，经过优化后的One-YOLOv5相比Ultralytics/YOLOv5的训练速度提升了20%左右。

然后我们再展示一下2卡DDP模式YOLOv5s FP32模型训练COCO数据集一个Epoch所需的耗时：

在DDP模式下，One-YOLOv5的性能依然领先，但还需要进一步，猜测可能是通信部分的开销比较大，后续我们会再研究一下。

2

 优化手段

我们深度分析了PyTorch的YOLOv5的执行序列，发现当前YOLOv5主要存在3个优化点。

第一，对于Upsample算子的改进，由于YOLOv5使用上采样是规整的最近邻2倍插值，所以我们可以实现一个特殊Kernel降低计算量并提升带宽。

第二，在YOLOv5中存在一个滑动更新模型参数的操作，这个操作启动了很多碎的CUDA Kernel，而每个CUDA Kernel的执行时间都非常短，所以启动开销不能忽略。我们使用水平并行CUDA Kernel的方式（MultiTensor）对其完成了优化，基于这个优化，One-YOLOv5获得了9%的加速。

第三，通过对YOLOv5nsys执行序列的观察发现，在ComputeLoss部分出现的bbox_iou是整个Loss计算部分的比较大的瓶颈，我们在bbox_iou函数部分完成了多个垂直的KernelFuse，使得它的开销从最初的3.xms降低到了几百个us。接下来将分别详细阐述