使用 TensorFlow 3D 理解 3D 場(chǎng)景
在過(guò)去幾年中 ,3D 傳感器(例如,激光雷達(dá)、深度感應(yīng)攝像頭和雷達(dá))的日益普及催生了對(duì)能夠處理這些設(shè)備捕獲的數(shù)據(jù)的場(chǎng)景理解技術(shù)的需求。這種技術(shù)可以使使用這些傳感器的機(jī)器學(xué)習(xí) (ML) 系統(tǒng)(如自動(dòng)駕駛汽車和機(jī)器人)在現(xiàn)實(shí)世界中導(dǎo)航和操作,并可以在移動(dòng)設(shè)備上創(chuàng)建改進(jìn)的增強(qiáng)現(xiàn)實(shí)體驗(yàn)。計(jì)算機(jī)視覺(jué)領(lǐng)域最近開(kāi)始在 3D 場(chǎng)景理解方面取得良好進(jìn)展,包括用于移動(dòng) 3D 物體檢測(cè)、透明物體檢測(cè)的模型等,但由于可應(yīng)用于 3D 數(shù)據(jù)的可用工具和資源有限,因此進(jìn)入該領(lǐng)域可能具有挑戰(zhàn)性。
為了進(jìn)一步提高對(duì) 3D 場(chǎng)景的理解并降低感興趣的研究人員的入門門檻,我們發(fā)布了TensorFlow 3D (TF 3D),這是一個(gè)高度模塊化且高效的庫(kù),旨在將 3D 深度學(xué)習(xí)功能引入 TensorFlow。TF 3D 提供了一組流行的操作、損失函數(shù)、數(shù)據(jù)處理工具、模型和指標(biāo),使更廣泛的研究社區(qū)能夠開(kāi)發(fā)、訓(xùn)練和部署最先進(jìn)的 3D 場(chǎng)景理解模型。
TF 3D 包含用于最先進(jìn)的 3D語(yǔ)義分割、3D對(duì)象檢測(cè)和 3D實(shí)例分割的訓(xùn)練和評(píng)估管道,并支持分布式訓(xùn)練。它還支持其他潛在應(yīng)用,如 3D 對(duì)象形狀預(yù)測(cè)、點(diǎn)云配準(zhǔn)和點(diǎn)云致密化。此外,它還為標(biāo)準(zhǔn) 3D 場(chǎng)景理解數(shù)據(jù)集的訓(xùn)練和評(píng)估提供了統(tǒng)一的數(shù)據(jù)集規(guī)范和配置。它目前支持Waymo Open、ScanNet和Rio數(shù)據(jù)集。但是,用戶可以自由轉(zhuǎn)換其他流行的數(shù)據(jù)集,例如NuScenes和吉滴,到了類似的格式,并在現(xiàn)有的預(yù)或自定義創(chuàng)建管道使用它們,并可以利用TF 3D的各種3D的深度學(xué)習(xí)研究和應(yīng)用,從快速原型和嘗試新的想法,部署實(shí)時(shí)間推理系統(tǒng)。
左側(cè)顯示了來(lái)自Waymo Open Dataset的幀上 TF 3D 中的 3D 對(duì)象檢測(cè)模型的示例輸出 。 右側(cè)顯示了來(lái)自ScanNet 數(shù)據(jù)集的場(chǎng)景上的 3D 實(shí)例分割模型的示例輸出 。
在這里,我們將展示 TF 3D 中提供的高效且可配置的稀疏卷積主干,這是在各種 3D 場(chǎng)景理解任務(wù)上獲得最先進(jìn)結(jié)果的關(guān)鍵。此外,我們將介紹 TF 3D 當(dāng)前支持的三個(gè)管道中的每一個(gè):3D 語(yǔ)義分割、3D 對(duì)象檢測(cè)和 3D 實(shí)例分割。
3D 稀疏卷積網(wǎng)絡(luò)
傳感器捕獲的 3D 數(shù)據(jù)通常包含一個(gè)場(chǎng)景,該場(chǎng)景包含一組感興趣的對(duì)象(例如汽車、行人等),這些對(duì)象主要被開(kāi)放空間包圍,這是有限的(或沒(méi)有)興趣。因此,3D 數(shù)據(jù)本質(zhì)上是稀疏的。在這樣的環(huán)境中,卷積的標(biāo)準(zhǔn)實(shí)現(xiàn)將是計(jì)算密集型的并消耗大量?jī)?nèi)存。因此,在 TF 3D 中,我們使用子流形稀疏卷積和池化操作,旨在更有效地處理 3D 稀疏數(shù)據(jù)。稀疏卷積模型是應(yīng)用于大多數(shù)戶外自動(dòng)駕駛(例如 Waymo、NuScenes)和室內(nèi)基準(zhǔn)測(cè)試(例如 ScanNet)的最先進(jìn)方法的核心。
我們還使用各種CUDA技術(shù)來(lái)加速計(jì)算(例如,散列、在共享內(nèi)存中分區(qū)/緩存過(guò)濾器,以及使用位操作)。在 Waymo Open 數(shù)據(jù)集上的實(shí)驗(yàn)表明,這種實(shí)現(xiàn)比具有預(yù)先存在的 TensorFlow 操作的精心設(shè)計(jì)的實(shí)現(xiàn)快 20 倍左右。
TF 3D 然后使用 3D 子流形稀疏U-Net 架構(gòu)為每個(gè)體素提取一個(gè)特征。U-Net 架構(gòu)已被證明是有效的,它讓網(wǎng)絡(luò)同時(shí)提取粗略特征和精細(xì)特征,并將它們組合起來(lái)進(jìn)行預(yù)測(cè)。U-Net 網(wǎng)絡(luò)由三個(gè)模塊組成,一個(gè)編碼器、一個(gè)瓶頸和一個(gè)解碼器,每個(gè)模塊由多個(gè)稀疏卷積塊組成,可能具有池化或反池化操作。

