英偉達最新開源了一個用於三維深度學習研究的PyTorch庫——Kaolin。Kaolin提供了可用於三維深度學習系統的可微三維模塊。具有加載和預處理常見的三維數據集的功能，同時實現了處理網格、點雲、符號距離函數和體素網格的函數，從而減少了編寫樣板代碼的負擔。Kaolin封裝了常用的可微圖形模塊，包括渲染、照明、陰影和視圖扭曲。Kaolin還支持一系列損失函數和評價指標，以實現對模型的無縫評估，並提供可視化功能以呈現三維結果。重要的是，Kaolin提供了一個全面的Model zoo，收納了許多最先進的三維深度學習架構，可以作爲未來研究工作的起點。本文是AI前線第101篇論文導讀，我們將對英偉達的最新論文《Kaolin: A PyTorch Library for Accelerating 3D Deep Learning Research》進行詳細解讀。
1 介紹
三維深度學習與某些複雜任務高度相關，如機器人、自動駕駛、增強和虛擬現實，因此受到了廣泛的關注和認可。深度學習和不斷髮展的基礎計算設施使得人們能夠分析高度複雜的三維數據。儘管人們對這一領域的研究興趣不斷增強，但由於缺乏標準化的工具，在這一領域開展研究面臨着陡峭的學習曲線。目前還沒有一個能夠實現簡單加載3D數據集、將3D數據轉換成不同數據表示、與現代機器學習框架融合，並訓練和評估深度學習結構的系統。這對於新接觸3D深度學習領域的研究人員很不友好，他們必須從各種代碼庫中編譯互不匹配的代碼片段才能實現一些基礎任務。這也導致state-of-the-art的方法之間缺乏有效的對比。
爲了助力3D深度學習的入門和深入研究，英偉達的研究人員發佈了Kaolin，一個基於PyTorch的三維深度學習庫。Kaolin提供了快速構建3D深度學習應用程序所需的所有核心模塊，從加載和預處理數據、將其轉換爲常用的三維表示（網格、體素、符號距離函數、點雲等），到基於這些表示實現深度學習任務，以及計算評價標準和三維數據可視化，Kaolin使3D深度學習的整個開發週期變得直觀且簡單。此外，Kaolin的Model zoo中實現了一系列常用的3D任務方法以及預訓練的模型。基於Model zoo可以輕鬆實現新方法，因此它可以作爲未來3D深度學習研究的大本營。最後，隨着三維深度學習應用中幾何結構顯式建模和其他物理過程（照明、陰影、投影等）的可微渲染的出現，Kaolin具有一個通用的、模塊化的可微渲染器，方便擴展到其他常見的可微渲染方法，便於將來的研究和開發。
2 Kaolin概覽
Kaolin提供了：1）加載和預處理常用的三維數據集的功能，2）一個大型模型庫，其中包含了用於點雲、網格、體素網格、符號距離函數和RGB-D圖像的三維任務的常用神經網絡結構和損失函數，3）實現了幾種現有的可微渲染器，並以模塊化的方式支持多個着色器，4）最常用的3D度量指標，便於評估研究結果，5）可視化3D結果的功能。Kaolin的函數高度優化，在現有的3D 深度學習研究代碼上有顯著的加速效果。
Kaolin旨在爲構建三維深度學習系統和操作三維數據提供高效、簡易的工具。藉助Kaolin提供的全面的模板代碼，3D深度學習的研究人員可以專注於開發其應用程序的新方面。下面將簡要介紹這個3D深度學習庫的各個主要功能。
2.1 3D 數據表示
由於不同的三維數據類型具有不同屬性，三維深度學習中數據表示的選擇會對結果產生很大影響。爲了確保選擇數據表示的高度靈活性，Kaolin支持所有流行的3D表示：
- 多邊形網格
- 點雲
- 體素網格
- 符號距離函數與水平集
- 深度圖像（2.5D）
每個表示類型都存儲爲一個PyTorch張量集合，集成爲一個獨立的類。這允許在通用函數上進行運算符重載，以實現數據增廣和修改。每個類內都提供了不同表示之間的有效（可微）轉換。例如，基於重參數化（reparameterization），Kaolin提供了可微的曲面採樣機制，使多邊形網格表示能夠轉換爲點雲表示。每種數據表示也有不同的網絡體系結構支持，例如支持網格的圖卷積網絡和MeshCNN、支持體素表示的3D卷積、支持點雲數據的PointNet和PointNet++。下面的一段示例代碼演示瞭如何簡單地將網格模型加載到Kaolin中，然後轉換爲點雲數據，並在兩種表示中進行渲染：
2.2 數據集
Kaolin爲許多常用的3D數據集提供了完全支持，減少了文件處理、解析和擴展到單個函數調用中所需要的工作量。通過對PyTorch的Dataset和DataLoader類的擴展，提供了對所有數據的訪問。這使得加載3D數據與加載MNIST一樣簡單直觀，並且直接爲用戶提供了PyTorch所支持的批數據加載。所有數據都可導入和導出爲通用場景描述（USD）格式，USD提供了一種通用語言，用於定義、打包、組裝和編輯圖形應用程序中的三維數據。
當前支持的數據集包括ShapeNet、PartNet、SHREC、ModelNet、ScanNet、HumanSeg以及許多常見和自定義的數據集。通過ShapeNet，Kaolin提供了一個龐大的CAD模型庫，包括數十個類中的上萬個對象。通過ScanNet，Kaolin提供了1500多個RGD-B視頻，包括250多萬個獨特的深度圖，其中包括對攝像機姿態、曲面重建和語義分割的完整註釋。這些大型的三維數據集以及更多的信息都可以通過調用單個函數輕鬆訪問。例如，訪問ModelNet，將其提供給PyTorch數據加載器，並加載一個batch的體素模型，只需要如下幾行代碼：
2.3 三維幾何函數
Kaolin的核心是一系列有效的三維幾何函數，能夠對三維內容進行操作。例如，剛體變換用參數化（歐拉角、李羣和四元數）實現。還實現了可微的圖像變形層，如GVNN（Neural network library for geometric vision）中定義的透視變形層。幾何子模塊能夠進行三維剛體、仿射和投影變換，以及三維——二維投影和二維——三維反投影。目前支持正投影和透視（針孔）投影。
2.4 模塊化可微渲染器
可微渲染可以使研究人員用常見的二維監督來執行三維任務。然而，開發可微渲染工具並非易事，所涉及的操作計算繁重且複雜。Kaolin提供了一個靈活的模塊化可微渲染器來解決這些問題，讓研究人員能夠簡單使用常見的微分渲染方法。Kaolin定義了一個抽象基礎類：DifferentiableRenderer（可微渲染器），包含渲染管道中每個組件的抽象方法（幾何變換、照明、着色、光柵化和投影）。通過這個抽象類，可以簡單直觀地組合部件、交換模塊和開發新技術。
Kaolin支持多種照明（環境光、定向光、高光）、着色（Lambertian、Phong、Cosine）、投影（透視、正交、扭曲）和光柵化模式。DifferentiableRenderer類的架構如圖4所示。該類是用CUDA實現的，以獲得最佳性能（表2）。爲了減少了開發工作量，多個公開可用的可微渲染器可用作DifferentiableRenderer類的具體實例。下面的代碼中，DIB Renderer被實例化，並使用Kaolin將網格渲染到圖像：
2.5 損失函數和度量指標
3D深度學習應用的一個共同挑戰在於定義和實現用於評估性能和監督神經網絡的工具。例如，對比網格或點雲等曲面表示可能需要對比數千個點或三角形的位置，而CUDA函數是必需的。因此，Kaolin爲每個3D表示提供了一系列常用的3D度量指標。其中包括：用於體素的十字交叉距離，用於點雲數據的倒角距離和（二次近似）推土機距離，用於網格的點到面損失，以及許多其他網格度量，如拉普拉斯，平滑度和邊緣長度正則化。
2.6 Model zoo
新接觸3D深度學習領域的研究人員在選擇3D表示、模型結構、損失函數等方面面臨一系列問題。Kaolin提供了豐富的基線集合，以及用於各種3D任務（分類、分割、圖像三維重建、超分辨率和可微渲染）的最新結構來幫助研究人員。除了源代碼之外，Kaolin的開發人員根據流行的基準發佈了這些任務的預訓練模型，可以作爲未來研究的基線。開發人員希望這將有助於促進該領域的評價方法和評價指標的標準化。
目前在這個模型庫中實現的方法包括Pixel2Mesh、GEOMetrics和AtlasNet，用於從單張圖像重建網格對象；NM3DR、Soft-Rasterizer 和Dib-Renderer，用於同樣的任務，僅需要二維監督；MeshCNN，用於網格通用學習；PointNet和PointNet++用於點雲上的通用學習；3D-GAN、3D-IWGAN和3DR2N2用於體素分佈的學習，以及用於學習水平集和SDF的Occupancy Networks和DeepSDF等等。圖5給出了可通過Kaolin的Model zoo得到的結果。
2.7 可視化
數據可視化在任何計算機視覺任務中都是一個重要方面。然而，3D數據的可視化並不簡單。雖然目前有Python包用於可視化某些數據類型，例如體素和點雲，但沒有支持所有流行的3D表示的可視化工具包。Kaolin的一個關鍵特性是對其所有表示類型都有可視化支持。這是通過輕量級可視化庫實現Trimesh，和用於運行時間可視化的pptk實現的。由於所有數據都可以導出爲USD，因此在更密集的圖形應用程序中，三維結果也更容易可視化，並且具有相當高的保真度（如圖5所示）。對於無頭應用程序，例如在沒有附加顯示的服務器上運行時，開發人員提供了壓縮程序將圖像和動畫渲染到磁盤上，以便以後進行可視化。
3 產品路線圖
Kaolin是加速3D深度學習研究的一個重要步驟，但開發人員並沒有就此停止腳步。他們打算圍繞Kaolin建立一個強大的開源社區，並歡迎其他3D深度學習研究人員進行貢獻。本節主要介紹Kaolin作爲開源軟件的總體路線圖。
Model zoo：研究人員希望能不斷提升Model zoo，特別是考慮到Kaolin提供的各項功能可以減少實現新方法所需的時間（大多數方法可以在一兩天內實現）。
可微渲染：研究人員計劃擴展支持新的可微渲染工具，並加入用於其他任務的功能，如域隨機化、材料恢復等。
LiDAR數據集：研究人員計劃增加幾個大規模的語義和實例分割數據集。例如，支持S3DIS和nuScenes是未來版本的高優先級任務。
三維目標檢測：目前，Kaolin的Model zoo中沒有用於三維目標檢測的模型。這是未來版本的重點領域。
自動混合精度：爲了使三維神經網絡結構更加緊湊和快速，開發人員正在研究自動混合精度（AMP）在常見三維結構（PointNet、MeshCNN、體素U-Net等）中的適用性。Nvidia Apex支持流行的二維深度學習架構的大多數AMP模式，而Kaolin的開發人員想探索如何將這種支持擴展到三維。
次要燈光效果：Kaolin的可微分渲染類目前僅支持主要燈光效果，這限制了應用程序理解更復雜的場景信息（如陰影）的能力。未來的版本計劃增加對路徑跟蹤和光線跟蹤的支持，以便支持次要效果。
查看論文原文：
Kaolin: A PyTorch Library for Accelerating 3D Deep Learning Research