Skip to content

YChuan1115/tf_nearest_neighbour

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

9 Commits
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Tensorflow Nearest Neighbour ops

Memory-efficient tensorflow implementation of nearest neighbour algorithm. C++/Cuda code is from Shapenet2017 Chamfer evaluation code. I include the source files for convenience. These files were provided without author/license information, but I am happy to ammend them with such if anyone can provide them.

I had issues with the compile/python script provided so created this repo with the result of my changes. Most changes came from information in the Compile the op using your system compiler of the adding an op page.

Installation

  1. Clone this repository
cd path/to/parent/dir
git clone https://github.com/jackd/tf_nearest_neighbour.git
  1. Compile the operation
cd src
./compile.sh
cd ..
  1. Add the parent directory to your python path
export PYTHONPATH=path/to/parent/dir:$PYTHONPATH

Consider adding this to your .bashrc. 4. Run the test script:

./scripts//test_nn_distance.py

Your output should look something like

..
----------------------------------------------------------------------
Ran 2 tests in 1.391s

OK

If you installed Tensorflow from source

The compile script uses D_GLIBCXX_USE_CXX11_ABI=0. I'm no expert, but I believe this makes it compatible with the tensorflow package installed via pip. If you build tensorflow from source and are using a gcc version >= 5, you'll need to add --cxxopt="-D_GLIBCXX_USE_CXX11_ABI=0" to the bazel build command as per the instructions on the adding an op page when building tensorflow.

Alternatives

The nearest neighbour is computed by taking the pair-wise distance between each point in each cloud. While this algorithm is trivial to implement, and this implementation is memory-efficient, it does not scale well to large point clouds. For clouds of size N and M, it scales like O(N*M).

Different algorithms exist that scale better with large N and M. For example, to find the one-directional nearest neighbours, KDTrees can be built in O(M log M) time and queries in O(N log M) time. benchmark.py contains a simple wrapping of scipy's cKDTree for tensorflow. These are not compiled/run on the GPU, so this leaves plenty of improvement potential. Note the timings do not include the time taken to build the tree, since this can be done on the CPU during preprocessing.

Requirements

  • tensorflow v1.4 or later
  • cuda/cudnn

About

Efficient tensorflow nearest neighbour op

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • C++ 47.4%
  • Python 33.8%
  • Cuda 16.6%
  • Shell 2.2%