BEVFormer-accelerate：基於EasyCV加速BEVFormer

語言: CN / TW / HK
時間 2022-12-13 16:02:17
主題:

作者：賀弘 夕陌 謙言 臨在

導言

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

BEVFormer演算法思想

image.png

如上圖所示，BEVFormer由如下三個部分組成：

  • backbone：用於從6個角度的環檢視像中提取多尺度的multi-camera feature
  • BEV encoder：該模組主要包括Temporal self-Attention 和 Spatial Cross-Attention兩個部分。
    • Spatial Cross-Attention結合多個相機的內外參資訊對對應位置的multi-camera feature進行query，從而在統一的BEV視角下將multi-camera feature進行融合。
    • Temporal self-Attention將History BEV feature和 current BEV feature通過 self-attention module進行融合。
    • 通過上述兩個模組，輸出同時包含多視角和時序資訊的BEV feature進一步用於下游3D檢測和分割任務
  • 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

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

image.png

由上圖可以看出，使用上述優化方式可以大幅提升模型收斂速度，僅需要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（見下圖），歡迎大家體驗！

image.png

使用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.cpp和mmcv/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的結果。

EasyCV（https://github.com/alibaba/EasyCV）會持續跟進業界sota方法，歡迎大家關注和使用，歡迎大家各種維度的反饋和改進建議以及技術討論，同時我們十分歡迎和期待對開源社群建設感興趣的同行一起參與共建。

EasyCV往期分享

EasyCV開源地址：https://github.com/alibaba/EasyCV

使用 EasyCV Mask2Former輕鬆實現影象分割 https://zhuanlan.zhihu.com/p/585259059

EasyCV DataHub 提供多領域視覺資料集下載，助力模型生產 https://zhuanlan.zhihu.com/p/572593950

EasyCV帶你復現更好更快的自監督演算法-FastConvMAE https://zhuanlan.zhihu.com/p/566988235

基於EasyCV復現DETR和DAB-DETR，Object Query的正確開啟方式 https://zhuanlan.zhihu.com/p/543129581

基於EasyCV復現ViTDet：單層特徵超越FPN https://zhuanlan.zhihu.com/p/528733299

MAE自監督演算法介紹和基於EasyCV的復現 https://zhuanlan.zhihu.com/p/515859470

EasyCV開源｜開箱即用的視覺自監督+Transformer演算法庫 https://zhuanlan.zhihu.com/p/50521999

「」