Mamba再下一城，殺入Occupancy！更快更強(qiáng)的MambaOcc來(lái)了（中科院&美團(tuán)）

作者：Yonglin Tian等 2024-09-03 10:20:00

今天為大家分享中科院&美團(tuán)最新的占用網(wǎng)絡(luò)工作—MambaOcc！基于Mamba框架的新型占用率預(yù)測(cè)方法！

本文經(jīng)自動(dòng)駕駛之心公眾號(hào)授權(quán)轉(zhuǎn)載，轉(zhuǎn)載請(qǐng)聯(lián)系出處。

寫(xiě)在前面&筆者的個(gè)人理解

在自動(dòng)駕駛系統(tǒng)當(dāng)中，感知作為自動(dòng)駕駛車(chē)輛檢測(cè)周?chē)o態(tài)和動(dòng)態(tài)障礙物的重要途徑，其準(zhǔn)確性和穩(wěn)定性的重要程度不言而喻。然而，當(dāng)自動(dòng)駕駛汽車(chē)在開(kāi)放場(chǎng)景中行駛時(shí)，準(zhǔn)確和穩(wěn)定的感知模塊變得尤其具有挑戰(zhàn)性，因?yàn)樗鼈儽仨殭z測(cè)看不見(jiàn)的或不規(guī)則的物體。

最近，由于占用預(yù)測(cè)具有更加細(xì)粒度的通用感知能力，它在自動(dòng)駕駛系統(tǒng)的感知和規(guī)劃流程中引起了越來(lái)越廣泛的關(guān)注。雖然占用預(yù)測(cè)網(wǎng)絡(luò)相比于原有的3D障礙物感知算法具有更加細(xì)粒度和通用的感知能力，但是占用預(yù)測(cè)網(wǎng)絡(luò)需要確定當(dāng)前感知的3D場(chǎng)景中每個(gè)體素的狀態(tài)，這會(huì)導(dǎo)致感知模型開(kāi)發(fā)過(guò)程中對(duì)計(jì)算和內(nèi)存的需求很高。

基于柵格預(yù)測(cè)的相關(guān)優(yōu)勢(shì)以及現(xiàn)有存在的諸多不足，我們的核心思路是提高基于BEV空間的占用預(yù)測(cè)的性能，同時(shí)減少參數(shù)數(shù)量和計(jì)算成本?？紤]到先前的研究工作中強(qiáng)調(diào)了Transformers算法模型在長(zhǎng)距離建模方面的優(yōu)勢(shì)，但它們的計(jì)算負(fù)擔(dān)也非常的大。最近，狀態(tài)空間模型 (SSM)（例如 Mamba）已成為長(zhǎng)距離建模的更有效解決方案。這一發(fā)展促使我們探索狀態(tài)空間模型在改進(jìn)占用預(yù)測(cè)任務(wù)方面的潛力。

因此，本文提出了一種基于Mamba框架的新型占用率預(yù)測(cè)方法，旨在實(shí)現(xiàn)輕量級(jí)，同時(shí)提供高效的遠(yuǎn)距離信息建模，我們稱之為MambaOcc算法模型。

論文鏈接：https://arxiv.org/pdf/2408.11464

網(wǎng)絡(luò)模型的整體架構(gòu)&細(xì)節(jié)梳理

在詳細(xì)介紹本文提出的MambaOcc算法模型的技術(shù)細(xì)節(jié)之前，下圖展示了我們提出的MambaOcc算法的整體網(wǎng)絡(luò)結(jié)構(gòu)。通過(guò)下圖可以看出，MambaOcc算法模型主要包括四個(gè)模塊，分別是基于Mamba的圖像主干網(wǎng)絡(luò)用于圖像特征的提取，用于獲取BEV形式特征和聚合多幀特征的視角轉(zhuǎn)換模塊以及時(shí)間融合模塊，帶有自適應(yīng)局部重排模塊LAR-SS2S混合BEV編碼器模塊以及占用預(yù)測(cè)頭模塊。

MambaOcc算法模型的整體網(wǎng)絡(luò)結(jié)構(gòu)

整體而言，我們采用四種方向的視覺(jué)Mamba來(lái)提取圖像特征。同時(shí)為了減輕與3D體素相關(guān)的高計(jì)算負(fù)擔(dān)，我們使用BEV特征作為占用預(yù)測(cè)的中間特征表示，并設(shè)計(jì)了結(jié)合卷積層和Mamba層的混合BEV編碼器?？紤]到Mamba架構(gòu)在特征提取過(guò)程中對(duì)標(biāo)記排序的敏感性，我們引入了一個(gè)利用可變形卷積層的局部自適應(yīng)重新排序模塊。該模塊旨在動(dòng)態(tài)更新每個(gè)位置的上下文信息，使模型能夠更好地捕獲和利用數(shù)據(jù)中的局部依賴關(guān)系。這種方法不僅可以緩解標(biāo)記序列相關(guān)的問(wèn)題，還可以通過(guò)確保在提取過(guò)程中優(yōu)先考慮相關(guān)的上下文信息來(lái)提高占用預(yù)測(cè)的整體準(zhǔn)確性。

