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

導讀

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

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

TensorFlow 3D：TensorFlow與3D深度學習合體

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

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

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

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

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

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

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

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

圖源：Waymo Open Dataset on GitHub

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

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

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

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

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

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

三維語義分割

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

三維實例分割

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

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

三維目標檢測

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

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

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

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

參考資料：

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

編輯：jq