谷歌AI研發(fā)TensorFlow3D操作速度竟提高二十倍

導(dǎo)讀

Google AI發(fā)布了TensorFlow 3D，將3D深度學(xué)習(xí)能力引入TensorFlow，加入3D稀疏卷積網(wǎng)絡(luò)，在Waymo Open數(shù)據(jù)集上的實驗表明，這種實現(xiàn)比預(yù)先設(shè)計好的TensorFlow操作提速「20倍」。隨著自動駕駛汽車與機(jī)器人的深入發(fā)展，激光雷達(dá)、深度傳感攝像機(jī)、雷達(dá)等3D傳感器已經(jīng)成為了獲取道路數(shù)據(jù)的必要設(shè)備。而利用這些傳感器的機(jī)器學(xué)習(xí)系統(tǒng)則顯得尤為重要，因為它可以幫助硬件在現(xiàn)實世界中進(jìn)行導(dǎo)航等操作。

近期，包括目標(biāo)檢測、透明目標(biāo)檢測等模型的3D場景理解方面取得了很大進(jìn)展，但是由于3D數(shù)據(jù)可用的工具和資源有限，這個領(lǐng)域仍面臨挑戰(zhàn)。

TensorFlow 3D：TensorFlow與3D深度學(xué)習(xí)合體

為了進(jìn)一步提高對3D場景的建模，簡化研究人員的工作，Google AI發(fā)布了TensorFlow 3D (TF 3D) ，一個高度模塊化、高效的庫，旨在將3D深度學(xué)習(xí)能力引入TensorFlow. TF 3D提供了一系列當(dāng)下常用的操作、損失函數(shù)、數(shù)據(jù)處理工具、模型和度量，使更多的研究團(tuán)隊能夠開發(fā)、培訓(xùn)和部署最先進(jìn)的3D場景理解模型。TF 3D包含用于最先進(jìn)的3D語義分割、3D目標(biāo)檢測和3D實例分割的培訓(xùn)和評估任務(wù)，還支持分布式訓(xùn)練。

另外，TF 3D還支持其他潛在的應(yīng)用，如三維物體形狀預(yù)測、點(diǎn)云配準(zhǔn)和點(diǎn)云增密。此外，它提供了一個統(tǒng)一的數(shù)據(jù)集規(guī)范和訓(xùn)練、評價標(biāo)準(zhǔn)三維場景理解數(shù)據(jù)集的配置。

目前，TF 3D支持Waymo Open、 ScanNet和Rio數(shù)據(jù)集。然而，用戶可以自由地將其他流行的數(shù)據(jù)集，如NuScenes和Kitti，轉(zhuǎn)換成類似的格式，并將其用于已有或自定義的pipeline模型中，還可以利用TF 3D進(jìn)行各種3D深度學(xué)習(xí)研究和應(yīng)用，從快速原型設(shè)計到部署實時推理系統(tǒng)。

左邊顯示的是TF 3D中3D物體檢測模型在Waymo Open Dataset的一幀畫面上的輸出示例。右邊是ScanNet數(shù)據(jù)集上3D實例分割模型的輸出示例。

在這里，我們將介紹在TF 3D中提供的高效且可配置的稀疏卷積骨干，這是在各種3D場景理解任務(wù)中獲得最先進(jìn)結(jié)果的關(guān)鍵。此外，我們將逐一介紹TF 3D目前支持的3個流水線任務(wù): 3D語義分割、3D目標(biāo)檢測分割和3D實例分割。

3D稀疏卷積網(wǎng)絡(luò)

傳感器采集到的3D數(shù)據(jù)通常包含一個場景，該場景包含一組感興趣的物體（如汽車、行人等），其周圍大多是開放空間。所以，3D數(shù)據(jù)本質(zhì)上是稀疏的。在這樣的環(huán)境中，卷積的標(biāo)準(zhǔn)實現(xiàn)將需要大量的計算、消耗大量的內(nèi)存。因此，在TF 3D 中，我們采用了流形稀疏卷積（submanifold sparse convolution）和池操作，這些操作可以更有效地處理3D稀疏數(shù)據(jù)。稀疏卷積模型是大多數(shù)戶外自動駕駛(如Waymo，NuScenes)和室內(nèi)基準(zhǔn)測試(如 ScanNet)中應(yīng)用的sota方法的關(guān)鍵。

