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

作者｜BBuf

很快乐为大家带来One-YOLOv5的最新进展，在《一个更快的YOLOv5面世，附送全面中文解析教程》(mp.weixin.qq.com/s/imTnKQVWc…) 发布后收到了许多算法工程师朋友的重视，非常感谢。

不过，或许你也在思考一个问题：尽管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链接：
github.com/Oneflow-Inc…

欢迎你给咱们在GitHub上点个Star，咱们会用更多高质量技术共享来回馈社区。对 One-YOLOv5 感兴趣的小伙伴能够添加bbuf23333进入One-YOLOv5微信沟通群，或许直接扫二维码：

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

1

效果展现

咱们展现一下别离运用One-YOLOv5以及Ultralytics/YOLOv5在RTX 3090单卡上运用YOLOv5s FP32模型练习COCO数据集的一个Epoch所需的耗时：

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

能够看到，在单卡形式下，通过优化后的One-YOLOv5比较Ultralytics/YOLOv5的练习速度进步了20%左右。

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

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

在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。接下来将别离详细论述这三种优化。

2.1 对UpsampleNearest2D的特化改善

这儿直接展现咱们对UpsampleNearest2D进行调优的技术总结，大家能够结合下面的PR链接来对应下面的常识点进行总结。咱们在A100 40G上测验了UpsampleNearest2D算子的功能体现，这块卡的峰值带宽在1555Gb/s , 咱们运用的CUDA版别为11.8。

进行 Profile 的程序如下：

import oneflow as flow
x = flow.randn(16, 32, 80, 80, device="cuda", dtype=flow.float32).requires_grad_()
m = flow.nn.Upsample(scale_factor=2.0, mode="nearest")
y = m(x)
print(y.device)
y.sum().backward()

github.com/Oneflow-Inc… & github.com/Oneflow-Inc… 这两个 PR 别离针对 UpsampleNearest2D 这个算子（这个算子是 YOLO 系列算法许多运用的）的前后向进行了调优，下面展现了在 A100 上调优前后的带宽占用和核算时间比较：

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

上述效果运用 /usr/local/cuda/bin/ncu -o torch_upsample /home/python3 debug.py 得到profile文件后运用Nsight Compute翻开记载。

基于上述对 UpsampleNearest2D 的优化，OneFlow 在 FP32 和 FP16 状况下的功能和带宽都大幅超越之前未经优化的版别，并且比较于 PyTorch 也有较大幅度的抢先。

本次优化涉及到的常识点总结如下（by OneFlow 柳俊丞）：

  • 为常见的状况写特例，比方这儿便是为采样倍数为2的Nearest插值写特例，避免运用NdIndexHelper带来的额定核算开支，不必追求再一个kernel完成中同时具有通用型和高效性；

  • 整数除法开支大（可是编译器有的时分会优化掉一些除法），nchw中的nc不需求分隔，兼并在一起核算减少核算量；

  • int64_t除法的开支更大，用int32满意大部分需求，其实这儿还有一个快速整数除法的问题；

  • 反向Kernel核算过程中循环dx比较循环dy ，实践上将坐标换算的开支减少到本来的1/4；

  • CUDA GMEM的开支的也比较大，尽管编译器有或许做优化，可是显式的运用局部变量更好；

  • 一次Memset的开支也很大，和写一次相同，所以反向Kernel中对dx运用Memset清零的机遇需求留意；

  • atomicAdd开支很大，即使抛开为了完成原子性或许需求的锁总线等，atomicAdd需求把本来的值先读出来，再写回去；别的，half的atomicAdd 巨慢无比，慢到假如一个算法需求用到atomicAdd，那么比较于用half ，转成float ，再atomicAdd，再转回去还要慢许多；

  • 向量化访存。

对这个Kernel进行特化是优化的第一步，基于这个优化能够给YOLOv5的单卡 PipLine 带来1%的进步。

2.2 对bbox_iou函数进行优化 (笔直Fuse优化)

