[论文解读] Kaolin: A PyTorch Library for Accelerating 3D Deep Learning Research

Krishna Murthy Jatavallabhula, Edward J. Smith+7 more|arXiv (Cornell University)|Nov 12, 2019

Image Processing and 3D Reconstruction参考文献 41被引用 69

一句话总结

Kaolin 是一个基于 PyTorch 的库，通过提供可微分的 3D 模块、数据集工具、可微分渲染器、带预训练架构的模型库以及可视化工具，来加速 3D 深度学习研究。

ABSTRACT

We present Kaolin, a PyTorch library aiming to accelerate 3D deep learning research. Kaolin provides efficient implementations of differentiable 3D modules for use in deep learning systems. With functionality to load and preprocess several popular 3D datasets, and native functions to manipulate meshes, pointclouds, signed distance functions, and voxel grids, Kaolin mitigates the need to write wasteful boilerplate code. Kaolin packages together several differentiable graphics modules including rendering, lighting, shading, and view warping. Kaolin also supports an array of loss functions and evaluation metrics for seamless evaluation and provides visualization functionality to render the 3D results. Importantly, we curate a comprehensive model zoo comprising many state-of-the-art 3D deep learning architectures, to serve as a starting point for future research endeavours. Kaolin is available as open-source software at https://github.com/NVIDIAGameWorks/kaolin/.

研究动机与目标

• 通过在 PyTorch 中提供一个端到端工具包，降低进入 3D 深度学习的门槛。
• 支持多种 3D 表示（网格、点云、体素、SDF、深度图）的跨表示高效转换。
• 提供可微分渲染和包含预训练基线的全面模型库，便于快速实验。
• 提供针对 3D 任务定制的损失函数、评估指标和可视化工具。
• 促进开源社区的发展并制定扩展 3D DL 研究工具的路线图。

提出的方法

• 提供基于 PyTorch 张量的全面 3D 表示类，并具备可微分的跨表示转换。
• 实现一个模块化的可微分渲染器，包含抽象的 DifferentiableRenderer 基类和 CUDA 加速组件。
• 包含扩展 PyTorch Dataset/DataLoader 的数据集加载框架，支持 USD 导出/导入，并支持流行的 3D 数据集（ShapeNet、ModelNet、ScanNet 等）。
• 提供 3D 损失函数和度量库（例如体素的 IoU、Chamfer 距离、Earth Mover’s distance、点到表面损失）。
• 提供带有预训练架构的模型库（例如 Pixel2Mesh、MeshCNN、PointNet/PointNet++、Occupancy Networks、DeepSDF），以加速方法基准测试与开发。
• 通过轻量级工具实现对所有表示的可视化，并通过基于 USD 的导出实现高保真渲染。

实验结果

研究问题

• RQ1 如何通过一个统一的基于 PyTorch 的工具包来简化对多种 3D 表示的加载、转换和训练？
• RQ2 是否可以将可微分渲染以模块化和可扩展的方式集成，以实现带 2D 监督的 3D 学习？
• RQ3 应提供哪些基线和预训练模型以标准化评估并加速 3D DL 的研究？
• RQ4 哪些度量和可视化工具最能支持快速开发和跨 3D 任务的公平比较？

主要发现

• Kaolin 提供对网格、点云、体素网格、符号距离函数和深度图的广泛支持，具备可微分转换。
• 一个模块化的 DifferentiableRenderer 使得可以交换渲染组件（几何、光照、着色、投影、光栅化），并包含用于性能的 CUDA 实现。
• 模型库包含最先进的 3D 架构及其预训练模型，可在分类、分割、重建和可微分渲染等任务中进行快速基准测试。
• Kaolin 将 ShapeNet、PartNet、SHREC、ModelNet、ScanNet、HumanSeg 等数据集与 USD 导出/导入集成，实现数据处理的流线化。
• 本文通过示例和详细生态系统展示了显著的加速和实用性，减少了 3D DL 工作流中的样板代码。
• 可视化和数据导出能力支持端到端的实验，以及在需要时更高保真的渲染。