消费级显卡的春天，GTX 3090 YOLOv5s单卡完整训练COCO数据集缩短11.35个小时

前语本文介绍了在单卡上凭借对YOLOv5的功能剖析以及几个简略的优化将GTX 3090 FP32 YOLOv5s的练习速度进步了近20%。关于需求迭代300个Epoch的COCO数据集来说比较 ultralytics/yolov5 咱们缩短了11.35个小时的练习时刻。

本文转载自GiantPandaCV

作者 | BBuf

欢迎重视大众号CV技能指南，专心于核算机视觉的技能总结、最新技能盯梢、经典论文解读、CV招聘信息。

核算机视觉入门1v3辅导班

0x1. 成果展现

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

消费级显卡的春天，GTX 3090 YOLOv5s单卡完整训练COCO数据集缩短11.35个小时

能够看到在单卡形式下，通过咱们的优化比较于 ultralytics/yolov5 的练习速度，咱们进步了 20% 左右。

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

在DDP形式下的功能进步起伏没有单卡这么多，猜想可能是通信部分的开支比较大，后续咱们会再研讨一下。

0x2. 优化手段

咱们在这一节完结技能揭秘。咱们深度剖析了PyTorch的YOLOv5的履行序列，咱们发现当前YOLOv5主要存在3个优化点。第一个便是关于Upsample算子的改善，因为YOLOv5运用上采样是规整的最近邻2倍插值，所以咱们能够完成一个特别Kernel下降核算量并进步带宽。第二个便是在YOLOv5中存在一个滑动更新模型参数的操作，这个操作启动了许多碎的CUDA Kernel，而每个CUDA Kernel的履行时刻都十分短，所以启动开支不能忽略。咱们运用水平并行CUDA Kernel的办法（MultiTensor）对其完结了优化，根据这个优化One-YOLOv5获得了9%的加速。第三个优化点来源于对YOLOv5 nsys履行序列的观察，咱们发现在ComputeLoss部分出现的bbox_iou是整个Loss核算部分一个比较大的瓶颈，咱们在bbox_iou函数部分完结了多个笔直的Kernel Fuse，使得它的开支从开始的3.x ms下降到了几百个us。接下来咱们别离详细论述这几种优化：

0x2.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 上调优前后的带宽占用和核算时刻比较：

框架

数据类型

Op类型

带宽利用率

耗时

PyTorch

Float32

UpsampleNearest2D forward

28.30%

111.42us

PyTorch

Float32

UpsampleNearest2D backward

60.16%

65.12us

OneFlow 未优化

Float32

UpsampleNearest2D forward

12.54%

265.82us

OneFlow 未优化

Float32

UpsampleNearest2D backward

18.4%

260.22us

OneFlow 优化后

Float32

UpsampleNearest2D forward

52.18%

61.44us

OneFlow 优化后

Float32

UpsampleNearest2D backward

77.66%

50.56us

PyTorch

Float16

UpsampleNearest2D forward

16.99%

100.38us

PyTorch

Float16

UpsampleNearest2D backward

31.56%

57.38us

OneFlow 未优化

Float16

UpsampleNearest2D forward

7.07%

262.82us

OneFlow 未优化

Float16

UpsampleNearest2D backward

41.04%

558.88us

OneFlow 优化后

Float16

UpsampleNearest2D forward

43.26%

35.36us

OneFlow 优化后

Float16

UpsampleNearest2D backward

44.82%

40.26us

上述成果运用 /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%的进步。

0x2.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核算很快，但拜访Global Memory以及屡次Kernel Launch的开支也是比较大的，所以咱们做了几个fuse来下降Kernel Launch的开支以及削减拜访Global Memrory来进步带宽。

在这儿插入图片描述

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

详细来说咱们这儿做了如下的几个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%左右，是一个十分有用的优化。

0x2.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'

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

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

而且水平方向的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 的代码完成了，感兴趣的能够留言后续独自讲解。

0x3. 运用办法

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

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

0x4. 总结

目前，yolov5s网络当以BatchSize=16的配置在GeForce RTX 3090上（这儿指定BatchSize为16时）练习COCO数据集时，OneFlow比较PyTorch能够节约 11.35 个小时。咱们相信这篇文章提到的优化技巧也能够对更多的从事方针检测的学生或者工程师带来启示。欢迎我们star one-yolov5项目：github.com/Oneflow-Inc…

欢迎重视大众号CV技能指南，专心于核算机视觉的技能总结、最新技能盯梢、经典论文解读、CV招聘信息。

核算机视觉入门1v3辅导班

【技能文档】《从零建立pytorch模型教程》122页PDF下载

QQ沟通群：444129970。群内有大佬负责解答我们的日常学习、科研、代码问题。

其它文章

DiffusionDet：用于方针检测的分散模型

CV小常识评论与剖析(7) 寻觅论文立异点的新办法

CV小常识剖析与评论(6）论文立异的一点误区

一文看尽深度学习中的各种注意力机制

MMYOLO 想你所想：练习进程可视化

顶刊TPAMI 2023！Food2K：大规模食物图像识别

用于精确方针检测的多网格冗余边界框标注

2023最新半监督语义切割总述 | 技能总结与展望！

本来Transformer便是一种图神经网络，这个概念你清楚吗？

快速完成常识蒸馏算法，运用 MMRazor 就够啦！

常识蒸馏的搬迁学习应用

TensorFlow 真的要被 PyTorch 比下去了吗？

深入剖析MobileAI图像超分最佳计划：ABPN

3D方针检测中点云的稀少性问题及解决计划

一文深度剖析分散模型终究学到了什么？

OpenMMLab教程【零】OpenMMLab介绍与装置

代码实战：YOLOv5完成钢材外表缺陷检测

TensorRT教程（六）运用Python和C++部署YOLOv5的TensorRT模型

超全汇总 | 核算机视觉/自动驾驶/深度学习资料合集！

高精度语义地图构建的一点考虑

点云切割练习哪家强？监督，弱监督，无监督还是半监督？

核算机视觉入门1v3辅导班

核算机视觉沟通群

用于超大图像的练习战略：Patch Gradient Descent

CV小常识评论与剖析（5）究竟什么是Latent Space?

【免费送书活动】关于语义切割的亿点考虑

新计划：从过错中学习，点云切割中的自我规范化层次语义表示

经典文章：Transformer是如何进军点云学习范畴的？

CVPR 2023 Workshop | 首个大规模视频全景切割比赛

如何更好地应对下游小样本图像数据？不平衡数据集的建模的技巧和策

Transformer沟通群

深度理解变分自编码器(VAE) | 从入门到精通

U-Net在2022年相关研讨的论文引荐

用少于256KB内存完成边缘练习，开支不到PyTorch千分之一

PyTorch 2.0 重磅发布：一行代码提速 30%

Hinton 最新研讨：神经网络的未来是前向-前向算法

聊聊核算机视觉入门

声明：本站所有文章，如无特殊说明或标注，均为本站原创发布。任何个人或组织，在未征得本站同意时，禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益，可联系我们进行处理。

消费级显卡的春天，GTX 3090 YOLOv5s单卡完整训练COCO数据集缩短11.35个小时

0x1. 成果展现

0x2. 优化手段

0x2.1 对UpsampleNearest2D的特化改善

0x2.2 对bbox_iou函数进行优化 (笔直Fuse优化)

0x2.3 对模型滑动均匀更新进行优化（水平Fuse优化）

0x3. 运用办法

0x4. 总结

近期文章

近期评论