BEVFormer-accelerate:基於 EasyCV 加速 BEVFormer

導言

BEVFormer是一種純視覺的自動駕駛感知算法,通過融合環視相機圖像的空間和時序特徵顯式的生成具有強表徵能力的BEV特徵,並應用於下游3D檢測、分割等任務,取得了SOTA的結果。我們在EasyCV開源框架(https://github.com/alibaba/EasyCV)中,對BEVFomer算法進行集成,並從訓練速度、算法收斂速度角度對代碼進行了一些優化。同時,我們進一步使用推理優化工具PAI-Blade對模型進行優化,相比於原始模型在A100配置下能取得40%的推理速度提升。本文將從以下幾個部分進行介紹:1、BEVFormer算法思想 2、訓練速度和算法收斂速度優化 3、使用PAI-Blade優化推理速度。

BEVFormer算法思想

如上圖所示,BEVFormer由如下三個部分組成:

  1. backbone:用於從6個角度的環視圖像中提取多尺度的multi-camera feature
  2. BEV encoder:該模塊主要包括Temporal self-Attention 和 Spatial Cross-Attention兩個部分。
    1. Spatial Cross-Attention結合多個相機的內外參信息對對應位置的multi-camera feature進行query,從而在統一的BEV視角下將multi-camera feature進行融合。
    2. Temporal self-Attention將History BEV feature和 current BEV feature通過 self-attention module進行融合。
    3. 通過上述兩個模塊,輸出同時包含多視角和時序信息的BEV feature進一步用於下游3D檢測和分割任務
  3. Det&Seg Head:用於特定任務的task head

BEVFormer訓練優化

訓練加速優化

我們從數據讀取和減少內存拷貝消耗等角度對訓練代碼進行優化。

  • 數據讀取
    • 使用更高效的圖片解碼庫 turbojpeg
    • BEVFormer在訓練過程中,需要時序上的數據作爲輸入,將串形的讀取方式優化爲並行讀取。
    • 先做resize再做其他預處理,減少了額外像素帶來的計算開銷
  • 內存拷貝優化
    • 使用pin_memery=True,並修復了mmcv DataContainer pin_memory的bug
    • 將代碼中的numpy操作替換爲torch.tensor,避免不必要的h2d拷貝
  • other
    • 使用torch.backends.cudnn.benchmark=True(ps:需要保證在輸入數據沒有動態性的情況下使用,否則反而會增加訓練耗時)
    • 修復了torch.cuda.amp混合精度在LayerNorm層失效的bug

我們在A100 80G的機器上,使用fp16對比吞吐量如下:

Setting throughput(samples/s)
BEVFormer-tiny bs=32 3.55
EasyCV BEVFormer-tiny bs=32 9.84(+177%)
BEVFormer-base bs=5 0.727
EasyCV BEVFormer-base bs=5 0.8(+10%)

精度收斂優化

我們使用額外的數據增廣方式和不同的損失函數來優化模型。同時加入額外的訓練策略來進一步提升模型收斂速度及精度。

  • 數據增廣方式
    • rand scale(採用不同分辨率的輸入進行訓練,實驗中發現該操作會引入至少20%的額外訓練時間,因此在下述實驗中,均沒有采用)
    • rand_flip(以50%的概率隨機翻轉圖片)
  • 損失函數
    • 使用smooth l1 loss或 balance l1 loss代替l1 loss。(在mini dataset的實驗中,這兩個損失都可以提升精度,下面的實驗中採用balance l1 loss)
  • 訓練策略
    • 使用one2many Branch

這個做法來自於H-Deformable-DETR,在DETR系列的檢測模型中採用one2one的匹配方式來分配GT Boxes,這種做法雖然讓模型在測試的時候,能夠避免冗餘的NMS後處理操作,但是隻有少數的Query會被分配給正樣本,導致訓練時模型收斂速度相比於one2many的方式會慢很多。因此,在訓練過程中加入auxiliary Query,同一個GT Box會匹配多個auxiliary Query,並使用attention mask將one2one branch和one2many branch的信息隔離開。通過這樣的方式,能夠顯著的提升訓練過程中的收斂速度,同時在測試過程中只需要保持one2one branch進行預測。(在實驗中,使用額外加入1800個auxiliary Query,每個GT box匹配4個query進行訓練)

    • CBGS in one2many Branch

我們的實驗是在NuScenes數據集上進行的,在該數據集的3D檢測任務上有10類標籤,但是這10類標籤之間的樣本極度不均衡,很多算法會採用CBGS操作進行類間樣本均衡,但是這個操作會將整個數據集擴大4.5倍,雖然有一定的精度提升,但是也帶來了巨大的訓練成本。我們考慮在one2many Branch上進行樣本均衡操作,即對於實例數量較多的樣本使用較少的auxiliary Query進行匹配,而對於長尾的樣本使用較多的auxiliary Query進行匹配。通過CBGS in one2many Branch的方式,訓練時間和base保持一致的基礎上會進一步提升收斂速度,最終的精度也有一定的提升。(實驗中匹配框數量變化:[4, 4, 4, 4, 4, 4, 4, 4, 4, 4] -> [2, 3, 7, 7, 9, 6, 7, 6, 2, 5])

我們在單機8卡A100 80G下進行實驗,如下表所示:

config setting NDS mAP throughput(samples/s)
官方 BEVFormer-base 52.44 41.91 3.289
EasyCV BEVFormer-base 52.66 42.13 3.45
EasyCV BEVFormer-base-one2manybranch 53.02(+0.58) 42.48(+0.57) 3.40
EasyCV BEVFormer-base-cbgs_one2manybranch 53.28(+0.84) 42.63(+0.72) 3.41

模型收斂速度如下圖所示:

由上圖可以看出,使用上述優化方式可以大幅提升模型收斂速度,僅需要75%的訓練時間就可以達到base的最終精度。同時最終的NDS相比於base也有0.8的提升。

詳細配置,訓練log和模型權重,參考:https://github.com/alibaba/EasyCV/blob/master/docs/source/model_zoo_det3d.md

在阿里雲機器學習平臺PAI上使用BEVFormer模型

PAI-DSW(Data Science Workshop)是阿里雲機器學習平臺PAI開發的雲上IDE,面向各類開發者,提供了交互式的編程環境。在DSW Gallery中(鏈接),提供了各種Notebook示例,方便用戶輕鬆上手DSW,搭建各種機器學習應用。我們也在DSW Gallery中上架了BEVFormer進行3D檢測的Sample Notebook(見下圖),歡迎大家體驗!

使用PAI-Blade進行推理加速

PAI-Blade是由阿里雲機器學習平臺PAI開發的模型優化工具,可以針對不同的設備不同模型進行推理加速優化。PAI-Blade遵循易用性,魯棒性和高性能爲原則,將模型的部署優化進行高度封裝,設計了統一簡單的API,在完成Blade環境安裝後,用戶可以在不瞭解ONNX、TensorRT、編譯優化等技術細節的條件下,通過簡單的代碼調用方便的實現對模型的高性能部署。更多PAI-Blade相關技術介紹可以參考 [PAI-Blade介紹]

PAI-EasyCV中對Blade進行了支持,用戶可以通過PAI-EasyCV的訓練config 中配置相關export 參數,從而對訓練得到的模型進行導出。

對於BEVFormer模型,我們在A100機器下進行進行推理速度對比,使用PAI-Blade優化後的模型能取得42%的優化加速。

Name Backend Median(FPS) Mean(FPS) Median(ms) Mean(ms)
easycv TensorRT 3.68697 3.68651 0.271226 0.271259
easycv script TensorRT 3.8131 3.79859 0.262254 0.26337
blade TensorRT 5.40248 5.23383(+42%) 0.1851 0.192212

環境準備

我們提供一個PAI-Blade + PAI-EasyCV 的鏡像包供用戶可以直接使用,鏡像包地址:easycv-blade-torch181-cuda111.tar

用戶也可以基於Blade每日發佈的鏡像自行搭建推理環境 [PAI-Blade社區鏡像發佈]

自行搭建環境時需要注意:BEVFomer-base使用resnet101-dcn作爲image backbone,DCN算子使用的是mmcv中的自定義算子,爲了導出TorchScript,我們對該接口進行了修改。所以mmcv需要源碼編譯。

  1. clone mmcv源碼
$ git clone https://github.com/open-mmlab/mmcv.git
  1. 替換mmcv文件

替換時請注意mmcv的版本,注意接口要匹配。mmcv1.6.0版本已驗證。

參考easycv/thirdparty/mmcv/目錄下的修改文件。用mmcv/ops/csrc/pytorch/modulated_deform_conv.cppmmcv/ops/modulated_deform_conv.py去替換mmcv中的原文件。

  1. 源碼編譯

mmcv源碼編譯請參考:https://mmcv.readthedocs.io/en/latest/get_started/build.html

導出Blade模型

導出Blade的模型的配置可以參考文件bevformer_base_r101_dcn_nuscenes.py中的export字段,配置如下:

export = dict(
    type='blade',
    blade_config=dict(
        enable_fp16=True,
        fp16_fallback_op_ratio=0.0,
        customize_op_black_list=[
            'aten::select', 'aten::index', 'aten::slice', 'aten::view',
            'aten::upsample', 'aten::clamp'
        ]
    )
)

導出命令:

$ cd ${EASYCV_ROOT}
$ export PYTHONPATH='./'
$ python tools/export.py configs/detection3d/bevformer/bevformer_base_r101_dcn_nuscenes.py bevformer_base.pth bevformer_export.pth

Blade模型推理

推理腳本:

from easycv.predictors import BEVFormerPredictor

blade_model_path = 'bevformer_export.pth.blade'
config_file = 'configs/detection3d/bevformer/bevformer_base_r101_dcn_nuscenes.py'

predictor = BEVFormerPredictor(
            model_path=blade_model_path,
            config_file=config_file,
            model_type='blade',
      )

inputs_file = 'nuscenes_infos_temporal_val.pkl'  # 以NuScenes val數據集文件爲例
input_samples = mmcv.load(inputs_file)['infos']
predict_results = predictor(input_samples)
print(predict_results)

NuScenes數據集準備請參考:NuScenes數據集準備

展望

我們在EasyCV框架中,集成了BEVFormer算法,並從訓練加速、精度收斂和推理加速角度對算法進行了一些改進。近期,也湧現了許多新的BEV感知算法,如BEVFormerv2。在BEVFormerv2中通過Perspective Supervision的方式,讓算法能夠不受限於使用一些在深度估計或3D檢測上的預訓練backbone,而直接使用近期更有效的大模型BackBone(如ConvNext、DCNv3等),同時採用two-stage的檢測方式進一步增強模型能力,在Nuscenes數據集的camera-based 3D檢測任務取得sota的結果。

原文鏈接

本文爲阿里雲原創內容,未經允許不得轉載。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章