/kaolin

A PyTorch Library for Accelerating 3D Deep Learning Research

Primary LanguageC++Apache License 2.0Apache-2.0

WIP License

Kaolin: A PyTorch Library for Accelerating 3D Deep Learning Research

Documentation | Paper

Kaolin is an active work-in-progress, we are working toward a stable beta release. Please Watch our github repository for releases to be notified. To start trying our code today see instructions below.

Kaolin is 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.

Fun fact: The name Kaolin, it’s from Kaolinite, a form of plasticine (clay) that is sometimes used in 3D modeling.

Table of Contents

Functionality

Currently, the (beta) release contains several processing functions for 3D deep learning on meshes, voxels, signed distance functions, and pointclouds. Loading of several popular datasets (eg. ShapeNet, ModelNet, SHREC) are supported out-of-the-box. We also implement several 3D conversion and transformation operations (both within and across the aforementioned representations).

Kaolin supports several 3D tasks such as:

Model Zoo

Kaolin curates a large model zoo containing reference implementations of popular 3D DL architectures. Head over here to check them out.

NOTE: For several of these models, the implementation is due to the original authors. We build a bridge to our library and wherever possible, we introduce optimization. If you use any of the models in the model zoo, or the graphics packages (eg. differentiable renderers such as NMR, SoftRas, DIB-R), please cite the original papers in addition to Kaolin. For convenience, BibTeX citation formats for each of the original papers are included in the documentation for each model provided.

Installation and Usage

NOTE: The API is currently somewhat unstable, as we're making constant changes. (It's a beta release)

Requirements

  • Linux
  • Python == 3.6
  • CUDA >= 10.0.130 (with nvcc installed)
  • Display Driver >= 410.48

Windows support is in the works and is currently considered experimental.

Dependencies

  • numpy >= 1.17, <1.18.4
  • PyTorch >=1.2, <1.5 and Torchvision >=0.4.0, <0.6.0 (see pytorch.org for installation instructions)

Installation

We highly recommend installing Kaolin inside of a virtual environment (such as ones created using conda or virtualenv). In this example, we show how to create a conda virtual environment for installing kaolin.

$ conda create --name kaolin python=3.6
$ conda activate kaolin

Dependencies

Install PyTorch and Torchvision by following instructions from https://pytorch.org/. Numpy will be installed as part of the Pytorch installation. Note that the setup file does not automatically install these dependencies.

Install Kaolin

Now, you can install the library. From the root directory of this repo (i.e., the directory containing this README file), run

For General Use
$ python setup.py build_ext --inplace   # optional, allows importing kaolin from the kaolin root directory
$ python setup.py install
For Development (modifying kaolin code)
$ python setup.py develop

Note, if modifying or adding Cython files, ensure that Cython is installed and set the following environment variable USE_CYTHON=1.

During installation, the packman package manager will download the nv-usd package to ~/packman-repo/ containing the necessary packages for reading and writing Universal Scene Description (USD) files.

Note, if you are using an heterogeneous GPUs setup set the architectures for which you want to compile the cuda code using the environment variable.

Example:

$ export TORCH_CUDA_ARCH_LIST="7.0 7.5"

Verify installation

To verify that kaolin has been installed, fire up your python interpreter, and execute the following commands.

>>> import kaolin as kal
>>> print(kal.__version__)

Building the Documentation

To delve deeper into the library, build the documentation. From the root directory of the repository (i.e., the directory containing this README file), execute the following.

$ cd docs
$ sphinx-build . _build

Running Unittests

To run unittests, from the root directory of the repository (i.e., the directory containing this README file), execute the following commands.

$ . setenv.sh
$ pytest tests/

Main Modules

NOTE: For several of these models, the implementation is due to the original authors. We build a bridge to our library and wherever possible, we introduce optimization. If you use any of the models in the model zoo, or the graphics packages (eg. differentiable renderers such as NMR, SoftRas, DIB-R), please cite the original papers in addition to Kaolin. For convenience, BibTeX citation formats for each of the original papers are included in the documentation for each model provided.

Getting Started

Take a look at some of our examples!! Examples include differentiable renderers, voxel superresolution, etc. Begin here.

