目標檢測、例項分割、旋轉框樣樣精通!詳解高效能檢測演算法 RTMDet
1. 簡介
近幾年來,目標檢測模型,尤其是單階段目標檢測模型在工業場景中已經得到廣泛應用。對於檢測演算法來說,模型的精度以及執行效率是實際使用時最受關注的指標。因此,我們對目前的單階段目標檢測器進行了全面的改進:從增強模型的特徵提取能力和對各個元件的計算量進行均衡化,到使用動態的軟標籤來優化訓練策略,再到對資料增強的效率和效能進行改進, 我們對演算法的各個方面都提出了新的優化方案。
通過這些改進,我們得到了從 tiny 到 extra-large 大小的一系列高效能檢測模型,我們將這套實時目標檢測模型(R eal- T ime M odels for object Det ection)命名為 RTMDet。其中,RTMDet-x 取得了 52.8 mAP 的精度和 300+FPS 的執行速度,同時,RTMDet-tiny 在僅有 4M 引數量的情況下達到了 41.1 mAP 的精度,超越了同級別的所有模型。RTMDet 不僅僅在目標檢測這一任務上效能優異,在實時例項分割以及旋轉目標檢測這兩個任務中也同樣達到了 SOTA 的水平!
除了擁有超高的精度,RTMDet 的工程優化也非常全面:基於 MMDeploy,RTMDet 全系列已經支援了 ONNXRuntime 以及 TensorRT 等推理框架的部署。同時,為了節約廣大煉丹師的寶貴時間,RTMDet 的訓練效率也同樣進行了大幅的優化,我們在 MMYOLO 中提供了高效訓練版本的程式碼,RTMDet-s 訓練 300 epoch 僅需 14 × 8 GPU 時!
檢測和例項分割模型程式碼已開源至:
https://github.com/open-mmlab/mmdetection/tree/3.x/configs/rtmdet
基於 MMYOLO 的高效訓練版本:
https://github.com/open-mmlab/mmyolo/tree/dev/configs/rtmdet
基於 MMRotate 的旋轉框檢測模型:
https://github.com/open-mmlab/mmrotate/tree/1.x/configs/rotated_rtmdet
技術報告連結:
https://arxiv.org/abs/2212.07784
話不多說,接下來就讓我們來詳細瞭解 RTMDet 在各方面的改進吧~
2. 演算法和工程優化詳解
2.1 模型結構設計
擁有一個高效的模型結構是設計實時目標檢測器最關鍵的問題之一。自 YOLOv4 將 Cross Stage Partial Network 的結構引入 DarkNet 之後,CSPDarkNet 因其簡潔高效而被廣泛應用於 YOLO 系列的各個改進版中。
而 RTMDet 也將 CSPDarkNet 作為基線,並使用同樣構建單元組成的 CSPPAFPN 進行多尺度的特徵融合,最後將特徵輸入給不同的檢測頭,進行目標檢測、例項分割和旋轉框檢測等任務。整體的模型結構如下圖所示:
圖 1. RTMDet 模型結構圖
整體的巨集觀架構與 YOLO 系列並沒有太大區別,RTMDet 對模型的修改主要聚焦於兩點:增加基本構建單元的特徵提取能力,以及探索模型不同元件間的計算效率和精度的平衡。
首先讓我們回顧一下 DarkNet 基本的構建單元,如圖 2.(a) 所示,它由 1 個 1x1 的卷積和 1 個 3x3 的卷積組成,以及殘差連線組成。為了提升基本單元的特徵提取能力,我們選擇引入更大的卷積核來增大感受野。在經過對不同 kernel size 的消融實驗(表 1)後我們發現,在基本單元中加入一個 5x5 的深度可分離卷積,取得了最好的計算效率與精度的平衡。新的基礎單元結構如 圖 2.(b) 所示。
表 1. 不同 kernel size 的精度與速度的對比實驗
值得注意的是,最近提出的 YOLO 的不同改進版,如 YOLOv6,v7,以及 PPYOLO-E,都選擇在基本單元中加入 RepVGG 的重引數化卷積模組,如圖 2.(c),圖 2.(d) 所示。
圖 2. CSPDarknet、RTMDet 、PPYOLOE、YOLOv6 的基本構建單元
重引數化通過訓練時使用多分支的結構,然後再在推理時融合成單分支的方式,在不增加推理計算量的情況下提升效能。
然而,重引數化也有其弊端,比如會明顯增加訓練的視訊記憶體開銷,同時也會導致使用低精度對模型進行量化時產生較大的效能下降。儘管可以通過 QAT 等方式解決量化掉點的問題,但這也使得從模型訓練到部署的流程變得更為複雜。
因此,我們認為引入大卷積核增加感受野的方式相比引入重引數化來說,不論是從訓練還是部署都能夠提供更好的效果。
除了基本構建單元的修改之外,RTMDet 還對整個模型的不同解析度層級之間、以及 backbone 和 neck 之間的計算量分配上進行了全面的調整。由於在基本單元中引入了深度可分離卷積,使得模型整體的層數相比於 CSPDarkNet 來說變得更深,從而導致推理速度的減慢。
為了解決這一問題,我們對模型不同解析度層級的基本單元數量進行了調整,從原本 C2~C5 分別為 3-9-9-3 個 block,調整為了 3-6-6-3 個 block
圖 3. Backbone 結構改進
同時為了保持模型的整體計算量不變,我們略微增大了模型的寬度,並且在每個 stage 之後增加一個 ChannelAttention 模組來提供通道注意力(如圖 4)。
圖 4. 增加通道注意力前後的 C5 stage 特徵對比
如表 2 所示,在經過調整之後,模型的推理速度更快,並且能夠保持相似的精度。
表 2. Backbone 結構對比實驗
而對於用來進行多尺度特徵融合的 neck 模組來說,以往的方法,如 EfficientDet、GiraffeDet 等,都是通過增加更多的特徵融合次數,也就是加入更多的連線來提升效能。
然而,過多的特徵層之間的連線會顯著增加視訊記憶體的佔用,也會因為訪存問題而增加推理耗時。因此,我們選擇僅通過增加 neck 模組的計算量佔比來提升其效能,當 backbone 與 neck 的引數量調整至相似時,整個檢測器的推理速度更快,效能也更高:
表 3. Backbone neck 引數量分配對比實驗
除了 backbone 和 neck 之外,我們對模型的檢測頭也進行了調整。YOLO 系列往往在不同的特徵層級上分別使用獨立的檢測頭,這就導致了模型引數的利用率低下。因為對於不同的特徵層級來說,其檢測出的物體的特徵,在相對的尺度大小下應當是相近的,而學術界常用的檢測器(如 RetinaNet、FCOS)使用共享引數的檢測頭也正印證了這一點。
但是,由於不同層級之間特徵的統計量仍存在差異,Normalization layer 依然是必須的,由於直接在共享引數的檢測頭中引入 BN 會導致其滑動平均值產生誤差,而引入 GN 又會增加推理時的開銷,因此我們參考 NASFPN 的做法,讓檢測頭共享卷積層,而 BN 則分別獨立計算。
在使用共享的檢測頭後,模型的引數量得到了減少,而且效能非但沒有下降,反而還得到了略微的提升:
表 4. Head 結構對比
通過採用以上這些策略,我們又對通道和深度的超引數進行縮放,得到了 tiny、s、m、l、x 五種不同大小的模型,其引數量分別從 4M 到 90M 遞增,以提供給不同的應用場景。
2.2 訓練策略優化
正負樣本的標籤分配策略是目標檢測訓練過程中最重要的一環,近幾年來,標籤分配策略從最初的基於 anchor IoU 的靜態匹配方式,逐漸演進為使用代價函式進行動態標籤分配。然而,目前主流的動態標籤分配策略,如匈牙利匹配、OTA 等均使用與損失函式一致的函式計算代價矩陣。
我們經過實驗發現,與損失函式完全一致的代價矩陣計算方式並不是最優的。因此,RTMDet 基於 YOLOX 的 SimOTA 進行改進,使用了動態的軟標籤分配策略,其代價矩陣計算公式如下:
它由三個代價函式組成,首先是分類代價函式,傳統的分類代價往往使用 0-1 分佈的二值化標籤進行損失計算,這很容易導致,一個擁有很高分類置信度但是擁有錯誤檢測框的低質量預測結果得到一個很低的分類代價,反之也是如此。
因此我們參考 GFL,將預測框與 Ground Truth 的 IoU 得分作為軟標籤,並對不同得分的匹配進行了重新加權,使得分類代價的匹配結果更為準確和穩定,其公式如下:
而對於迴歸代價,使用與損失函式一致的 GIoU,很容易導致一個低質量的匹配和一個高質量的匹配區分度不夠高,這是因為對於 GIoU 函式來說,一個完全正確的檢測框和一個完全錯誤的檢測框,他們之間的 IoU 得分也只相差了 1 而已。因此,我們選擇取對數來增大低質量與高質量匹配之間的差異:
最後,我們還引入了一個“軟化”的中心先驗代價。與之前的眾多方法使用固定的先驗區域(如 ATSS,SimOTA)不同,中心先驗代價將位置先驗也引入代價矩陣的計算之中,一方面能夠穩定收斂,另一方面也能夠使匹配的正樣本區域更為靈活。
為了公平比較,我們在標準的 ResNet-50 12 epoch 的 setting 下與其他標籤分配策略進行了對比,結果顯示我們的方法取得了最優的精度:
表 5. R50 1x setting 下的標籤分配策略效能對比
為了驗證這套方法的通用性,我們也在 300 epoch 和強資料增強的情況下與 YOLOX 的 SimOTA 進行了對比,結果顯示,我們的方法也同樣得到了更高的精度:
表 6. 使用 RTMDet-s 訓練 300 epoch 的標籤分配策略效能對比
我們使用 MMYOLO 中提供的 GradCAM++ 視覺化工具(demo/boxam_vis_demo.py
)可視化了 neck 部分的特徵圖,從下圖中可以看出,與 YOLOv6 相比,使用了動態軟標籤分配策略的 RTMDet 的檢測目標的特徵響應非常集中,且沒有誤檢:
圖 5. 使用 GradCAM++ 對 neck 部分特徵響應的視覺化(左:RTMDet-l,右:YOLOv6-l)
2.3 資料增強改進
資料增強不僅關乎模型的精度,也對訓練的效率有著巨大的影響。隨著 GPU 計算效能的不斷增加,模型前向和反向傳播的速度在不斷提升。
然而,當我們在訓練工業級的模型時,由於引入大量的資料增強,CPU 的計算能力以及儲存的 IO 往往成為了制約模型訓練速度的瓶頸。尤其是當使用了 YOLO 系列中廣泛使用的 Mosaic 和 MixUp 這樣涉及到多張圖片混合的資料增強,由於需要讀取和處理的圖片數量成倍增加,資料增強的耗時也大幅增加。
圖 6. Mosaic MixUp 資料增強效果圖
為了解決多圖混合資料增強的耗時問題,我們在這些資料增強中引入了快取機制。我們使用一個快取佇列將歷史的圖片儲存下來,當需要進行圖片混合操作時,不再通過 dataset 重新載入圖片,而是從快取佇列中隨機選取歷史的圖片進行混合。通過實驗發現,當快取佇列足夠大,且使用隨機出隊的方式時,訓練得到的模型精度與傳統的 Mosaic & MixUp 並無區別。由於引入了快取機制,兩種資料增強的執行效率得到了大幅提升。
通過使用 MMDetection 中的 benchmark 工具 tools/analysis_tools/benchmark.py
,我們將 RTMDet 與以高效訓練著稱的 YOLOv5 的資料增強進行了對比:
YOLOv5 的資料增強(MixUp 概率僅為 0.1)的吞吐量: Overall fps: 30.2 img/s, times per img: 33.1 ms/img;RTMDet 的資料增強(MixUp 全程開啟)的吞吐量:Overall fps: 52.4 img/s, times per img: 19.1 ms/img。
從 benchmark 的結果來看, RTMDet 在全程開啟影象混合的資料增強的情況下,資料處理吞吐量依然顯著高於 YOLOv5,每秒處理的圖片數量達到了 YOLOv5 的 1.7 倍!使用快取機制加速的 Mosaic 和 MixUp 已經在 MMYOLO 中得到了全面的支援,不僅 RTMDet 可以使用,YOLOv5,v6,v7上都可以通過 use_cached=True
開關直接開啟,真正做到了一鍵加速訓練!
除了優化資料增強的執行效率,我們也對資料增強的邏輯進行了優化。我們分析了 YOLOX 的資料增強後發現了其中的問題:YOLOX 使用強弱兩個階段的資料增強,但其第一個訓練階段引入了旋轉和切變等操作,導致標註框產生了比較大的誤差,使其第二個訓練階段需要對迴歸分支額外增加一個 L1 loss 來對迴歸任務進行修正。這就造成了資料增強與模型的損失函式設計產生了耦合。
為了使資料增強和模型解耦,得到更通用的增強策略,我們選擇在第一個訓練階段不引入導致標註誤差的資料增強,而是通過增加 Mosaic MixUp 混合圖片的數量來提升強度,我們將總體混合的圖片數量從 YOLOX 使用的 5 張增加至 8 張。
得益於上文提到的快取機制,增加混合圖片的數量並不會導致訓練速度變慢。在第二個訓練階段,我們關閉 Mosaic 和 MixUp,轉而使用 Large Scale Jitter(LSJ),使模型在更符合原資料集特徵分佈的狀態下微調。通過對比實驗可以看出,我們的資料增強顯著優於之前的方法:
表 7. 資料增強訓練精度對比
2.4 最終效果
通過上文的種種改進,我們最終得到了 RTMDet 系列模型,我們可以將實驗拆解,並逐個應用在我們的對照組模型 YOLOX 之上,來解析如何從 YOLOX 一步一步修改為 RTMDet:
表 8. 從 YOLOX 到 RTMDet step-by-step 對比
首先,為了解決 CosineLr 學習率衰減過快以及 SGD 收斂不穩定的問題,我們將優化器以及學習率分別更換為了 AdamW 和 FlatCosineLR,得到了 0.4% 的提升。
然後我們使用新的基礎單元構建而成的 backbone 和 neck 替換了原有的模型結構,這一改進又提升了 1.2% 的 AP,並且推理速度只減慢了 0.02ms。在使用共享權重的檢測頭後,引數量過大的問題得到了解決,而模型精度和推理速度都沒有下降。在此基礎之上,我們又增加了動態軟標籤分配策略以及改進後的資料增強,分別帶來了 1.1% 和 1.3% 的精度提升。
最後,錦上添花的一點,我們又將 backbone 在 imagenet 上進行了預訓練,也使精度得到了略微的提升。不過預訓練模型僅在 tiny 和 s 模型上有精度提升的效果,在更大的模型結構上則體現不出優勢。
綜合以上這些修改,RTMDet 相比於 YOLOX 在相似的模型大小以及推理速度上提升了 4.3% AP!
通過調整深度以及寬度的係數進行模型縮放,我們最終得到了 tiny/s/m/l/x 五種不同大小的模型,在不同量級上均超越了同級別的模型:
表 9. RTMDet 與其他同級別模型的效能對比(推理測速使用 3090 GPU,TensorRT 8.4.3)
- 另外,我們還在工業界常用的 T4 GPU 上進行了測速,效果如下表所示:
- | Model | latency(ms) | FPS | | ----------- | --------------- | ------- | | RTMDet-tiny | 2.34 | 427.35 | | RTMDet-s | 2.96 | 337.84 | | RTMDet-m | 6.41 | 156.01 | | RTMDet-l | 10.32 | 96.90 | | RTMDet-x | 18.80 | 53.19 |
表 10. RTMDet 在 T4 GPU 上的推理速度(測速使用 TensorRT 8.4,FP16,batchsize=1)
3. 多項任務取得 SOTA
為了驗證 RTMDet 演算法的通用性,我們通過僅增加任務頭的方式,對模型進行了最小限度的修改,將其拓展至了例項分割以及旋轉目標檢測任務上。在僅進行 了 非常簡單的修改的情況下,RTMDet 也依然取得了 SOTA 的效果!
圖 7. RTMDet 在三種不同任務上的視覺化效果
3.1 例項分割
傳統的例項分割往往採用雙階段的方式來預測 mask,但近兩年來,基於 kernel 的方法正逐漸興起。為了保持單階段檢測器的簡潔高效和易於部署的特性,我們參考 CondInst 為 RTMDet 增加了 mask head 以及 kernel head。
mask head 由 4 個卷積層組成,通過 neck 輸出的多尺度特徵預測出維度為 8 的 mask prototype feature。而 kernel head 則為每個 instance 預測 169 維的向量,從而組成 3 個動態卷積的卷積核,與 mask feature 進行互動,最終得到每個 instance 的 mask。
圖 8. RTMDet-Ins 例項分割分支示意圖
為了最大化利用 mask 的標註資訊,我們還將標籤分配中的中心先驗代價的檢測框中心修改為了 mask 的重心。為了與其他例項分割方法做公平對比,我們使用 ResNet50 FPN 12 epoch 的標準 setting 進行了實驗,同時參考 CondInst 加入了語義分割的輔助分支加快收斂。
實驗結果表明,我們的方法在標準 setting 下,不僅超越了 CondInst,SOLOv2 等單階段模型,也超越了 Cascade Mask R-CNN 這樣的多階段模型。在使用與 RTMDet 一致的模型結構以及訓練策略後,得到的 RTMDet-Ins 模型,精度不僅大幅超越 YOLOv5-seg,也超越了前幾天剛剛推出的 YOLOv8,取得了實時例項分割的 SOTA。
表 11. RTMDet-Ins 與其他例項分割演算法的效能對比
3.2 旋轉目標檢測
旋轉目標檢測是遙感領域使用最廣泛的方法,而將 RTMDet 拓展至此任務也非常簡單,得益於 MMDetection 3.0 版本,我們只需要增加回歸分支的輸出特徵維度,增加角度分量,同時更換 box 的編解碼器, 便能讓 RTMDet 支援預測旋轉框。由於旋轉框與普通目標檢測任務僅有迴歸分支有差異,因此旋轉框的模型也能夠載入目標檢測的預訓練模型並從中受益。
我們在遙感領域最常用的 DOTA 資料集上驗證了我們的方法,結果表明,我們的模型不僅超越了眾多的多階段方法,也超越了之前最優的演算法 PPYOLOE-R,成為了 SOTA。
圖 9. RTMDet-R 與其他旋轉目標檢測方法在 DOTA 資料集上的對比
為了驗證模型的泛化能力,我們還在 DOTA 1.5 以及 HRSC 兩個資料集上進行了訓練,也同樣取得了最優的結果。
表 12. RTMDet-R 在 DOTA 1.5 上的效能
表 13. RTMDet-R 在 HRSC 上的效能,僅使用 tiny 模型就超越了之前的多階段演算法
4. 總結
通過在模型結構的基本單元引入大 kernel 深度可分離卷積增大感受野,並平衡不同解析度層級間以及 backbone 和 neck 間的計算量、引數量平衡,改進標籤分配以及資料增強策略,RTMDet 不論是在目標檢測任務上,還是在例項分割以及旋轉目標檢測任務上,均取得了優異的效能。我們希望我們探索出的這些改進方案能夠對設計實時的目標檢測模型有所啟發,也希望我們在工程及演算法上的優化能夠在實際的工業場景中有所應用。
最後,別忘了給 MMDetection、MMYOLO、MMRotate 點個 Star 哦~
目標檢測和例項分割專案地址:https://github.com/open-mmlab/mmdetection/tree/3.x/configs/rtmdet
MMYOLO 高效訓練版本地址:https://github.com/open-mmlab/mmyolo/tree/dev/configs/rtmdet
旋轉框檢測專案地址:https://github.com/open-mmlab/mmrotate/tree/1.x/configs/rotated_rtmdet
- 目標檢測、例項分割、旋轉框樣樣精通!詳解高效能檢測演算法 RTMDet
- 商湯OpenMMLab YOLO系列工具箱;慕尼黑工大『計算機視覺深度學習進階課』;伯克利『深度無監督學習課程』;前沿論文 | ShowMeAI資訊日報
- 模型部署入門教程(一):模型部署簡介
- 開源演算法平臺OpenMMLab來了,Gitee深度學習CV領域再添一員猛將
- logging 詳解第三期:Logging 不為人知的二三事
- 手把手帶你訓練 CVPR2022 影片超分模型
- ResNet 高精度預訓練模型在 MMDetection 中的最佳實踐
- 用 OpenMMLab 輕鬆搭建主幹網路,多種視覺任務一網打盡
- PyTorch 零基礎入門 GAN 模型之 cGAN
- 三句話,讓 logger 言聽計從
- logging 詳解第一期:是誰偷偷動了我的 logger
- TorchScript 系列解讀(一):初識 TorchScript
- PyTorch 原始碼解讀之 torch.utils.data:解析資料處理全流程
- PyTorch 原始碼解讀之 BN & SyncBN:BN 與 多卡同步 BN 詳解
- PyTorch 原始碼解讀之 torch.autograd:梯度計算詳解
- PyTorch 原始碼解讀之 cpp_extension:揭祕 C /CUDA 運算元實現和呼叫全流程
- 自監督學習系列(一):基於 Pretext Task
- 困擾我 48 小時的深拷貝,今天終於...
- 旋轉~跳躍~檢測王者 MMDetection 的好兄弟來啦!
- 一戶一墩?墩墩生成器安排了!