Save off indices where possible
from __future__ import absolute_import
import pynndescent
import numpy as np
+import os
+from sklearn.externals import joblib
+from ann_benchmarks.constants import INDEX_DIR
from ann_benchmarks.algorithms.base import BaseANN
class PyNNDescent(BaseANN):
- def __init__(self, metric, n_neighbors=10, n_trees=8, leaf_size=40, queue_size=2.0):
+ def __init__(self, metric, n_neighbors=10, n_trees=8, leaf_size=20, queue_size=2.0):
self._n_neighbors = int(n_neighbors)
self._n_trees = int(n_trees)
self._leaf_size = int(leaf_size)
self._queue_size = float(queue_size)
- self._pynnd_metric = {'angular': 'cosine', 'euclidean': 'euclidean'}[metric]
+ self._pynnd_metric = {'angular': 'cosine',
+ 'euclidean': 'euclidean',
+ 'hamming': 'hamming',
+ 'jaccard': 'jaccard'}[metric]
+ self._index_name = os.path.join(INDEX_DIR, 'pynndescent_%s_%i_%i_%i' % (self._pynnd_metric,
+ self._n_neighbors,
+ self._n_trees,
+ self._leaf_size))
+
self.name = 'PyNNDescent(n_neighbors=%d,n_trees=%d,leaf_size=%d,queue_size=%.2f)' % \
(self._n_neighbors, self._n_trees, self._leaf_size, self._queue_size)
+ d = os.path.dirname(self._index_name)
+ if not os.path.exists(d):
+ os.makedirs(d)
+
def fit(self, X):
+ if os.path.exists(self._index_name):
+ print('Loading index from file')
+ self._index = joblib.load(self._index_name)
+ else:
self._index = pynndescent.NNDescent(X,
n_neighbors=self._n_neighbors,
n_trees=self._n_trees,
leaf_size=self._leaf_size,
metric=self._pynnd_metric)
+ joblib.dump(self._index, self._index_name)
def query(self, v, n):
ind, dist = self._index.query(np.array([v]), k=n, queue_size=self._queue_size)
