作者|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个优化点。文章来源:https://www.toymoban.com/news/detail-451400.html
第一,对于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。接下来将分别详细阐述文章来源地址https://www.toymoban.com/news/detail-451400.html
到了这里,关于一块RTX 3090加速训练YOLOv5s,时间减少11个小时,速度提升20%的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!