Note: We will (very soon) host our docs online. Stay tuned for the link. Until then, please follow instructions from above to build docs.

Contributors

Krishna Murthy Jatavallabhula, Edward Smith, Jean-Francois Lafleche, Clement Fuji Tsang, Artem Rozantsev, Wenzheng Chen, Tommy Xiang, Rev Lebaredian, Gavriel State, Sanja Fidler,

Acknowledgements

Acknowledgements

We would like to thank Amlan Kar for suggesting the need for this library. We also thank Ankur Handa for his advice during the initial and final stages of the project. Many thanks to Joanh Philion, Daiqing Li, Mark Brophy, Jun Gao, and Huan Ling who performed detailed internal reviews, and provided constructive comments. We also thank Gavriel State for all his help during the project.

Most importantly, we thank all 3D DL researchers who have made their code available as open-source. The field could use a lot more of it!

License and Copyright

LICENSE | COPYRIGHT

If you find this library useful, consider citing the following paper:

@article{kaolin2019arxiv,
    author = {J., {Krishna Murthy} and Smith, Edward and Lafleche, Jean-Francois and {Fuji Tsang}, Clement and Rozantsev, Artem and Chen, Wenzheng and Xiang, Tommy and Lebaredian, Rev and Fidler, Sanja},
    title = {Kaolin: A PyTorch Library for Accelerating 3D Deep Learning Research},
    journal = {arXiv:1911.05063},
    year = {2019},
}

Contributors for borrowed sources

Here is a list of all authors on relevant research papers that Kaolin borrows code from. Without the efforts of these folks (and their willingness to release their implementations under permissive open-source licenses), Kaolin would not have been possible.

  • Kornia: Edgar Riba, Dmytro Mishkin, Daniel Ponsa, Ethan Rublee and Gary Bradski. [Paper] [Code]
  • Occupancy Networks: Lars Mescheder, Michael Oechsle, Michael Niemeyer, Sebastian Nowozin, Andreas Geiger. [Paper] [Code]
  • Multi-View Silhouette and Depth Decomposition for High Resolution 3D Object Representation: Edward Smith, Scott Fujimoto, David Meger. [Paper] [Code]
  • Pytorch Chamfer Distance: Christian Diller. [Code]
  • GEOMetrics: Edward Smith, Scott Fujimoto, Adriana Romero, David Meger. [Paper] [Code]
  • DeepSDF: Jeong Joon Park, Peter Florence, Julian Straub, Richard Newcombe, Steven Lovegrove. [Paper] [Code]
  • PointGAN: Fei Xia. [Code]
  • AtlasNet: Thibault Groueix, Matthew Fisher, Vladimir G. Kim, Bryan C. Russell, Mathieu Aubry. [Paper] [Code]
  • PointNet: Charles R. Qi, Hao Su, Kaichun Mo, Leonidas J. Guibas. Also, Fei Xia (reimplementation). [Paper] [Code]
  • MeshCNN: Rana Hanocka, Amir Hertz, Noa Fish, Raja Giryes, Shachar Fleishman, Daniel Cohen-Or. [Paper] [Code]
  • DGCNN: Muhan Zhang, Zhicheng Cui, Marion Neumann, Yixin Chen. [Paper] [Code]
  • Neural 3D Mesh Renderer: Hiroharu Kato, Yoshitaka Ushiku, Tatsuya Harada. Also, Nikos Kolotouros (for reimplementation). [Paper] [Code]
  • SoftRasterizer: Shichen Liu, Tianye Li, Weikai Chen, Hao Li. [Paper] [Code]
  • DIB-R: Wenzheng Chen, Jun Gao, Huan Ling, Edward J. Smith, Jaakko Lehtinen, Alec Jacobson, Sanja Fidler. [Paper] [Code]
  • PointNet++: Charles R. Qi, Li (Eric) Yi, Hao Su, Leonidas J. Guibas. Also, Erik Wijmans (reimplementation). [Paper] [Code]
  • Learning a Probabilistic Latent Space of Object Shapes via 3D Generative-Adversarial Modeling': Jiajun Wu, Chengkai Zhang, Tianfan Xue, William T. Freeman, Joshua B. Tenenbaum. [Paper] [Code]