bufferkdtree¶

Quick Overview¶

The main approach provided by the bufferkdtree package is an efficient many-core implementation (suitable for, e.g., GPUs) for processing huge amounts of nearest neighor queries by means of so-called buffer k-d trees. Such trees depict modifications of standard k-d trees that can be used to make use of the massive parallelism provided by today’s many-core devices (such as GPUs).

Buffer k-d trees aim at scenarios, where you are given both a large reference (e.g., 1,000,000 points) and a huge query set (e.g., 10,000,000 points or more) with an input space of moderate dimensionality (e.g., from d=5 to d=25 dimensions). The general workflow is sketched below; the key idea is to “delay” the processing of nearest neighbor queries until enough work is gathered that can be processed by the many-core device.

Workflow: Initially, all queries are given in the input queue. The computation of nearest neighbors takes place in iterations. In each iteration, the procedure FindLeafBatch removes query indices from both queues and distributes them to the buffers (or removes them if no further processing is needed). In case enough work has been gathered, the procedure ProcessAllBuffers is invoked, which updates the nearest neighbors and reinserts all query indices into reinsert. The process stops as soon as both queues and all buffers are empty. For each query (index), an associated stack is stored, which is used to traverse the overall tree.

Note that each query is traversed in the same manner as for a standard k-d tree traversal (given the same tree depth). However, in contrast to the original traversal, queries are now grouped together before the nearest neighbors are updated in each leaf. This greatly improves the performance on today’s many-core devices, since similar memory regions are processed by neighbored threads.

A detailed description of the techniques used and an experimental evaluation of the implementation using massive astronomical data sets are provided in:

Fabian Gieseke, Justin Heinermann, Cosmin Oancea, and Christian Igel. Buffer k-d Trees: Processing Massive Nearest Neighbor Queries on GPUs. In: Proceedings of the 31st International Conference on Machine Learning (ICML) 32(1), 2014, 172-180. [pdf] [bibtex]

Installation¶

$ sudo pip install bufferkdtree

$ git clone https://github.com/gieseke/bufferkdtree.git

$ cd bufferkdtree
$ python setup.py install --user

$ sudo python setup.py build
$ sudo python setup.py install

$ sudo apt-get install python2.7 python-dev swig build-essential python-numpy python-setuptools

$ sudo zypper install python python-devel swig python-numpy python-setuptools

export C_INCLUDE_PATH=$C_INCLUDE_PATH:/usr/local/cuda/include

$ sudo apt-get install python-virtualenv python-pip

$ mkdir ~/.virtualenvs
$ cd ~/.virtualenvs
$ virtualenv bufferkdtree
$ source bufferkdtree/bin/activate
$ pip install numpy==1.6.1
$ pip install bufferkdtree

Examples¶

The following two examples sketch the use of the different implementations and can both be found in the examples subdirectory of the bufferkdtree package.

import numpy
from bufferkdtree.neighbors import NearestNeighbors

n_neighbors = 10
plat_dev_ids = {0:[0]}
n_jobs = 1
verbose = 0

X = numpy.random.uniform(low=-1, high=1, size=(10000,10))

# (1) apply buffer k-d tree implementation
nbrs_buffer_kd_tree = NearestNeighbors(algorithm="buffer_kd_tree", \
                        tree_depth=9, \
                        plat_dev_ids=plat_dev_ids, \
                        verbose=verbose)    
nbrs_buffer_kd_tree.fit(X)
dists, inds = nbrs_buffer_kd_tree.kneighbors(X, n_neighbors=n_neighbors)
print("\nbuffer_kd_tree output\n" + unicode(dists[0]))

# (2) apply brute-force implementation
nbrs_brute = NearestNeighbors(algorithm="brute", \
                        plat_dev_ids=plat_dev_ids, \
                        verbose=verbose)    
nbrs_brute.fit(X)
dists, inds = nbrs_brute.kneighbors(X, n_neighbors=n_neighbors)
print("\nbrute output\n" + unicode(dists[0]))

# (3) apply k-d tree mplementation
nbrs_kd_tree = NearestNeighbors(algorithm="kd_tree", \
                        n_jobs=n_jobs, \
                        verbose=verbose)    
