harrydobbs/frnn

Fixed Radius Nearest Neighbor Search on GPU

FRNN

A Fixed Radius 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
pip install .

# install FRNN
cd ../../ # back to the {FRNN} directory
# this might take a while since I instantiate all combinations of D and K
pip install -e .
# might need to restart the bash to make importing this package work

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 small point clouds (e.g. < 10,000 points), the bruteforce way (e.g. pytorch3d's KNN) might be faster.

TODO

• support large D (not fully optimized yet)
• support large K (not fully optimized yet)
• try use z-order for the grid cell indices
• speedup and interface for the same query and reference point cloud
• collect all points within radius
• cpp standalone implementation

If you want a new feature, just open an issue or send me an email about it.

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.