英伟达重磅开源Kaolin：基于PyTorch的3D深度学习加速工具

英伟达最新开源了一个用于三维深度学习研究的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概览

图1：Kaolin是一个PyTorch库，旨在加速三维深度学习研究。

Kaolin提供了：1）加载和预处理常用的三维数据集的功能，2）一个大型模型库，其中包含了用于点云、网格、体素网格、符号距离函数和RGB-D图像的三维任务的常用神经网络结构和损失函数，3）实现了几种现有的可微渲染器，并以模块化的方式支持多个着色器，4）最常用的3D度量指标，便于评估研究结果，5）可视化3D结果的功能。Kaolin的函数高度优化，在现有的3D 深度学习研究代码上有显著的加速效果。

表1：Kaolin是第一个全面的3D深度学习库。通过对各种表示、数据集和模型的广泛支持，它对现有的三维库TensorFlow Graphics、Kornia和GVNN等进行了补充。

Kaolin旨在为构建三维深度学习系统和操作三维数据提供高效、简易的工具。借助Kaolin提供的全面的模板代码，3D深度学习的研究人员可以专注于开发其应用程序的新方面。下面将简要介绍这个3D深度学习库的各个主要功能。

图2：用Kaolin训练3D深度学习模型很简单。上图演示了用5行代码训练和测试PointNet++分类器分类汽车和飞机的代码。

2.1 3D 数据表示

图3：Kaolin提供了有效的PyTorch操作，对不同的三维表示进行转换。网格、点云和体素网格仍然是最常用的三维表示，但Kaolin还支持符号距离函数（SDF）、正交深度图（ODM）和RGB-D图像。

由于不同的三维数据类型具有不同属性，三维深度学习中数据表示的选择会对结果产生很大影响。为了确保选择数据表示的高度灵活性，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（可微渲染器），包含渲染管道中每个组件的抽象方法（几何变换、照明、着色、光栅化和投影）。通过这个抽象类，可以简单直观地组合部件、交换模块和开发新技术。

图4：模块化可微渲染器：Kaolin拥有一个灵活的模块化可微渲染器，允许子操作的简单交换，以组成新的变体。

表2：Kaolin加速现有开源代码。

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得到的结果。

图5：Kaolin的应用场景：基于二维监督的三维目标预测、基于GAN的三维内容创建、三维分割、自动标记、基于三维监督的三维目标预测等等。

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