nbrs_kd_tree.fit(X)
dists, inds = nbrs_kd_tree.kneighbors(X, n_neighbors=n_neighbors)
print("\nkd_tree output\n" + unicode(dists[0]) + "\n")

Nearest Neighbors
=================

This example demonstrates the use of the different
implementations given on a small artifical data set.


buffer_kd_tree output
[ 0.          1.0035212   1.09866345  1.11734533  1.13440645  1.17730558
  1.1844281   1.20736992  1.2085104   1.21593559]

brute output
[ 0.          1.0035212   1.09866357  1.11734521  1.13440645  1.17730546
  1.18442798  1.20736992  1.2085104   1.21593571]

kd_tree output
[ 0.          1.0035212   1.09866357  1.11734521  1.13440645  1.17730546
  1.18442798  1.20736992  1.2085104   1.21593571]

import time
import generate
from bufferkdtree.neighbors import NearestNeighbors

# parameters
plat_dev_ids = {0:[0,1,2,3]}
n_jobs = 8
verbose = 0
n_neighbors=10

def run_algorithm(algorithm="buffer_kd_tree", tree_depth=None, leaf_size=None):

    nbrs = NearestNeighbors(n_neighbors=n_neighbors, \
                            algorithm=algorithm, \
                            tree_depth=tree_depth, \
                            leaf_size=leaf_size, \
                            n_jobs = n_jobs, \
                            plat_dev_ids=plat_dev_ids, \
                            verbose=verbose)

    start_time = time.time()
    nbrs.fit(Xtrain)
    end_time = time.time()
    print("Fitting time: %f" % (end_time-start_time))

    start_time = time.time()
    dists, inds = nbrs.kneighbors(Xtest)
    end_time = time.time()
    print("Testing time: %f" % (end_time-start_time))

print("Parsing data ...")
Xtrain, Ytrain, Xtest = generate.get_data_set(NUM_TRAIN=2000000, NUM_TEST=10000000)
print("-------------------------------- DATA --------------------------------")
print("Number of training patterns:\t %i" % Xtrain.shape[0])
print("Number of test patterns:\t %i" % Xtest.shape[0])
print("Dimensionality of patterns:\t%i" % Xtrain.shape[1])

Nearest Neighbors
=================

This example demonstrates the use of both tree-based
implementations on a large-scale data set.

-------------------------------- DATA --------------------------------
Number of training patterns: 2000000
Number of test patterns:     10000000
Dimensionality of patterns:  10
----------------------------------------------------------------------

print("\n\nRunning the GPU version ...")
run_algorithm(algorithm="buffer_kd_tree", tree_depth=9)

print("\n\nRunning the CPU version ...")
run_algorithm(algorithm="kd_tree", leaf_size=32)

Running the GPU version ...
Fitting time: 1.394939
Testing time: 11.148126

Running the CPU version ...
Fitting time: 0.681938
Testing time: 314.787735

Reference¶

All neighbor implementations can be invoked via the main NearestNeighbors class, which exhibits a similar structure as the corresponding class from the scikit-learn package:

class bufferkdtree.neighbors.NearestNeighbors(n_neighbors=5, algorithm='buffer_kd_tree', tree_depth=None, leaf_size=30, splitting_type='cyclic', n_train_chunks=1, plat_dev_ids={0: [0]}, allowed_train_mem_percent_chunk=0.15, allowed_test_mem_percent=0.55, n_jobs=1, verbose=0, **kwargs)¶

The ‘NearestNeighbors’ provides access to all nearest neighbor implementations. It has simmilar parameters as the corresponding implementation of the scikit-learn package.

The main method is the “buffer_kd_tree”, which can be seen as mix between the “brute” and the “kd_tree” implementations.

Parameters:

