Flex-Convolution (Million-Scale Point-Cloud Learning Beyond Grid-Worlds)

Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch

Abstract

Traditional convolution layers are specifically designed to exploit the natural data representation of images -- a fixed and regular grid. However, unstructured data like 3D point clouds containing irregular neighborhoods constantly breaks the grid-based data assumption. Therefore applying best-practices and design choices from 2D-image learning methods towards processing point clouds are not readily possible. In this work, we introduce a natural generalization flex-convolution of the conventional convolution layer along with an efficient GPU implementation. We demonstrate competitive performance on rather small benchmark sets using fewer parameters and lower memory consumption and obtain significant improvements on a million-scale real-world dataset. Ours is the first which allows to efficiently process 7 million points concurrently.

The following figure shows the raw network semantic segmentation prediction on a real-world example:

This repository contains the source code of our FlexConv Layer from our 2018 ACCV paper "Flex-Convolution (Million-Scale Point-Cloud Learning Beyond Grid-Worlds)".

Provided novel operations

In the following we summarize the operations described in our paper along with a highly tuned (online, exhaustive) nearest neighbor search layer for 3d point-clouds. All layers follow the tf.layers interface and can be directly used in your TensorFlow model. We further provide unit-tests to verify the correctness of our implementation.

# some point-cloud data
features = np.random.randn(B, Din, N).astype(np.float32)
positions = np.random.randn(B, Dp, N).astype(np.float32)

# To find neighborhoods of K neighbors (not used in our paper, just for your convenience):
neighborhoods = knn_bruteforce(positions, K=8)

# To apply our Flex-Convolution operation to a given set of points with some input-features, position and neighborhood information:
new_features = flex_convolution(features, positions, neighborhoods, out_channels=32)
new_features = flex_convolution_transpose(features, positions, neighborhoods, out_channels=32)

# To apply max-pooling for each neighborhood:
new_features = flex_pooling(features, neighborhoods)

Build Instructions

Our CUDA kernel implementations use CUB primitives. You might install this header-only library by

user@host $ # apt-get install unzip
user@host $ cd /tmp
user@host $ wget https://github.com/NVlabs/cub/archive/v1.8.0.zip
user@host $ unzip v1.8.0.zip -d $HOME/libs
user@host $ export CUB_INC=$HOME/libs/cub-1.8.0/
user@host $ rm /tmp/v1.8.0.zip

We provide GPU-tailored CUDA implementations of our novel FlexConv, FlexPool, FlexDeconv, NearestNeighbor operations in TensorFlow, which require a compilation/linking step. To build our operations just use

user@host $ pip install tensorflow-gpu --user                # optional if not yet installed
user@host $ cd user_ops
user@host $ cmake . -DPYTHON_EXECUTABLE=python2 && make -j   # switch the python version when necessary
user@host $ python test_all.py                               # run all unit-tests to verify the operations
user@host $ cd ..
user@host $ python example.py                                # fully functional toy-example

Deep learning on point-clouds is a complex matter and an active research area. Hence, our internal codebase reflects that complexity and we try our best to provide a usable implementation. To provide a simple training example, we demonstrate training on a very basic 3D-MNIST dataset which deliberately omits fancy parts to give you an idea how to actually train such a model with our operations:

user@host $ python basic_mnist_3d.py --gpu 0

Benchmark

We benchmarked the inference time of entire network on the 2D-3D-S dataset and with a recent test on a NVIDIA V100 GPU, we were able to process ~18 Million Points (the paper stated 7 Million Points on 1080 GTX).

ShapeNet Segmentation

ShapeNet part segmentation results per category and mIoU (%) for different methods and inference speed (on a NVIDIA GeForce GTX 1080 Ti).

2D-3D-S dataset

Class specific average precision (AP) on the 2D-3D-S dataset.

More Resources

• Arxiv Pre-Print
• Video
• Project-Page

Citation

If you use the code in this repository, please cite our paper:

@inproceedings{accv2018/Groh,
  author    = {Fabian Groh and
               Patrick Wieschollek and
               Hendrik P. A. Lensch
               },
  title     = {Flex-Convolution (Million-Scale Point-Cloud Learning Beyond Grid-Worlds)},
  booktitle = {Asian Conference on Computer Vision (ACCV)},
  month     = {Dezember},
  year      = {2018}
}

References

[1] C. Qi, H. Su, K. Mo, L. Guibas, "PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation", Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2017.
[2] C. Qi and L. Yi, H. Su, L. Guibas, "PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space", Advances in Neural Information Processing Systems (NIPS) 2017.
[3] H. Su, V. Jampani, D.Sun, S. Maji, E. Kalogerakis, M.-H. Yang, J. Kautz, "SPLATNet: Sparse Lattice Networks for Point Cloud Processing", Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018.
[4] R. Klokov, V. Lempitsky, "Escape from Cells: Deep Kd-Networks for the Recognition of 3D Point Cloud Models", Proceedings of the IEEE International Conference on Computer Vision (ICCV) 2017.
[5] W. Wang, R. Yu, Q. Huang, U. Neumann, "Sgpn: Similarity group proposal network for 3d point cloud instance segmentation", Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018.
[6] I. Armeni, A. Sax, A.-R. Zamir, S. Savarese, "Joint 2D-3D-Semantic Data for Indoor Scene Understanding", ArXiv e-prints 2017.