一塊RTX 3090加速訓練YOLOv5s,時間減少11個小時,速度提升20%

語言: CN / TW / HK

作者|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。接下來將分別詳細闡述這三種優化。

 

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()

https://github.com/Oneflow-Inc/oneflow/pull/9415 & https://github.com/Oneflow-Inc/oneflow/pull/9424 這兩個 PR 分別針對 UpsampleNearest2D 這個算子(這個算子是 YOLO 系列算法大量使用的)的前後向進行了調優,下面展示了在 A100 上調優前後的帶寬佔用和計算時間比較:

 

上述結果使用  /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來提升帶寬。

 

 

經過我們的Kernel Fuse之後的耗時只需要600+us。

 

 

具體來説我們這裏做了如下的幾個fuse:

  • fused_get_boundding_boxes_coord:https://github.com/Oneflow-Inc/oneflow/pull/9433

  • fused_get_intersection_area: https://github.com/Oneflow-Inc/oneflow/pull/9485

  • fused_get_iou: https://github.com/Oneflow-Inc/oneflow/pull/9475

  • fused_get_convex_diagonal_squared: https://github.com/Oneflow-Inc/oneflow/pull/9481

  • fused_get_center_dist: https://github.com/Oneflow-Inc/oneflow/pull/9446

  • fused_get_ciou_diagonal_angle: https://github.com/Oneflow-Inc/oneflow/pull/9465

  • fused_get_ciou_result: https://github.com/Oneflow-Inc/oneflow/pull/9462


然後我們在One-YOLOv5的train.py中擴展了一個  --bbox_iou_optim  選項,只要訓練的時候帶上這個選項就會自動調用上面的fuse kernel來對bbox_iou函數進行優化了,具體請看: https://github.com/Oneflow-Inc/one-yolov5/blob/main/utils/metrics.py#L224-L284 。對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'

以下是未優化前的這個函數的時序圖:

 

 

這部分的CUDAKernel的執行速度大概為7.4ms,而經過我們水平Fuse優化(即MultiTensor),這部分的耗時情況降低了127us。

 

 

並且水平方向的Kernel Fuse也同樣降低了Kernel Launch的開銷,使得前後2個Iter的間隙也進一步縮短了。最終這個優化為YOLOv5的整體訓練速度提升了10%左右。本優化實現的pr如下: https://github.com/Oneflow-Inc/oneflow/pull/9498

此外,對於Optimizer部分同樣可以水平並行,所以我們在One-YOLOv5裏設置了一個 multi_tensor_optimizer 標誌,打開這個標誌就可以讓 optimizer 以及 EMA 的 update以水平並行的方式運行。

關於MultiTensor這個知識可以看 zzk 的這篇文章: https://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項目:
https://github.com/Oneflow-Inc/one-yolov5

One-YOLOv5的優化工作實際上不僅包含性能,我們目前也付出了很多心血在文檔和源碼解讀上,後續會繼續放出《YOLOv5全面解析教程》的其他文章,並將儘快發佈新版本。

 

5

致謝

 

感謝同事柳俊丞在這次調優中提供的 idea 和技術支持,感謝胡伽魁同學實現的一些fuse kernel,感謝鄭澤康和宋易承的MultiTensorUpdate實現,感謝馮文的精度驗證工作以及文檔支持,以及小糖對One-YOLOv5的推廣,以及幫助本項目發展的工程師如趙露陽、樑德澎等等。本項目未來會繼續發力做出更多的成果。
 

其他人都在看

歡迎Star、試用OneFlow最新版本:https://github.com/Oneflow-Inc/oneflow/


 

 

本文分享自微信公眾號 - OneFlow(OneFlowTechnology)。
如有侵權,請聯繫 [email protected] 刪除。
本文參與“OSC源創計劃”,歡迎正在閲讀的你也加入,一起分享。