n_neighbors : int (default 5) Number of neighbors used algorithm : {“brute”, “kd_tree”, “buffer_kd_tree”}, optional (default=”buffer_kd_tree”) The algorithm that shall be used to compute the nearest neighbors. One of - ‘brute’: brute-force search - ‘kd_tree’: k-d tree based search - ‘buffer_kd_tree’: buffer k-d tree based search (with GPUs) tree_depth : int or None, optional (default=None) Passed to the ‘kd_tree’ and ‘buffer_kd_tree’ implementation. In case ‘tree_depth’ is specified, a tree of such a depth is built (‘tree_depth’ has priority over ‘leaf_size’). leaf_size : int, optional (default=30) Passed to the ‘kd_tree’ and ‘buffer_kd_tree’ implementation. In case ‘leaf_size’ is set, the corresponding tree depth is computed (is ignored in case tree_depth is not None). splitting_type : {‘cyclic’}, optional (default=’cyclic’) Passed to the ‘kd_tree’ and ‘buffer_kd_tree’ implementation. The splitting rule that shall be used to construct the kd tree. Currently, only “cyclic” is supported. n_train_chunks : int, optional (default=1) Passed to the ‘buffer_kd_tree’ implementation. The number of chunks the training patterns shall be processed in; only needed in case the training patterns do not fit on the GPU (in case n_train_chunks is too small, it is increased automatically). plat_dev_ids : dict, optional (default={0:[0]}) Passed to the ‘brute’ and the ‘buffer_kd_tree’ implementation. The platforms and devices that shall be used. E.g., plat_dev_ids={0:[0,1]} makes use of platform 0 and the first two devices. allowed_train_mem_percent_chunk : float, optional (default=0.15) Passed to the ‘buffer_kd_tree’ implementation. The amount of memory (OpenCL) used for the training patterns (in percent). allowed_test_mem_percent : float, optional (default=0.55) Passed to the ‘buffer_kd_tree’ implementation. The amount of memory (OpenCL) used for the test/query patterns (in percent). n_jobs : int, optional (default=1) Passed to the ‘kd_tree’ implementation. The number of threads used for the querying phase. verbose : int, optional (default=0) The verbosity level (0=no output, 1=output)

Notes

The brute-force implementatation is only used for comparison in relatively low-dimensional spaces; the performance is suboptimal for higher dimensional feature spaces (but even superior over other matrix based implementations making use e.g., CUBLAS).

The performance of the GPU implementations depends on the corresponding architecture. An important ingredient is the support of automatic hardware caches.

Only single-precision is supported until now.

Examples

>>> import numpy
>>> from bufferkdtree.neighbors.base import NearestNeighbors
>>> X = numpy.random.uniform(low=-1,high=1,size=(10000,10))
>>> nbrs = NearestNeighbors(n_neighbors=10, algorithm="buffer_kd_tree", tree_depth=9, plat_dev_ids={0:[0]})    
>>> nbrs.fit(X)
>>> dists, inds = nbrs.kneighbors(X)   

fit(X)¶

Fit the model to the given data.

Parameters:	X : array-like, shape (n_samples, n_features) The set of training/reference points, where ‘n_samples’ is the number points and ‘n_features’ the number of features.
Returns:	self : instance of NearestNeighbors The object itself

kneighbors(X=None, n_neighbors=None, return_distance=True)¶

Finds the nearest neighbors for a given set of points.

Parameters:

X : array-like, shape (n_samples, n_features) The set of query points. If not provided, the neighbors of each point in the training data are returned (in this case, the query point itself is not considered its own neighbor. n_neighbors : int or None, optional (default=None) The number of nearest neighbors that shall be returned for each query points. If ‘None’, then the default values provided to the constructor is used. return_distance : bool, optional (default=True) If False, then the distances associated with each query will not be returned. Otherwise, they will be returned.

Returns:

dist : array The array containing the distances idx : array The array containing the indices

License¶

The source code is published under the GNU General Public License Version 2 (GPLv2).

Changes¶

Citations¶

If you wish to cite a paper that describes the techniques and the implementation for buffer k-d trees, please make use of the following work:

Fabian Gieseke, Justin Heinermann, Cosmin Oancea, and Christian Igel. Buffer k-d Trees: Processing Massive Nearest Neighbor Queries on GPUs. In: Proceedings of the 31st International Conference on Machine Learning (ICML) 32(1), 2014, 172-180. [pdf] [bibtex]