3D 稀疏體素 U-Net 架構(gòu)。請(qǐng)注意,水平箭頭接收體素特征并對(duì)其應(yīng)用子流形稀疏卷積。向下移動(dòng)的箭頭執(zhí)行子流形稀疏池。向上移動(dòng)的箭頭將收集匯集的特征,將它們與來(lái)自水平箭頭的特征連接起來(lái),并對(duì)連接的特征執(zhí)行子流形稀疏卷積。 上述稀疏卷積網(wǎng)絡(luò)是 TF 3D 中提供的 3D 場(chǎng)景理解管道的主干。下面描述的每個(gè)模型都使用這個(gè)主干網(wǎng)絡(luò)來(lái)提取稀疏體素的特征,然后添加一個(gè)或多個(gè)額外的預(yù)測(cè)頭來(lái)推斷感興趣的任務(wù)。用戶可以通過(guò)改變編碼器/解碼器層數(shù)和每層的卷積數(shù),以及通過(guò)修改卷積濾波器的大小來(lái)配置 U-Net 網(wǎng)絡(luò),從而可以通過(guò)不同的主干配置
3D語(yǔ)義分割
的3D語(yǔ)義分割模型僅具有一個(gè)輸出頭用于預(yù)測(cè)每個(gè)體素的語(yǔ)義分?jǐn)?shù),其被映射回點(diǎn)來(lái)預(yù)測(cè)每點(diǎn)一個(gè)語(yǔ)義標(biāo)簽。
來(lái)自ScanNet 數(shù)據(jù)集的室內(nèi)場(chǎng)景的 3D 語(yǔ)義分割 。
3D Instance Segmentation
在3D Instance Segmentation 中,除了預(yù)測(cè)語(yǔ)義之外,目標(biāo)是將屬于同一對(duì)象的體素分組在一起。TF 3D 中使用的 3D 實(shí)例分割算法基于我們之前使用深度度量學(xué)習(xí)進(jìn)行2D 圖像分割的工作。該模型預(yù)測(cè)每個(gè)體素的實(shí)例嵌入向量以及每個(gè)體素的語(yǔ)義分?jǐn)?shù)。實(shí)例嵌入向量將體素映射到一個(gè)嵌入空間,其中對(duì)應(yīng)于同一對(duì)象實(shí)例的體素靠近在一起,而對(duì)應(yīng)于不同對(duì)象的體素相距很遠(yuǎn)。在這種情況下,輸入是點(diǎn)云而不是圖像,它使用 3D 稀疏網(wǎng)絡(luò)而不是 2D 圖像網(wǎng)絡(luò)。在推理時(shí),貪心算法 一次選擇一個(gè)實(shí)例種子,并使用體素嵌入之間的距離將它們分組。
3D 對(duì)象檢測(cè)
3D 對(duì)象檢測(cè)模型預(yù)測(cè)每個(gè)體素的大小、中心和旋轉(zhuǎn)矩陣以及對(duì)象語(yǔ)義分?jǐn)?shù)。在推理時(shí),使用框提議機(jī)制將數(shù)十萬(wàn)個(gè)每體素框預(yù)測(cè)減少為幾個(gè)準(zhǔn)確的框提議,然后在訓(xùn)練時(shí),將框預(yù)測(cè)和分類損失應(yīng)用于每體素預(yù)測(cè)。我們應(yīng)用Huber 損失關(guān)于預(yù)測(cè)框角和真實(shí)框角之間的距離。由于根據(jù)大小、中心和旋轉(zhuǎn)矩陣估計(jì)框角的函數(shù)是可微的,因此損失將自動(dòng)傳播回那些預(yù)測(cè)的對(duì)象屬性。我們使用動(dòng)態(tài)框分類損失將與真實(shí)情況強(qiáng)烈重疊的框分類為正,將非重疊框分類為負(fù)。
我們?cè)?ScanNet 數(shù)據(jù)集上的 3D 對(duì)象檢測(cè)結(jié)果。 在我們最近的論文“ DOPS:學(xué)習(xí)檢測(cè) 3D 對(duì)象并預(yù)測(cè)其 3D 形狀”中,我們?cè)敿?xì)描述了用于 TF 3D 中對(duì)象檢測(cè)的單階段弱監(jiān)督學(xué)習(xí)算法。此外,在后續(xù)工作中,我們通過(guò)提出基于稀疏 LSTM 的多幀模型來(lái)擴(kuò)展 3D 對(duì)象檢測(cè)模型以利用時(shí)間信息。我們繼續(xù)證明,在Waymo Open 數(shù)據(jù)集中,這個(gè)時(shí)間模型比逐幀方法的性能高 7.5% 。

DOPS論文中介紹的3D物體檢測(cè)和形狀預(yù)測(cè)模型 。3D 稀疏 U-Net 用于為每個(gè)體素提取特征向量。對(duì)象檢測(cè)模塊使用這些特征來(lái)提出 3D 框和語(yǔ)義分?jǐn)?shù)。同時(shí),網(wǎng)絡(luò)的另一個(gè)分支預(yù)測(cè)用于輸出每個(gè)對(duì)象的網(wǎng)格的形狀嵌入。 準(zhǔn)備好開(kāi)始了嗎? 我們當(dāng)然發(fā)現(xiàn)此代碼庫(kù)對(duì)我們的 3D 計(jì)算機(jī)視覺(jué)項(xiàng)目很有用,我們希望您也能如此。歡迎對(duì)代碼庫(kù)做出貢獻(xiàn),請(qǐng)繼續(xù)關(guān)注我們對(duì)框架的進(jìn)一步更新。要開(kāi)始使用,請(qǐng)?jiān)L問(wèn)我們的 github 存儲(chǔ)庫(kù)。