英伟达重磅开源 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库，旨在加速三维深度学习研究。

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

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

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

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

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

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

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

2.1 3D 数据表示

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

图 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拥有一个灵活的模块化可微渲染器，允许子操作的简单交换，以组成新的变体。

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

表2：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的三维内容创建、三维分割、自动标记、基于三维监督的三维目标预测等等。

图 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