A Fixed Radius Nearest Neighbors Search implemented on CUDA with similar interface as pytorch3d.ops.knn_points.
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.
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
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)
For small point clouds (e.g. < 10,000 points), the bruteforce way (e.g. pytorch3d's KNN) might be faster.
- 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.
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.