Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/lxxue/frnn
Fixed Radius Nearest Neighbor Search on GPU
https://github.com/lxxue/frnn
cuda nearest-neighbor-search pytorch
Last synced: 1 day ago
JSON representation
Fixed Radius Nearest Neighbor Search on GPU
- Host: GitHub
- URL: https://github.com/lxxue/frnn
- Owner: lxxue
- Created: 2020-07-10T18:17:34.000Z (over 4 years ago)
- Default Branch: master
- Last Pushed: 2024-03-06T11:00:13.000Z (12 months ago)
- Last Synced: 2025-02-08T23:11:18.263Z (10 days ago)
- Topics: cuda, nearest-neighbor-search, pytorch
- Language: Cuda
- Homepage:
- Size: 77.6 MB
- Stars: 184
- Watchers: 5
- Forks: 24
- Open Issues: 5
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# FRNN
A Fixed Radius Nearest Neighbors Search implemented on CUDA with similar interface as [pytorch3d.ops.knn_points](https://pytorch3d.readthedocs.io/en/latest/modules/ops.html#pytorch3d.ops.knn_points).
## Performance
![Performance](./images/teaser.png)
## Algorithm Walkthrough & Experiment Results
[FRNN Presentation](./frnn_pre.pdf)
## 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](https://github.com/lxxue/FRNN/blob/59a4c8fdc786c64afd991919af39f1a65d4ec2ff/frnn/frnn.py#L154-L224)```
# 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](https://github.com/lxxue/FRNN/blob/59a4c8fdc786c64afd991919af39f1a65d4ec2ff/frnn/frnn.py#L268-L292)```
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
- [x] support large D (not fully optimized yet)
- [x] 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 implementationIf you want a new feature, just open an issue or send me an email about it.
## Acknowledgement
The code is build on the [algorithm](https://on-demand.gputechconf.com/gtc/2014/presentations/S4117-fast-fixed-radius-nearest-neighbor-gpu.pdf) introduced by Rama C. Hoetzlein. I use the [parallel prefix_sum](https://github.com/lxxue/prefix_sum) routines implemented by [mattdean1](https://github.com/mattdean1/cuda). I also learn (copy & paste) a lot from [Pytorch3D's KNN](https://github.com/facebookresearch/pytorch3d/blob/master/pytorch3d/csrc/knn/knn.cu) implementations.