通过对nsys的剖析，咱们发现无论是One-YOLOv5还是Ultralytics/YOLOv5，在核算Loss的阶段都有一个耗时比较严重的bbox_iou函数，这儿贴一下bbox_iou部分的代码：

def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7):
    # Returns Intersection over Union (IoU) of box1(1,4) to box2(n,4)
    # Get the coordinates of bounding boxes
    if xywh:  # transform from xywh to xyxy
        (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
        w1_, h1_, w2_, h2_ = w1 / 2, h1 / 2, w2 / 2, h2 / 2
        b1_x1, b1_x2, b1_y1, b1_y2 = x1 - w1_, x1 + w1_, y1 - h1_, y1 + h1_
        b2_x1, b2_x2, b2_y1, b2_y2 = x2 - w2_, x2 + w2_, y2 - h2_, y2 + h2_
    else:  # x1, y1, x2, y2 = box1
        b1_x1, b1_y1, b1_x2, b1_y2 = box1.chunk(4, -1)
        b2_x1, b2_y1, b2_x2, b2_y2 = box2.chunk(4, -1)
        w1, h1 = b1_x2 - b1_x1, (b1_y2 - b1_y1).clamp(eps)
        w2, h2 = b2_x2 - b2_x1, (b2_y2 - b2_y1).clamp(eps)
    # Intersection area
    inter = (b1_x2.minimum(b2_x2) - b1_x1.maximum(b2_x1)).clamp(0) * \
            (b1_y2.minimum(b2_y2) - b1_y1.maximum(b2_y1)).clamp(0)
    # Union Area
    union = w1 * h1 + w2 * h2 - inter + eps
    # IoU
    iou = inter / union
    if CIoU or DIoU or GIoU:
        cw = b1_x2.maximum(b2_x2) - b1_x1.minimum(b2_x1)  # convex (smallest enclosing box) width
        ch = b1_y2.maximum(b2_y2) - b1_y1.minimum(b2_y1)  # convex height
        if CIoU or DIoU:  # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1
            c2 = cw ** 2 + ch ** 2 + eps  # convex diagonal squared
            rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4  # center dist ** 2
            if CIoU:  # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47
                v = (4 / math.pi ** 2) * (torch.atan(w2 / h2) - torch.atan(w1 / h1)).pow(2)
                with torch.no_grad():
                    alpha = v / (v - iou + (1 + eps))
                return iou - (rho2 / c2 + v * alpha)  # CIoU
            return iou - rho2 / c2  # DIoU
        c_area = cw * ch + eps  # convex area
        return iou - (c_area - union) / c_area  # GIoU https://arxiv.org/pdf/1902.09630.pdf
    return iou  # IoU

以One-YOLOv5的原始履行序列图为例，咱们发现bbox_iou函数这部分每一次运行都需求花2.6ms左右，并且能够看到这儿有许多的小Kernel被调度，尽管每个小Kernel核算很快，但拜访GlobalMemory以及多次KernelLaunch的开支也比较大，所以咱们做了几个fuse来下降Kernel Launch的开支以及减少拜访Global Memrory来进步带宽。

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

通过咱们的Kernel Fuse之后的耗时只需求600+us。

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

具体来说咱们这儿做了如下的几个fuse：

  • fused_get_boundding_boxes_coord：github.com/Oneflow-Inc…

  • fused_get_intersection_area: github.com/Oneflow-Inc…

  • fused_get_iou: github.com/Oneflow-Inc…

  • fused_get_convex_diagonal_squared: github.com/Oneflow-Inc…

  • fused_get_center_dist: github.com/Oneflow-Inc…

  • fused_get_ciou_diagonal_angle: github.com/Oneflow-Inc…

  • fused_get_ciou_result: github.com/Oneflow-Inc…

然后咱们在One-YOLOv5的train.py中扩展了一个 –bbox_iou_optim 选项，只需练习的时分带上这个选项就会主动调用上面的fuse kernel来对bbox_iou函数进行优化了，具体请看：github.com/Oneflow-Inc… 。对bbox_iou这个函数的一系列笔直Fuse优化使得YOLOv5全体的练习速度进步了8%左右，是一个非常有效的优化。

2.3 对模型滑动均匀更新进行优化（水平Fuse优化）

在 YOLOv5 中会运用EMA（指数移动均匀）对模型的参数做均匀, 一种给予近期数据更高权重的均匀方法, 以求进步测验目标并添加模型鲁棒。这儿的核心操作如下代码所示：

def update(self, model):
        # Update EMA parameters
        self.updates += 1
        d = self.decay(self.updates)
        msd = de_parallel(model).state_dict()  # model state_dict
        for k, v in self.ema.state_dict().items():
            if v.dtype.is_floating_point:  # true for FP16 and FP32
                v *= d
                v += (1 - d) * msd[k].detach()
        # assert v.dtype == msd[k].dtype == flow.float32, f'{k}: EMA {v.dtype} and model {msd[k].dtype} must be FP32'

以下是未优化前的这个函数的时序图：

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

这部分的CUDAKernel的履行速度大概为7.4ms，而通过咱们水平Fuse优化（即MultiTensor），这部分的耗时状况下降了127us。

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

并且水平方向的Kernel Fuse也相同下降了Kernel Launch的开支，使得前后2个Iter的间隙也进一步缩短了。终究这个优化为YOLOv5的全体练习速度进步了10%左右。本优化完成的pr如下：github.com/Oneflow-Inc…

此外，关于Optimizer部分相同能够水平并行，所以咱们在One-YOLOv5里设置了一个multi_tensor_optimizer标志，翻开这个标志就能够让 optimizer 以及 EMA 的 update以水平并行的方法运行。

关于MultiTensor这个常识能够看 zzk 的这篇文章：zhuanlan.zhihu.com/p/566595789 。zzk 在 OneFlow 中也完成了一套 MultiTensor 计划，上面的 PR 9498 也是基于这套 MultiTensor 计划完成的。介于篇幅原因咱们就不翻开MultiTensor的代码完成了，感兴趣朋友的能够留言后续单独讲解。

3

运用方法

上面已经说到所有的优化都集中于bbox_iou_optim和multi_tensor_optimizer这两个扩展的Flag，只需咱们练习的时分翻开这两个Flag就能够享受到上述优化了。其他的运行指令和One-YOLOv5没有变化，以One-YOLOv5在RTX 3090上练习YOLOv5为例，指令为：

python train.py
--batch 16 
--cfg models/yolov5s.yaml 
--weights '' 
--data coco.yaml 
--img 640 
--device 0 
--epoch 1 
--bbox_iou_optim 
--multi_tensor_optimizer

4

总结

现在，YOLOv5s网络当以BatchSize=16的配置在GeForce RTX 3090上（这儿指定BatchSize为16时）练习COCO数据集时，OneFlow比较PyTorch能够节省 11.35 个小时。期望这篇文章说到的优化技巧能够对更多的从事目标检测的工程师带来启发。
欢迎Star One-YOLOv5项目：
github.com/Oneflow-Inc…

One-YOLOv5的优化工作实践上不仅包含功能，咱们现在也付出了许多汗水在文档和源码解读上，后续会持续放出《YOLOv5全面解析教程》(mp.weixin.qq.com/s/imTnKQVWc… )的其他文章，并将尽快发布新版别。

5

称谢

感谢同事柳俊丞在这次调优中提供的 idea 和技术支持，感谢胡伽魁(mp.weixin.qq.com/s/XchwBXvHR…) 同学完成的一些fuse kernel，感谢郑泽康(mp.weixin.qq.com/s/5Rx0YzCXI…) 和宋易承的MultiTensorUpdate完成，感谢冯文的精度验证工作以及文档支持，以及小糖对One-YOLOv5的推广，以及协助本项目发展的工程师如赵露阳、梁德澎(mp.weixin.qq.com/s/jsObI7kOo…) 等等。本项目未来会持续发力做出更多的效果。

欢迎 Star、试用 OneFlow 最新版别：
github.com/Oneflow-Inc…