Skip to content

Commit

Permalink
switch BruteForce to single precision math
Browse files Browse the repository at this point in the history
  • Loading branch information
piskvorky committed Jun 13, 2015
1 parent cf2b2d9 commit 8b321a2
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions ann_benchmarks.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,26 +176,29 @@ def query(self, v, n):

class BruteForce(BaseANN):
"""kNN search that uses a linear scan = brute force."""
def __init__(self, metric):
def __init__(self, metric, precision=numpy.float32):
if metric not in ('angular', 'euclidean'):
raise NotImplementedError("BruteForce doesn't support metric %s" % metric)
self._metric = metric
self._precision = precision
self.name = 'BruteForce()'

def fit(self, X):
"""Initialize the search index."""
self.lengths = (X ** 2).sum(-1) # precompute (squared) length of each vector
lens = (X ** 2).sum(-1) # precompute (squared) length of each vector
if self._metric == 'angular':
# for cossim, normalize index vectors to unit length
self.index = numpy.ascontiguousarray(X / numpy.sqrt(self.lengths)[..., numpy.newaxis])
X /= numpy.sqrt(lens)[..., numpy.newaxis] # normalize index vectors to unit length
self.index = numpy.ascontiguousarray(X, dtype=self._precision)
elif self._metric == 'euclidean':
self.index = numpy.ascontiguousarray(X)
self.index = numpy.ascontiguousarray(X, dtype=self._precision)
self.lengths = numpy.ascontiguousarray(lens, dtype=self._precision)
else:
assert False, "invalid metric" # shouldn't get past the constructor!

def query(self, v, n):
"""Find indices of `n` most similar vectors from the index to query vector `v`."""
# HACK we ignore query length as that's a constant not affecting the final ordering
v = numpy.ascontiguousarray(v, dtype=self._precision) # use same precision for query as for index
if self._metric == 'angular':
# argmax_a cossim(a, b) = argmax_a dot(a, b) / |a||b| = argmin_a -dot(a, b)
dists = -numpy.dot(self.index, v)
Expand Down

0 comments on commit 8b321a2

Please sign in to comment.