谷歌還應(yīng)用了各種CUDA技術(shù)來加快計算速度(如hash、在共享內(nèi)存中分區(qū)/緩存過濾器以及使用位操作)。在Waymo Open數(shù)據(jù)集上的實驗表明，這種實現(xiàn)比預(yù)先設(shè)計好的TensorFlow操作要快「20倍」左右。

圖源：Waymo Open Dataset on GitHub

然后，TF 3D使用3D流形稀疏U-Net架構(gòu)來提取每個voxel的特征。通過讓網(wǎng)絡(luò)提取粗細(xì)特征并將它們組合起來進(jìn)行預(yù)測，U-Net架構(gòu)已被證明是有效的。

U-Net網(wǎng)絡(luò)由編碼器、瓶頸和解碼器三個模塊組成，每個模塊都由許多稀疏卷積塊組成，并可能進(jìn)行池化或非池化操作。

一個3D稀疏體素U-Net架構(gòu)。注意，一個水平的箭頭接收體素特征，并對其應(yīng)用流形稀疏卷積。向下移動的箭頭會執(zhí)行流形稀疏池化。向上移動的箭頭將收集池化的特征，與水平方向箭頭的特征進(jìn)行concat，并對concat后的特征進(jìn)行流形稀疏卷積。

上述稀疏卷積網(wǎng)絡(luò)是TF 3D提供的3D場景理解pipeline模型的backbone。

下面描述的每個模型使用這個骨干網(wǎng)絡(luò)提取稀疏體素特征，然后添加一個或多個額外的預(yù)測頭來推斷感興趣的任務(wù)。

用戶可以通過改變編碼器/解碼器層數(shù)和每層卷積的數(shù)量來配置U-Net網(wǎng)絡(luò)，并通過修改卷積濾波器的尺寸，從而能夠通過不同的網(wǎng)絡(luò)配置來權(quán)衡的速度和精度。

三維語義分割

三維語義分割模型只有一個輸出，用于預(yù)測每一個點(diǎn)的語義分?jǐn)?shù)，將其映射回點(diǎn)，預(yù)測每一個點(diǎn)的語義標(biāo)簽。從ScanNet數(shù)據(jù)集對室內(nèi)場景進(jìn)行3D語義分割。

三維實例分割

在三維實例分割中，除了要預(yù)測語義，更重要的是將同一對象的體素組合在一起。在TF 3D中使用的3D實例分割算法是基于用深度度量學(xué)習(xí)方法進(jìn)行的2D圖像分割工作。這種模型預(yù)測能預(yù)測每個體素的實例嵌入向量以及每個體素的語義評分。

實例嵌入向量將體素映射到一個嵌入空間，其中對應(yīng)于同一對象實例的體素相距很近，而對應(yīng)于不同對象的體素相距很遠(yuǎn)。在這種情況下，輸入是一個點(diǎn)云而不是一個圖像，并且他將使用一個三維稀疏網(wǎng)絡(luò)而不是一個二維圖像網(wǎng)絡(luò)。在推理過程中利用貪心算法選取實例種子，并利用體素嵌入的距離函數(shù)將不同的體素聚合到對應(yīng)的實例上去。

三維目標(biāo)檢測

目標(biāo)檢測模型可以預(yù)測每個體素的大小、中心和旋轉(zhuǎn)矩陣以及對象的語義評分。在推理時，推選機(jī)制將給出的多個候選框處理為少數(shù)幾個精確的3D目標(biāo)框。在訓(xùn)練時使用了預(yù)測與GT間的Huber Loss距離來計算損失。由于利用大小、中心和旋轉(zhuǎn)矩陣估算框邊角是可差分過程，損失可以自然地傳遞到預(yù)測過程的權(quán)重中。研究人員利用動態(tài)框分類損失來對預(yù)測的框進(jìn)行正例和負(fù)例進(jìn)行區(qū)分。

ScanNet數(shù)據(jù)集上的3D物體檢測結(jié)果。

TF 3D只是市場上的3D深度學(xué)習(xí)擴(kuò)展之一。2020年，F(xiàn)acebook推出了PyTorch3D，專注于3D渲染和虛擬現(xiàn)實。另一個是英偉達(dá)的Kaolin，這是一個模塊化的可分辨渲染的應(yīng)用，如高分辨率模擬環(huán)境。

從這個概述來看，TF 3D應(yīng)用程序似乎更專注于機(jī)器人感知和映射，而其他選項則更專注于3D模擬和渲染。為了實現(xiàn)3D渲染，Google推出了TensorFlow Graphics.

參考資料：

https://ai.googleblog.com/2021/02/3d-scene-understanding-with-tensorflow.html

編輯：jq