VM-Backbone（視覺(jué)Mamba主干網(wǎng)絡(luò)）

View Transformation and Temporal Fusion（視角轉(zhuǎn)換以及時(shí)序融合）

在MambaOcc算法模型中，我們采用LSS算法模型實(shí)現(xiàn)從圖像平面到BEV平面的空間視圖變換。首先，將圖像主干網(wǎng)絡(luò)的輸出特征組織成2D格式的地圖。然后通過(guò)深度預(yù)測(cè)網(wǎng)絡(luò)生成每個(gè)像素的一系列離散深度。最后，使用體素池化在預(yù)定義的BEV平面上聚合每個(gè)網(wǎng)格內(nèi)的深度預(yù)測(cè)。

在使用了時(shí)間域融合的情況下，視角轉(zhuǎn)換模塊提供了一種方便的方式融合來(lái)自不同視角以及不同時(shí)間戳的圖像特征。利用來(lái)自前幾幀的 BEV空間特征，首先基于自車(chē)運(yùn)動(dòng)信息進(jìn)行特征轉(zhuǎn)換操作。然后，應(yīng)用采樣和插值操作來(lái)生成與當(dāng)前幀BEV空間特征圖對(duì)齊的特征。最后，將對(duì)齊的特征合并到一起來(lái)實(shí)現(xiàn)融合時(shí)序的上下文特征信息。

LAR-SS2D Hybrid BEV Encoder（LAR-SS2D混合BEV編碼器）

在BEV空間特征的提取方面，我們首先設(shè)計(jì)了基于Mamba網(wǎng)絡(luò)模型的架構(gòu)，該架構(gòu)由三個(gè)塊組成，其中每個(gè)塊包含兩個(gè)SS2D組?？紤]到 SSM層對(duì)序列中token的順序很敏感，我們進(jìn)一步探索局部自適應(yīng)偽重排序機(jī)制來(lái)優(yōu)化上下文信息的嵌入。然后使用LAR組替換每個(gè)塊中的一個(gè)SS2D組。

通過(guò)上述的修改，我們建立了一個(gè)靈活的局部偽重排序機(jī)制。此外，所提出的重排序過(guò)程可以通過(guò)可變形卷積算子高效實(shí)現(xiàn)，從而確保較高的計(jì)算效率并保持較快的處理速度。

Occupancy預(yù)測(cè)頭（Occupancy Prediction Head）

我們采用了FlashOcc算法當(dāng)中的實(shí)現(xiàn)思路，我們同樣使用了channel-to-height操作從生成的BEV特征圖的通道維度當(dāng)中恢復(fù)出高度信息。這個(gè)過(guò)程允許我們?cè)谡麄€(gè)網(wǎng)絡(luò)的最后來(lái)獲得3D占用的特征表達(dá)。隨后，我們使用線性層來(lái)預(yù)測(cè)3D空間中每個(gè)位置的類(lèi)別，從而提供整個(gè)3D空間中詳細(xì)完整的占用預(yù)測(cè)信息。

實(shí)驗(yàn)結(jié)果&評(píng)價(jià)指標(biāo)

與其它SOTA算法的對(duì)比試驗(yàn)

為了驗(yàn)證我們提出的MambaOcc算法模型的有效性，我們?cè)贠cc3D-nuScenes數(shù)據(jù)集上進(jìn)行了相關(guān)實(shí)驗(yàn)，相關(guān)的實(shí)驗(yàn)結(jié)果如下表所示。

不同算法模型在Occ3D-nuScenes數(shù)據(jù)集上的結(jié)果匯總

通過(guò)實(shí)驗(yàn)結(jié)果可以看出，與最先進(jìn)的方法相比，我們提出的MambaOcc算法模型在計(jì)算效率和參數(shù)數(shù)量方面具有更加顯著的優(yōu)勢(shì)。與以Swin-Transformer為主干網(wǎng)絡(luò)的FlashOcc算法模型相比，MambaOcc取得了更好的性能，同時(shí)減少了42%的參數(shù)和39%的計(jì)算成本。此外，MambaOcc-Large比FlashOcc高出了0.77的mIoU，減少了14%的參數(shù)和32%的計(jì)算成本。與以ResNet-101為主干網(wǎng)絡(luò)的PanoOcc算法相比，MambaOcc的性能高出1.23 mIoU，同時(shí)減少了19%的參數(shù)。這些結(jié)果均表明，與基于CNN和Transformer的方法相比，所提出的Mamba框架在參數(shù)量、計(jì)算效率和感知能力方面具有顯著優(yōu)勢(shì)。

