/KPConv.pytorch

PyTorch reimplementation for "KPConv: Flexible and Deformable Convolution for Point Clouds" https://arxiv.org/abs/1904.08889

Primary LanguageC++

KPConv.pytorch

This repo is implementation for KPConv(https://arxiv.org/abs/1904.08889) in pytorch.

TODO

There are still some works to be done:

  • Deformable KPConv. Currently I have only implemented the rigid KPConv.
    • Regularization loss for the deformable convolution needs to be implemented. I have tried using the deformable convolution layer in part segmention on shapenet without the regularization term, the performance is similar with the rigid convolution counterparts.
  • Speed up. For current implementation, the collate_fn where the neighbor indices and pooling indices are calculated, is too slow. In the tf version, the author implement 2 tensroflow C++ wrapper which is quite efficient. I am planing to write C++ extention using pytorch.
    • But after I implemented the C++ extention, the evaluation time reduces significantly while the model forward and backward pass still cost about 0.8s per iteration.
  • Maybe other datasets.

Installation

  1. Create an environment from the environment.yml file,
conda env create -f environment.yml
  1. Compile the customized Tensorflow operators and C++ extension module following the installation instructions provided by the authors.
  2. Go to pytorch_ops dictionary and run python setup.py install to build and install the C++ extension for batch_find_neighbors function.

Experiments

Due to the time limitation, I have just implemented the experiments on ShapeNet(classification and part segmentation) and ModelNet40.

  • Shape Classification on ModelNet40 or ShapeNet.
python training_ModelNet.py[training_ShapeNetCls.py]
  • Part Segmentation on ShapeNet. (I have only implemented the single class part segmentation.)
python training_ShapeNetPart.py

Acknowledgment

Thank @HuguesTHOMAS for sharing the tensorflow version and valuable explainations.