A Fixed Nearest Neighbors Search implemented on CUDA

FRNN

A Fixed Nearest Neighbors Search implemented on CUDA with similar interface as pytorch3d.ops.knn_points.

Performance

Algorithm Walkthrough & Experiment Results

FRNN Presentation

Depenency

Tested with cuda 10.2, python 3.8 and pytorch 1.6.0 on ubuntu 18.04.

Should be also fine other versions of cuda/python/pytorch.

Install

git clone --recursive https://github.com/lxxue/FRNN.git
# install a prefix_sum routine first
cd FRNN/external/prefix_sum
python setup.py install

# install FRNN
cd ../../ # back to the {FRNN} directory
python setup.py install

Usage

For fixed nearest neighbors search: doc

  # first time there is no cached grid
  dists, idxs, nn, grid = frnn.frnn_grid_points(
        points1, points2, lengths1, lengths2, K, r, grid=None, return_nn=False, return_sorted=True
  )
  # if points2 and r don't change, we can reuse the grid
  dists, idxs, nn, grid = frnn.frnn_grid_points(
        points1, points2, lengths1, lengths2, K, r, grid=grid, return_nn=False, return_sorted=True
  )

For manually gather nearest neighbors from idxs generated via frnn_grid_points: doc

  nn = frnn.frnn_gather(points2, idxs, lengths2)

Note

For now this function only supports D=3 (point clouds) and K <= 32. Would add more supports for arbitrary D & K in the near future. For point clouds with fewer than 10,000 points, the speedup might not be that much.

Acknowledgement

The code is build on the algorithm introduced by Rama C. Hoetzlein. I use the parallel prefix_sum routines implemented by mattdean1. I also learn (copy & paste) a lot from Pytorch3D's KNN implementations.