此外，為了更加直觀的展示我們提出的MambaOcc算法模型的有效性，下圖可視化了MambaOcc算法模型占用預(yù)測(cè)的結(jié)果。如圖所示，MambaOcc可以為人類(lèi)和車(chē)輛等典型物體提供精確的感知結(jié)果，同時(shí)還能有效檢測(cè)電線桿、交通燈和路錐等結(jié)構(gòu)不規(guī)則的物體。

MambaOcc算法模型的占用預(yù)測(cè)結(jié)果可視化

此外，為了直觀的展示我們提出的MambaOcc和FlashOcc預(yù)測(cè)占用的效果對(duì)比，我們也對(duì)兩個(gè)模型的結(jié)果進(jìn)行了可視化，如下圖所示。

MambaOcc與FlashOcc結(jié)果可視化對(duì)比

通過(guò)可視化兩個(gè)模型的預(yù)測(cè)結(jié)果對(duì)比可以看出，MambaOcc算法模型在長(zhǎng)距離平面感知方面的卓越性能，能夠提供更全面的地面預(yù)測(cè)，而FlashOcc算法模型通常會(huì)將這些區(qū)域預(yù)測(cè)為空。

消融對(duì)比實(shí)驗(yàn)

為了清楚地展示提出的MambaOcc算法模型中每個(gè)組件的貢獻(xiàn)，我們?cè)谙卤碇姓故玖宋覀冞M(jìn)行的消融研究結(jié)果，以強(qiáng)調(diào)每個(gè)模塊的有效性。

通過(guò)匯總的消融實(shí)驗(yàn)結(jié)果可以看出，用Mamba網(wǎng)絡(luò)結(jié)構(gòu)替換CNN網(wǎng)絡(luò)架構(gòu)，可使mIoU顯著增加3.96，凸顯了Mamba網(wǎng)絡(luò)架構(gòu)的有效性。此外，我們提出的LAR-SS2D BEV編碼器模塊比基于CNN的編碼器額外增加了1.12的mIoU。此外，通過(guò)結(jié)合位置編碼，可以進(jìn)一步提高模型的預(yù)測(cè)性能。

此外，我們也進(jìn)行了相關(guān)實(shí)驗(yàn)來(lái)驗(yàn)證不同的圖像主干網(wǎng)絡(luò)初始化方法對(duì)于網(wǎng)絡(luò)模型占用預(yù)測(cè)的影響效果，相關(guān)的實(shí)驗(yàn)結(jié)果如下表所示。

通過(guò)實(shí)驗(yàn)結(jié)果可以明顯的看出，良好的參數(shù)初始化方法會(huì)顯著影響性能。使用ImageNet分類(lèi)預(yù)訓(xùn)練初始化占用預(yù)測(cè)網(wǎng)絡(luò)與隨機(jī)初始化相比，對(duì)于Mamba和卷積網(wǎng)絡(luò)，效果明顯更好。例如與使用隨機(jī)值初始化的相比，使用ImageNet預(yù)訓(xùn)練的VM-Backbone的MambaOcc在mIoU方面的性能高出10.01。

我們也對(duì)不同的BEV編碼器對(duì)于占用預(yù)測(cè)任務(wù)的影響進(jìn)行了相關(guān)的實(shí)驗(yàn)，實(shí)驗(yàn)結(jié)果匯總在下表所示。

如上表所示，BEV編碼器的結(jié)構(gòu)顯著影響了占用預(yù)測(cè)性能。純SS2D優(yōu)于純CNN，mIoU指標(biāo)提高了0.56?；旌螩NN-SS2D網(wǎng)絡(luò)架構(gòu)的性能優(yōu)于純CNN和純SS2D的網(wǎng)絡(luò)架構(gòu)，mIoU分別提高了0.77和0.21。所提出的LAR-SS2D混合架構(gòu)取得了最佳效果，比CNN-SS2D混合架構(gòu)高出0.48 mIoU。

除此之外，我們比較了LAR層中不同映射方法的效果。對(duì)于多對(duì)一映射，我們使用不同的條目數(shù)3×3和5×5進(jìn)行了實(shí)驗(yàn)，其中原始序列中多個(gè)位置的信息在映射到新序列中的相同位置之前進(jìn)行加權(quán)和融合。相關(guān)的實(shí)驗(yàn)結(jié)果匯總在下表中。