Info
5 |ANN-Benchmarks is a benchmarking environment for approximate nearest neighbor algorithms search. This website contains the current benchmarking results. Please visit http://github.com/erikbern/ann-benchmarks/ to get an overview over evaluated data sets and algorithms. Make a pull request on Github to add your own code or improvements to the 6 | benchmarking system. 7 |
8 |
9 |
37 | {% endfor %} 38 | {% endfor %} 39 |
57 | {% endfor %} 58 | {% endif %} 59 | {% endfor %} 60 | {% endblock %} 61 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/onng_ngt.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import sys 3 | import os 4 | import ngtpy 5 | import numpy as np 6 | import subprocess 7 | import time 8 | from ann_benchmarks.algorithms.base import BaseANN 9 | from ann_benchmarks.constants import INDEX_DIR 10 | 11 | 12 | class ONNG(BaseANN): 13 | def __init__(self, metric, object_type, epsilon, param): 14 | metrics = {'euclidean': '2', 'angular': 'C'} 15 | self._edge_size = int(param['edge']) 16 | self._outdegree = int(param['outdegree']) 17 | self._indegree = int(param['indegree']) 18 | self._metric = metrics[metric] 19 | self._object_type = object_type 20 | self._edge_size_for_search = int(param['search_edge']) if 'search_edge' in param.keys() else -2 21 | self._tree_disabled = (param['tree'] == False) if 'tree' in param.keys() else False 22 | self._build_time_limit = 4 23 | self._epsilon = epsilon 24 | print('ONNG: edge_size=' + str(self._edge_size)) 25 | print('ONNG: outdegree=' + str(self._outdegree)) 26 | print('ONNG: indegree=' + str(self._indegree)) 27 | print('ONNG: edge_size_for_search=' + str(self._edge_size_for_search)) 28 | print('ONNG: epsilon=' + str(self._epsilon)) 29 | print('ONNG: metric=' + metric) 30 | print('ONNG: object_type=' + object_type) 31 | 32 | def fit(self, X): 33 | print('ONNG: start indexing...') 34 | dim = len(X[0]) 35 | print('ONNG: # of data=' + str(len(X))) 36 | print('ONNG: dimensionality=' + str(dim)) 37 | index_dir = 'indexes' 38 | if not os.path.exists(index_dir): 39 | os.makedirs(index_dir) 40 | index = os.path.join( 41 | index_dir, 42 | 'ONNG-{}-{}-{}'.format(self._edge_size, self._outdegree, 43 | self._indegree)) 44 | anngIndex = os.path.join(index_dir, 'ANNG-' + str(self._edge_size)) 45 | print('ONNG: index=' + index) 46 | if (not os.path.exists(index)) and (not os.path.exists(anngIndex)): 47 | print('ONNG: create ANNG') 48 | t = time.time() 49 | args = ['ngt', 'create', '-it', '-p8', '-b500', '-ga', '-of', 50 | '-D' + self._metric, '-d' + str(dim), 51 | '-E' + str(self._edge_size), '-S0', 52 | '-e' + str(self._epsilon), '-P0', '-B30', 53 | '-T' + str(self._build_time_limit), anngIndex] 54 | subprocess.call(args) 55 | idx = ngtpy.Index(path=anngIndex) 56 | idx.batch_insert(X, num_threads=24, debug=False) 57 | idx.save() 58 | idx.close() 59 | print('ONNG: ANNG construction time(sec)=' + str(time.time() - t)) 60 | if not os.path.exists(index): 61 | print('ONNG: degree adjustment') 62 | t = time.time() 63 | args = ['ngt', 'reconstruct-graph', '-mS', 64 | '-o ' + str(self._outdegree), 65 | '-i ' + str(self._indegree), anngIndex, index] 66 | subprocess.call(args) 67 | print('ONNG: degree adjustment time(sec)=' + str(time.time() - t)) 68 | if os.path.exists(index): 69 | print('ONNG: index already exists! ' + str(index)) 70 | t = time.time() 71 | print(self._tree_disabled) 72 | #self.index = ngtpy.Index(index, read_only=True, tree_disabled=self._tree_disabled) 73 | self.index = ngtpy.Index(index, read_only=True) 74 | self.indexName = index 75 | print('ONNG: open time(sec)=' + str(time.time() - t)) 76 | else: 77 | print('ONNG: something wrong.') 78 | print('ONNG: end of fit') 79 | 80 | def set_query_arguments(self, epsilon): 81 | print("ONNG: epsilon=" + str(epsilon)) 82 | self._epsilon = epsilon - 1.0 83 | self.name = 'ONNG-NGT(%s, %s, %s, %s, %1.3f)' % ( 84 | self._edge_size, self._outdegree, 85 | self._indegree, self._edge_size_for_search, 86 | self._epsilon + 1.0) 87 | 88 | def query(self, v, n): 89 | results = self.index.search( 90 | v, n, self._epsilon, self._edge_size_for_search, 91 | with_distance=False) 92 | return results 93 | 94 | def freeIndex(self): 95 | print('ONNG: free') 96 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/elasticsearch.py: -------------------------------------------------------------------------------- 1 | """ 2 | ann-benchmarks interfaces for Elasticsearch. 3 | Note that this requires X-Pack, which is not included in the OSS version of Elasticsearch. 4 | """ 5 | import logging 6 | from time import sleep 7 | from urllib.error import URLError 8 | from urllib.request import Request, urlopen 9 | 10 | from elasticsearch import Elasticsearch 11 | from elasticsearch.helpers import bulk 12 | 13 | from ann_benchmarks.algorithms.base import BaseANN 14 | 15 | # Configure the elasticsearch logger. 16 | # By default, it writes an INFO statement for every request. 17 | logging.getLogger("elasticsearch").setLevel(logging.WARN) 18 | 19 | # Uncomment these lines if you want to see timing for every HTTP request and its duration. 20 | # logging.basicConfig(level=logging.INFO) 21 | # logging.getLogger("elasticsearch").setLevel(logging.INFO) 22 | 23 | def es_wait(): 24 | print("Waiting for elasticsearch health endpoint...") 25 | req = Request("http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s") 26 | for i in range(30): 27 | try: 28 | res = urlopen(req) 29 | if res.getcode() == 200: 30 | print("Elasticsearch is ready") 31 | return 32 | except URLError: 33 | pass 34 | sleep(1) 35 | raise RuntimeError("Failed to connect to local elasticsearch") 36 | 37 | 38 | class ElasticsearchScriptScoreQuery(BaseANN): 39 | """ 40 | KNN using the Elasticsearch dense_vector datatype and script score functions. 41 | - Dense vector field type: https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html 42 | - Dense vector queries: https://www.elastic.co/guide/en/elasticsearch/reference/master/query-dsl-script-score-query.html 43 | """ 44 | 45 | def __init__(self, metric: str, dimension: int): 46 | self.name = f"elasticsearch-script-score-query_metric={metric}_dimension={dimension}" 47 | self.metric = metric 48 | self.dimension = dimension 49 | self.index = f"es-ssq-{metric}-{dimension}" 50 | self.es = Elasticsearch(["http://localhost:9200"]) 51 | self.batch_res = [] 52 | if self.metric == "euclidean": 53 | self.script = "1 / (1 + l2norm(params.query_vec, \"vec\"))" 54 | elif self.metric == "angular": 55 | self.script = "1.0 + cosineSimilarity(params.query_vec, \"vec\")" 56 | else: 57 | raise NotImplementedError(f"Not implemented for metric {self.metric}") 58 | es_wait() 59 | 60 | def fit(self, X): 61 | body = dict(settings=dict(number_of_shards=1, number_of_replicas=0)) 62 | mapping = dict( 63 | properties=dict( 64 | id=dict(type="keyword", store=True), 65 | vec=dict(type="dense_vector", dims=self.dimension) 66 | ) 67 | ) 68 | self.es.indices.create(self.index, body=body) 69 | self.es.indices.put_mapping(mapping, self.index) 70 | 71 | def gen(): 72 | for i, vec in enumerate(X): 73 | yield { "_op_type": "index", "_index": self.index, "vec": vec.tolist(), 'id': str(i + 1) } 74 | 75 | (_, errors) = bulk(self.es, gen(), chunk_size=500, max_retries=9) 76 | assert len(errors) == 0, errors 77 | 78 | self.es.indices.refresh(self.index) 79 | self.es.indices.forcemerge(self.index, max_num_segments=1) 80 | 81 | def query(self, q, n): 82 | body = dict( 83 | query=dict( 84 | script_score=dict( 85 | query=dict(match_all=dict()), 86 | script=dict( 87 | source=self.script, 88 | params=dict(query_vec=q.tolist()) 89 | ) 90 | ) 91 | ) 92 | ) 93 | res = self.es.search(index=self.index, body=body, size=n, _source=False, docvalue_fields=['id'], 94 | stored_fields="_none_", filter_path=["hits.hits.fields.id"]) 95 | return [int(h['fields']['id'][0]) - 1 for h in res['hits']['hits']] 96 | 97 | def batch_query(self, X, n): 98 | self.batch_res = [self.query(q, n) for q in X] 99 | 100 | def get_batch_results(self): 101 | return self.batch_res 102 | 103 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/pynndescent.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import pynndescent 3 | from ann_benchmarks.algorithms.base import BaseANN 4 | import numpy as np 5 | import scipy.sparse 6 | 7 | 8 | class PyNNDescent(BaseANN): 9 | def __init__(self, metric, index_param_dict, n_search_trees=1): 10 | if "n_neighbors" in index_param_dict: 11 | self._n_neighbors = int(index_param_dict["n_neighbors"]) 12 | else: 13 | self._n_neighbors = 30 14 | 15 | if "pruning_degree_multiplier" in index_param_dict: 16 | self._pruning_degree_multiplier = float( 17 | index_param_dict["pruning_degree_multiplier"] 18 | ) 19 | else: 20 | self._pruning_degree_multiplier = 1.5 21 | 22 | if "diversify_prob" in index_param_dict: 23 | self._diversify_prob = float(index_param_dict["diversify_prob"]) 24 | else: 25 | self._diversify_prob = 1.0 26 | 27 | if "leaf_size" in index_param_dict: 28 | self._leaf_size = int(index_param_dict["leaf_size"]) 29 | else: 30 | leaf_size = 32 31 | 32 | self._n_search_trees = int(n_search_trees) 33 | 34 | self._pynnd_metric = { 35 | "angular": "dot", 36 | # 'angular': 'cosine', 37 | "euclidean": "euclidean", 38 | "hamming": "hamming", 39 | "jaccard": "jaccard", 40 | }[metric] 41 | 42 | def _sparse_convert_for_fit(self, X): 43 | lil_data = [] 44 | self._n_cols = 1 45 | self._n_rows = len(X) 46 | for i in range(self._n_rows): 47 | lil_data.append([1] * len(X[i])) 48 | if max(X[i]) + 1 > self._n_cols: 49 | self._n_cols = max(X[i]) + 1 50 | 51 | result = scipy.sparse.lil_matrix( 52 | (self._n_rows, self._n_cols), dtype=np.int 53 | ) 54 | result.rows[:] = list(X) 55 | result.data[:] = lil_data 56 | return result.tocsr() 57 | 58 | def _sparse_convert_for_query(self, v): 59 | result = scipy.sparse.csr_matrix((1, self._n_cols), dtype=np.int) 60 | result.indptr = np.array([0, len(v)]) 61 | result.indices = np.array(v).astype(np.int32) 62 | result.data = np.ones(len(v), dtype=np.int) 63 | return result 64 | 65 | def fit(self, X): 66 | if self._pynnd_metric == "jaccard": 67 | # Convert to sparse matrix format 68 | X = self._sparse_convert_for_fit(X) 69 | 70 | self._index = pynndescent.NNDescent( 71 | X, 72 | n_neighbors=self._n_neighbors, 73 | metric=self._pynnd_metric, 74 | low_memory=True, 75 | leaf_size=self._leaf_size, 76 | pruning_degree_multiplier=self._pruning_degree_multiplier, 77 | diversify_prob=self._diversify_prob, 78 | n_search_trees=self._n_search_trees, 79 | compressed=True, 80 | verbose=True, 81 | ) 82 | if hasattr(self._index, "prepare"): 83 | self._index.prepare() 84 | else: 85 | self._index._init_search_graph() 86 | if self._index._is_sparse: 87 | if hasattr(self._index, "_init_sparse_search_function"): 88 | self._index._init_sparse_search_function() 89 | else: 90 | if hasattr(self._index, "_init_search_function"): 91 | self._index._init_search_function() 92 | 93 | def set_query_arguments(self, epsilon=0.1): 94 | self._epsilon = float(epsilon) 95 | 96 | def query(self, v, n): 97 | if self._pynnd_metric == "jaccard": 98 | # convert index array to sparse matrix format and query 99 | v = self._sparse_convert_for_query(v) 100 | ind, dist = self._index.query(v, k=n, epsilon=self._epsilon) 101 | else: 102 | ind, dist = self._index.query( 103 | v.reshape(1, -1).astype("float32"), k=n, epsilon=self._epsilon 104 | ) 105 | return ind[0] 106 | 107 | def __str__(self): 108 | str_template = "PyNNDescent(n_neighbors=%d, pruning_mult=%.2f, diversify_prob=%.3f, epsilon=%.3f, leaf_size=%02d)" 109 | return str_template % ( 110 | self._n_neighbors, 111 | self._pruning_degree_multiplier, 112 | self._diversify_prob, 113 | self._epsilon, 114 | self._leaf_size, 115 | ) 116 | -------------------------------------------------------------------------------- /ann_benchmarks/plotting/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | import itertools 4 | import numpy 5 | from ann_benchmarks.plotting.metrics import all_metrics as metrics 6 | 7 | 8 | def get_or_create_metrics(run): 9 | if 'metrics' not in run: 10 | run.create_group('metrics') 11 | return run['metrics'] 12 | 13 | 14 | def create_pointset(data, xn, yn): 15 | xm, ym = (metrics[xn], metrics[yn]) 16 | rev_y = -1 if ym["worst"] < 0 else 1 17 | rev_x = -1 if xm["worst"] < 0 else 1 18 | data.sort(key=lambda t: (rev_y * t[-1], rev_x * t[-2])) 19 | 20 | axs, ays, als = [], [], [] 21 | # Generate Pareto frontier 22 | xs, ys, ls = [], [], [] 23 | last_x = xm["worst"] 24 | comparator = ((lambda xv, lx: xv > lx) 25 | if last_x < 0 else (lambda xv, lx: xv < lx)) 26 | for algo, algo_name, xv, yv in data: 27 | if not xv or not yv: 28 | continue 29 | axs.append(xv) 30 | ays.append(yv) 31 | als.append(algo_name) 32 | if comparator(xv, last_x): 33 | last_x = xv 34 | xs.append(xv) 35 | ys.append(yv) 36 | ls.append(algo_name) 37 | return xs, ys, ls, axs, ays, als 38 | 39 | 40 | def compute_metrics(true_nn_distances, res, metric_1, metric_2, 41 | recompute=False): 42 | all_results = {} 43 | for i, (properties, run) in enumerate(res): 44 | algo = properties['algo'] 45 | algo_name = properties['name'] 46 | # cache distances to avoid access to hdf5 file 47 | run_distances = numpy.array(run['distances']) 48 | if recompute and 'metrics' in run: 49 | del run['metrics'] 50 | metrics_cache = get_or_create_metrics(run) 51 | 52 | metric_1_value = metrics[metric_1]['function']( 53 | true_nn_distances, 54 | run_distances, metrics_cache, properties) 55 | metric_2_value = metrics[metric_2]['function']( 56 | true_nn_distances, 57 | run_distances, metrics_cache, properties) 58 | 59 | print('%3d: %80s %12.3f %12.3f' % 60 | (i, algo_name, metric_1_value, metric_2_value)) 61 | 62 | all_results.setdefault(algo, []).append( 63 | (algo, algo_name, metric_1_value, metric_2_value)) 64 | 65 | return all_results 66 | 67 | 68 | def compute_all_metrics(true_nn_distances, run, properties, recompute=False): 69 | algo = properties["algo"] 70 | algo_name = properties["name"] 71 | print('--') 72 | print(algo_name) 73 | results = {} 74 | # cache distances to avoid access to hdf5 file 75 | run_distances = numpy.array(run["distances"]) 76 | if recompute and 'metrics' in run: 77 | del run['metrics'] 78 | metrics_cache = get_or_create_metrics(run) 79 | 80 | for name, metric in metrics.items(): 81 | v = metric["function"]( 82 | true_nn_distances, run_distances, metrics_cache, properties) 83 | results[name] = v 84 | if v: 85 | print('%s: %g' % (name, v)) 86 | return (algo, algo_name, results) 87 | 88 | 89 | def generate_n_colors(n): 90 | vs = numpy.linspace(0.3, 0.9, 7) 91 | colors = [(.9, .4, .4, 1.)] 92 | 93 | def euclidean(a, b): 94 | return sum((x - y)**2 for x, y in zip(a, b)) 95 | while len(colors) < n: 96 | new_color = max(itertools.product(vs, vs, vs), 97 | key=lambda a: min(euclidean(a, b) for b in colors)) 98 | colors.append(new_color + (1.,)) 99 | return colors 100 | 101 | 102 | def create_linestyles(unique_algorithms): 103 | colors = dict( 104 | zip(unique_algorithms, generate_n_colors(len(unique_algorithms)))) 105 | linestyles = dict((algo, ['--', '-.', '-', ':'][i % 4]) 106 | for i, algo in enumerate(unique_algorithms)) 107 | markerstyles = dict((algo, ['+', '<', 'o', '*', 'x'][i % 5]) 108 | for i, algo in enumerate(unique_algorithms)) 109 | faded = dict((algo, (r, g, b, 0.3)) 110 | for algo, (r, g, b, a) in colors.items()) 111 | return dict((algo, (colors[algo], faded[algo], 112 | linestyles[algo], markerstyles[algo])) 113 | for algo in unique_algorithms) 114 | 115 | 116 | def get_up_down(metric): 117 | if metric["worst"] == float("inf"): 118 | return "down" 119 | return "up" 120 | 121 | 122 | def get_left_right(metric): 123 | if metric["worst"] == float("inf"): 124 | return "left" 125 | return "right" 126 | 127 | 128 | def get_plot_label(xm, ym): 129 | template = ("%(xlabel)s-%(ylabel)s tradeoff - %(updown)s and" 130 | " to the %(leftright)s is better") 131 | return template % {"xlabel": xm["description"], 132 | "ylabel": ym["description"], 133 | "updown": get_up_down(ym), 134 | "leftright": get_left_right(xm)} 135 | -------------------------------------------------------------------------------- /templates/chartjs.template: -------------------------------------------------------------------------------- 1 |
Benchmarking Results
10 |Results are split by distance measure and dataset. In the bottom, you can find an overview of an algorithm's performance on all datasets. Each dataset is annoted 11 | by (k = ...), the number of nearest neighbors an algorithm was supposed to return. The plot shown depicts Recall (the fraction 12 | of true nearest neighbors found, on average over all queries) against Queries per second. Clicking on a plot reveils detailled interactive plots, including 13 | approximate recall, index size, and build time.
14 | {% for type in ['non-batch', 'batch'] %} 15 | {% if len(dataset_with_distances[type]) > 0 %} 16 | {% if type == 'batch' %} 17 |Benchmarks for Batched Queries
18 | {% else %} 19 |Benchmarks for Single Queries
20 | {% endif %} 21 | 22 |Results by Dataset
23 | {% for distance_data in dataset_with_distances[type] %} 24 |Distance: {{ distance_data.name }}
25 | {% for entry in distance_data.entries %} 26 | 27 |
28 |
35 |
36 |
29 |
31 | {{entry.desc}}
30 |
32 | }}.png)
33 |
34 |
33 | 37 | {% endfor %} 38 | {% endfor %} 39 |
Results by Algorithm
40 |- Algorithms:
41 | {% for algo in algorithms[type].keys() %}
42 |
- {{algo}} 43 | {% endfor %} 44 |
48 |
55 |
56 |
49 |
51 | {{algo}}
50 |
52 | }}.png)
53 |
54 |
53 | 57 | {% endfor %} 58 | {% endif %} 59 | {% endfor %} 60 | {% endblock %} 61 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/onng_ngt.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import sys 3 | import os 4 | import ngtpy 5 | import numpy as np 6 | import subprocess 7 | import time 8 | from ann_benchmarks.algorithms.base import BaseANN 9 | from ann_benchmarks.constants import INDEX_DIR 10 | 11 | 12 | class ONNG(BaseANN): 13 | def __init__(self, metric, object_type, epsilon, param): 14 | metrics = {'euclidean': '2', 'angular': 'C'} 15 | self._edge_size = int(param['edge']) 16 | self._outdegree = int(param['outdegree']) 17 | self._indegree = int(param['indegree']) 18 | self._metric = metrics[metric] 19 | self._object_type = object_type 20 | self._edge_size_for_search = int(param['search_edge']) if 'search_edge' in param.keys() else -2 21 | self._tree_disabled = (param['tree'] == False) if 'tree' in param.keys() else False 22 | self._build_time_limit = 4 23 | self._epsilon = epsilon 24 | print('ONNG: edge_size=' + str(self._edge_size)) 25 | print('ONNG: outdegree=' + str(self._outdegree)) 26 | print('ONNG: indegree=' + str(self._indegree)) 27 | print('ONNG: edge_size_for_search=' + str(self._edge_size_for_search)) 28 | print('ONNG: epsilon=' + str(self._epsilon)) 29 | print('ONNG: metric=' + metric) 30 | print('ONNG: object_type=' + object_type) 31 | 32 | def fit(self, X): 33 | print('ONNG: start indexing...') 34 | dim = len(X[0]) 35 | print('ONNG: # of data=' + str(len(X))) 36 | print('ONNG: dimensionality=' + str(dim)) 37 | index_dir = 'indexes' 38 | if not os.path.exists(index_dir): 39 | os.makedirs(index_dir) 40 | index = os.path.join( 41 | index_dir, 42 | 'ONNG-{}-{}-{}'.format(self._edge_size, self._outdegree, 43 | self._indegree)) 44 | anngIndex = os.path.join(index_dir, 'ANNG-' + str(self._edge_size)) 45 | print('ONNG: index=' + index) 46 | if (not os.path.exists(index)) and (not os.path.exists(anngIndex)): 47 | print('ONNG: create ANNG') 48 | t = time.time() 49 | args = ['ngt', 'create', '-it', '-p8', '-b500', '-ga', '-of', 50 | '-D' + self._metric, '-d' + str(dim), 51 | '-E' + str(self._edge_size), '-S0', 52 | '-e' + str(self._epsilon), '-P0', '-B30', 53 | '-T' + str(self._build_time_limit), anngIndex] 54 | subprocess.call(args) 55 | idx = ngtpy.Index(path=anngIndex) 56 | idx.batch_insert(X, num_threads=24, debug=False) 57 | idx.save() 58 | idx.close() 59 | print('ONNG: ANNG construction time(sec)=' + str(time.time() - t)) 60 | if not os.path.exists(index): 61 | print('ONNG: degree adjustment') 62 | t = time.time() 63 | args = ['ngt', 'reconstruct-graph', '-mS', 64 | '-o ' + str(self._outdegree), 65 | '-i ' + str(self._indegree), anngIndex, index] 66 | subprocess.call(args) 67 | print('ONNG: degree adjustment time(sec)=' + str(time.time() - t)) 68 | if os.path.exists(index): 69 | print('ONNG: index already exists! ' + str(index)) 70 | t = time.time() 71 | print(self._tree_disabled) 72 | #self.index = ngtpy.Index(index, read_only=True, tree_disabled=self._tree_disabled) 73 | self.index = ngtpy.Index(index, read_only=True) 74 | self.indexName = index 75 | print('ONNG: open time(sec)=' + str(time.time() - t)) 76 | else: 77 | print('ONNG: something wrong.') 78 | print('ONNG: end of fit') 79 | 80 | def set_query_arguments(self, epsilon): 81 | print("ONNG: epsilon=" + str(epsilon)) 82 | self._epsilon = epsilon - 1.0 83 | self.name = 'ONNG-NGT(%s, %s, %s, %s, %1.3f)' % ( 84 | self._edge_size, self._outdegree, 85 | self._indegree, self._edge_size_for_search, 86 | self._epsilon + 1.0) 87 | 88 | def query(self, v, n): 89 | results = self.index.search( 90 | v, n, self._epsilon, self._edge_size_for_search, 91 | with_distance=False) 92 | return results 93 | 94 | def freeIndex(self): 95 | print('ONNG: free') 96 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/elasticsearch.py: -------------------------------------------------------------------------------- 1 | """ 2 | ann-benchmarks interfaces for Elasticsearch. 3 | Note that this requires X-Pack, which is not included in the OSS version of Elasticsearch. 4 | """ 5 | import logging 6 | from time import sleep 7 | from urllib.error import URLError 8 | from urllib.request import Request, urlopen 9 | 10 | from elasticsearch import Elasticsearch 11 | from elasticsearch.helpers import bulk 12 | 13 | from ann_benchmarks.algorithms.base import BaseANN 14 | 15 | # Configure the elasticsearch logger. 16 | # By default, it writes an INFO statement for every request. 17 | logging.getLogger("elasticsearch").setLevel(logging.WARN) 18 | 19 | # Uncomment these lines if you want to see timing for every HTTP request and its duration. 20 | # logging.basicConfig(level=logging.INFO) 21 | # logging.getLogger("elasticsearch").setLevel(logging.INFO) 22 | 23 | def es_wait(): 24 | print("Waiting for elasticsearch health endpoint...") 25 | req = Request("http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s") 26 | for i in range(30): 27 | try: 28 | res = urlopen(req) 29 | if res.getcode() == 200: 30 | print("Elasticsearch is ready") 31 | return 32 | except URLError: 33 | pass 34 | sleep(1) 35 | raise RuntimeError("Failed to connect to local elasticsearch") 36 | 37 | 38 | class ElasticsearchScriptScoreQuery(BaseANN): 39 | """ 40 | KNN using the Elasticsearch dense_vector datatype and script score functions. 41 | - Dense vector field type: https://www.elastic.co/guide/en/elasticsearch/reference/master/dense-vector.html 42 | - Dense vector queries: https://www.elastic.co/guide/en/elasticsearch/reference/master/query-dsl-script-score-query.html 43 | """ 44 | 45 | def __init__(self, metric: str, dimension: int): 46 | self.name = f"elasticsearch-script-score-query_metric={metric}_dimension={dimension}" 47 | self.metric = metric 48 | self.dimension = dimension 49 | self.index = f"es-ssq-{metric}-{dimension}" 50 | self.es = Elasticsearch(["http://localhost:9200"]) 51 | self.batch_res = [] 52 | if self.metric == "euclidean": 53 | self.script = "1 / (1 + l2norm(params.query_vec, \"vec\"))" 54 | elif self.metric == "angular": 55 | self.script = "1.0 + cosineSimilarity(params.query_vec, \"vec\")" 56 | else: 57 | raise NotImplementedError(f"Not implemented for metric {self.metric}") 58 | es_wait() 59 | 60 | def fit(self, X): 61 | body = dict(settings=dict(number_of_shards=1, number_of_replicas=0)) 62 | mapping = dict( 63 | properties=dict( 64 | id=dict(type="keyword", store=True), 65 | vec=dict(type="dense_vector", dims=self.dimension) 66 | ) 67 | ) 68 | self.es.indices.create(self.index, body=body) 69 | self.es.indices.put_mapping(mapping, self.index) 70 | 71 | def gen(): 72 | for i, vec in enumerate(X): 73 | yield { "_op_type": "index", "_index": self.index, "vec": vec.tolist(), 'id': str(i + 1) } 74 | 75 | (_, errors) = bulk(self.es, gen(), chunk_size=500, max_retries=9) 76 | assert len(errors) == 0, errors 77 | 78 | self.es.indices.refresh(self.index) 79 | self.es.indices.forcemerge(self.index, max_num_segments=1) 80 | 81 | def query(self, q, n): 82 | body = dict( 83 | query=dict( 84 | script_score=dict( 85 | query=dict(match_all=dict()), 86 | script=dict( 87 | source=self.script, 88 | params=dict(query_vec=q.tolist()) 89 | ) 90 | ) 91 | ) 92 | ) 93 | res = self.es.search(index=self.index, body=body, size=n, _source=False, docvalue_fields=['id'], 94 | stored_fields="_none_", filter_path=["hits.hits.fields.id"]) 95 | return [int(h['fields']['id'][0]) - 1 for h in res['hits']['hits']] 96 | 97 | def batch_query(self, X, n): 98 | self.batch_res = [self.query(q, n) for q in X] 99 | 100 | def get_batch_results(self): 101 | return self.batch_res 102 | 103 | -------------------------------------------------------------------------------- /ann_benchmarks/algorithms/pynndescent.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | import pynndescent 3 | from ann_benchmarks.algorithms.base import BaseANN 4 | import numpy as np 5 | import scipy.sparse 6 | 7 | 8 | class PyNNDescent(BaseANN): 9 | def __init__(self, metric, index_param_dict, n_search_trees=1): 10 | if "n_neighbors" in index_param_dict: 11 | self._n_neighbors = int(index_param_dict["n_neighbors"]) 12 | else: 13 | self._n_neighbors = 30 14 | 15 | if "pruning_degree_multiplier" in index_param_dict: 16 | self._pruning_degree_multiplier = float( 17 | index_param_dict["pruning_degree_multiplier"] 18 | ) 19 | else: 20 | self._pruning_degree_multiplier = 1.5 21 | 22 | if "diversify_prob" in index_param_dict: 23 | self._diversify_prob = float(index_param_dict["diversify_prob"]) 24 | else: 25 | self._diversify_prob = 1.0 26 | 27 | if "leaf_size" in index_param_dict: 28 | self._leaf_size = int(index_param_dict["leaf_size"]) 29 | else: 30 | leaf_size = 32 31 | 32 | self._n_search_trees = int(n_search_trees) 33 | 34 | self._pynnd_metric = { 35 | "angular": "dot", 36 | # 'angular': 'cosine', 37 | "euclidean": "euclidean", 38 | "hamming": "hamming", 39 | "jaccard": "jaccard", 40 | }[metric] 41 | 42 | def _sparse_convert_for_fit(self, X): 43 | lil_data = [] 44 | self._n_cols = 1 45 | self._n_rows = len(X) 46 | for i in range(self._n_rows): 47 | lil_data.append([1] * len(X[i])) 48 | if max(X[i]) + 1 > self._n_cols: 49 | self._n_cols = max(X[i]) + 1 50 | 51 | result = scipy.sparse.lil_matrix( 52 | (self._n_rows, self._n_cols), dtype=np.int 53 | ) 54 | result.rows[:] = list(X) 55 | result.data[:] = lil_data 56 | return result.tocsr() 57 | 58 | def _sparse_convert_for_query(self, v): 59 | result = scipy.sparse.csr_matrix((1, self._n_cols), dtype=np.int) 60 | result.indptr = np.array([0, len(v)]) 61 | result.indices = np.array(v).astype(np.int32) 62 | result.data = np.ones(len(v), dtype=np.int) 63 | return result 64 | 65 | def fit(self, X): 66 | if self._pynnd_metric == "jaccard": 67 | # Convert to sparse matrix format 68 | X = self._sparse_convert_for_fit(X) 69 | 70 | self._index = pynndescent.NNDescent( 71 | X, 72 | n_neighbors=self._n_neighbors, 73 | metric=self._pynnd_metric, 74 | low_memory=True, 75 | leaf_size=self._leaf_size, 76 | pruning_degree_multiplier=self._pruning_degree_multiplier, 77 | diversify_prob=self._diversify_prob, 78 | n_search_trees=self._n_search_trees, 79 | compressed=True, 80 | verbose=True, 81 | ) 82 | if hasattr(self._index, "prepare"): 83 | self._index.prepare() 84 | else: 85 | self._index._init_search_graph() 86 | if self._index._is_sparse: 87 | if hasattr(self._index, "_init_sparse_search_function"): 88 | self._index._init_sparse_search_function() 89 | else: 90 | if hasattr(self._index, "_init_search_function"): 91 | self._index._init_search_function() 92 | 93 | def set_query_arguments(self, epsilon=0.1): 94 | self._epsilon = float(epsilon) 95 | 96 | def query(self, v, n): 97 | if self._pynnd_metric == "jaccard": 98 | # convert index array to sparse matrix format and query 99 | v = self._sparse_convert_for_query(v) 100 | ind, dist = self._index.query(v, k=n, epsilon=self._epsilon) 101 | else: 102 | ind, dist = self._index.query( 103 | v.reshape(1, -1).astype("float32"), k=n, epsilon=self._epsilon 104 | ) 105 | return ind[0] 106 | 107 | def __str__(self): 108 | str_template = "PyNNDescent(n_neighbors=%d, pruning_mult=%.2f, diversify_prob=%.3f, epsilon=%.3f, leaf_size=%02d)" 109 | return str_template % ( 110 | self._n_neighbors, 111 | self._pruning_degree_multiplier, 112 | self._diversify_prob, 113 | self._epsilon, 114 | self._leaf_size, 115 | ) 116 | -------------------------------------------------------------------------------- /ann_benchmarks/plotting/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import absolute_import 2 | 3 | import itertools 4 | import numpy 5 | from ann_benchmarks.plotting.metrics import all_metrics as metrics 6 | 7 | 8 | def get_or_create_metrics(run): 9 | if 'metrics' not in run: 10 | run.create_group('metrics') 11 | return run['metrics'] 12 | 13 | 14 | def create_pointset(data, xn, yn): 15 | xm, ym = (metrics[xn], metrics[yn]) 16 | rev_y = -1 if ym["worst"] < 0 else 1 17 | rev_x = -1 if xm["worst"] < 0 else 1 18 | data.sort(key=lambda t: (rev_y * t[-1], rev_x * t[-2])) 19 | 20 | axs, ays, als = [], [], [] 21 | # Generate Pareto frontier 22 | xs, ys, ls = [], [], [] 23 | last_x = xm["worst"] 24 | comparator = ((lambda xv, lx: xv > lx) 25 | if last_x < 0 else (lambda xv, lx: xv < lx)) 26 | for algo, algo_name, xv, yv in data: 27 | if not xv or not yv: 28 | continue 29 | axs.append(xv) 30 | ays.append(yv) 31 | als.append(algo_name) 32 | if comparator(xv, last_x): 33 | last_x = xv 34 | xs.append(xv) 35 | ys.append(yv) 36 | ls.append(algo_name) 37 | return xs, ys, ls, axs, ays, als 38 | 39 | 40 | def compute_metrics(true_nn_distances, res, metric_1, metric_2, 41 | recompute=False): 42 | all_results = {} 43 | for i, (properties, run) in enumerate(res): 44 | algo = properties['algo'] 45 | algo_name = properties['name'] 46 | # cache distances to avoid access to hdf5 file 47 | run_distances = numpy.array(run['distances']) 48 | if recompute and 'metrics' in run: 49 | del run['metrics'] 50 | metrics_cache = get_or_create_metrics(run) 51 | 52 | metric_1_value = metrics[metric_1]['function']( 53 | true_nn_distances, 54 | run_distances, metrics_cache, properties) 55 | metric_2_value = metrics[metric_2]['function']( 56 | true_nn_distances, 57 | run_distances, metrics_cache, properties) 58 | 59 | print('%3d: %80s %12.3f %12.3f' % 60 | (i, algo_name, metric_1_value, metric_2_value)) 61 | 62 | all_results.setdefault(algo, []).append( 63 | (algo, algo_name, metric_1_value, metric_2_value)) 64 | 65 | return all_results 66 | 67 | 68 | def compute_all_metrics(true_nn_distances, run, properties, recompute=False): 69 | algo = properties["algo"] 70 | algo_name = properties["name"] 71 | print('--') 72 | print(algo_name) 73 | results = {} 74 | # cache distances to avoid access to hdf5 file 75 | run_distances = numpy.array(run["distances"]) 76 | if recompute and 'metrics' in run: 77 | del run['metrics'] 78 | metrics_cache = get_or_create_metrics(run) 79 | 80 | for name, metric in metrics.items(): 81 | v = metric["function"]( 82 | true_nn_distances, run_distances, metrics_cache, properties) 83 | results[name] = v 84 | if v: 85 | print('%s: %g' % (name, v)) 86 | return (algo, algo_name, results) 87 | 88 | 89 | def generate_n_colors(n): 90 | vs = numpy.linspace(0.3, 0.9, 7) 91 | colors = [(.9, .4, .4, 1.)] 92 | 93 | def euclidean(a, b): 94 | return sum((x - y)**2 for x, y in zip(a, b)) 95 | while len(colors) < n: 96 | new_color = max(itertools.product(vs, vs, vs), 97 | key=lambda a: min(euclidean(a, b) for b in colors)) 98 | colors.append(new_color + (1.,)) 99 | return colors 100 | 101 | 102 | def create_linestyles(unique_algorithms): 103 | colors = dict( 104 | zip(unique_algorithms, generate_n_colors(len(unique_algorithms)))) 105 | linestyles = dict((algo, ['--', '-.', '-', ':'][i % 4]) 106 | for i, algo in enumerate(unique_algorithms)) 107 | markerstyles = dict((algo, ['+', '<', 'o', '*', 'x'][i % 5]) 108 | for i, algo in enumerate(unique_algorithms)) 109 | faded = dict((algo, (r, g, b, 0.3)) 110 | for algo, (r, g, b, a) in colors.items()) 111 | return dict((algo, (colors[algo], faded[algo], 112 | linestyles[algo], markerstyles[algo])) 113 | for algo in unique_algorithms) 114 | 115 | 116 | def get_up_down(metric): 117 | if metric["worst"] == float("inf"): 118 | return "down" 119 | return "up" 120 | 121 | 122 | def get_left_right(metric): 123 | if metric["worst"] == float("inf"): 124 | return "left" 125 | return "right" 126 | 127 | 128 | def get_plot_label(xm, ym): 129 | template = ("%(xlabel)s-%(ylabel)s tradeoff - %(updown)s and" 130 | " to the %(leftright)s is better") 131 | return template % {"xlabel": xm["description"], 132 | "ylabel": ym["description"], 133 | "updown": get_up_down(ym), 134 | "leftright": get_left_right(xm)} 135 | -------------------------------------------------------------------------------- /templates/chartjs.template: -------------------------------------------------------------------------------- 1 |
{{xlabel}}/{{ylabel}}
2 |
3 |
4 |
85 |
86 | {% if args.latex %}
87 |
88 |
92 |
97 |
102 | {% endif %}
103 |
--------------------------------------------------------------------------------
/ann_benchmarks/algorithms/elastiknn.py:
--------------------------------------------------------------------------------
1 | """
2 | ann-benchmarks interfaces for elastiknn: https://github.com/alexklibisz/elastiknn
3 | Uses the elastiknn python client
4 | To install a local copy of the client, run `pip install --upgrade -e /path/to/elastiknn/client-python/`
5 | To monitor the Elasticsearch JVM using Visualvm, add `ports={ "8097": 8097 }` to the `containers.run` call in runner.py.
6 | """
7 | from sys import stderr
8 | from urllib.error import URLError
9 |
10 | import numpy as np
11 | from elastiknn.api import Vec
12 | from elastiknn.models import ElastiknnModel
13 | from elastiknn.utils import dealias_metric
14 |
15 | from ann_benchmarks.algorithms.base import BaseANN
16 |
17 | from urllib.request import Request, urlopen
18 | from time import sleep, perf_counter
19 |
20 | import logging
21 |
22 | # Mute the elasticsearch logger.
23 | # By default, it writes an INFO statement for every request.
24 | logging.getLogger("elasticsearch").setLevel(logging.WARN)
25 |
26 |
27 | def es_wait():
28 | print("Waiting for elasticsearch health endpoint...")
29 | req = Request("http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s")
30 | for i in range(30):
31 | try:
32 | res = urlopen(req)
33 | if res.getcode() == 200:
34 | print("Elasticsearch is ready")
35 | return
36 | except URLError:
37 | pass
38 | sleep(1)
39 | raise RuntimeError("Failed to connect to local elasticsearch")
40 |
41 |
42 | class Exact(BaseANN):
43 |
44 | def __init__(self, metric: str, dimension: int):
45 | self.name = f"eknn-exact-metric={metric}_dimension={dimension}"
46 | self.metric = metric
47 | self.dimension = dimension
48 | self.model = ElastiknnModel("exact", dealias_metric(metric))
49 | self.batch_res = None
50 | es_wait()
51 |
52 | def _handle_sparse(self, X):
53 | # convert list of lists of indices to sparse vectors.
54 | return [Vec.SparseBool(x, self.dimension) for x in X]
55 |
56 | def fit(self, X):
57 | if self.metric in {'jaccard', 'hamming'}:
58 | return self.model.fit(self._handle_sparse(X), shards=1)[0]
59 | else:
60 | return self.model.fit(X, shards=1)
61 |
62 | def query(self, q, n):
63 | if self.metric in {'jaccard', 'hamming'}:
64 | return self.model.kneighbors(self._handle_sparse([q]), n)[0]
65 | else:
66 | return self.model.kneighbors(np.expand_dims(q, 0), n)[0]
67 |
68 | def batch_query(self, X, n):
69 | if self.metric in {'jaccard', 'hamming'}:
70 | self.batch_res = self.model.kneighbors(self._handle_sparse(X), n)
71 | else:
72 | self.batch_res = self.model.kneighbors(X, n)
73 |
74 | def get_batch_results(self):
75 | return self.batch_res
76 |
77 |
78 | class L2Lsh(BaseANN):
79 |
80 | def __init__(self, L: int, k: int, w: int):
81 | self.name_prefix = f"eknn-l2lsh-L={L}-k={k}-w={w}"
82 | self.name = None # set based on query args.
83 | self.model = ElastiknnModel("lsh", "l2", mapping_params=dict(L=L, k=k, w=w))
84 | self.X_max = 1.0
85 | self.query_params = dict()
86 | self.batch_res = None
87 | self.sum_query_dur = 0
88 | self.num_queries = 0
89 | es_wait()
90 |
91 | def fit(self, X):
92 | print(f"{self.name_prefix}: indexing {len(X)} vectors")
93 |
94 | # I found it's best to scale the vectors into [0, 1], i.e. divide by the max.
95 | self.X_max = X.max()
96 | return self.model.fit(X / self.X_max, shards=1)
97 |
98 | def set_query_arguments(self, candidates: int, probes: int):
99 | # This gets called when starting a new batch of queries.
100 | # Update the name and model's query parameters based on the given params.
101 | self.name = f"{self.name_prefix}_candidates={candidates}_probes={probes}"
102 | self.model.set_query_params(dict(candidates=candidates, probes=probes))
103 | # Reset the counters.
104 | self.num_queries = 0
105 | self.sum_query_dur = 0
106 |
107 | def query(self, q, n):
108 | # If QPS after 100 queries is < 10, this setting is bad and won't complete within the default timeout.
109 | if self.num_queries > 100 and self.num_queries / self.sum_query_dur < 10:
110 | print("Throughput after 100 queries is less than 10 q/s. Terminating to avoid wasteful computation.", flush=True)
111 | exit(0)
112 | else:
113 | t0 = perf_counter()
114 | res = self.model.kneighbors(np.expand_dims(q, 0) / self.X_max, n)[0]
115 | dur = (perf_counter() - t0)
116 | self.sum_query_dur += dur
117 | self.num_queries += 1
118 | return res
119 |
120 | def batch_query(self, X, n):
121 | self.batch_res = self.model.kneighbors(X, n)
122 |
123 | def get_batch_results(self):
124 | return self.batch_res
125 |
--------------------------------------------------------------------------------
/plot.py:
--------------------------------------------------------------------------------
1 | import os
2 | import matplotlib as mpl
3 | mpl.use('Agg') # noqa
4 | import matplotlib.pyplot as plt
5 | import numpy as np
6 | import argparse
7 |
8 | from ann_benchmarks.datasets import get_dataset
9 | from ann_benchmarks.algorithms.definitions import get_definitions
10 | from ann_benchmarks.plotting.metrics import all_metrics as metrics
11 | from ann_benchmarks.plotting.utils import (get_plot_label, compute_metrics,
12 | create_linestyles, create_pointset)
13 | from ann_benchmarks.results import (store_results, load_all_results,
14 | get_unique_algorithms, get_algorithm_name)
15 |
16 |
17 | def create_plot(all_data, raw, x_log, y_log, xn, yn, fn_out, linestyles,
18 | batch):
19 | xm, ym = (metrics[xn], metrics[yn])
20 | # Now generate each plot
21 | handles = []
22 | labels = []
23 | plt.figure(figsize=(12, 9))
24 | for algo in sorted(all_data.keys(), key=lambda x: x.lower()):
25 | xs, ys, ls, axs, ays, als = create_pointset(all_data[algo], xn, yn)
26 | color, faded, linestyle, marker = linestyles[algo]
27 | handle, = plt.plot(xs, ys, '-', label=algo, color=color,
28 | ms=7, mew=3, lw=3, linestyle=linestyle,
29 | marker=marker)
30 | handles.append(handle)
31 | if raw:
32 | handle2, = plt.plot(axs, ays, '-', label=algo, color=faded,
33 | ms=5, mew=2, lw=2, linestyle=linestyle,
34 | marker=marker)
35 | labels.append(get_algorithm_name(algo, batch))
36 |
37 | if x_log:
38 | plt.gca().set_xscale('log')
39 | if y_log:
40 | plt.gca().set_yscale('log')
41 | plt.gca().set_title(get_plot_label(xm, ym))
42 | plt.gca().set_ylabel(ym['description'])
43 | plt.gca().set_xlabel(xm['description'])
44 | box = plt.gca().get_position()
45 | # plt.gca().set_position([box.x0, box.y0, box.width * 0.8, box.height])
46 | plt.gca().legend(handles, labels, loc='center left',
47 | bbox_to_anchor=(1, 0.5), prop={'size': 9})
48 | plt.grid(b=True, which='major', color='0.65', linestyle='-')
49 | if 'lim' in xm:
50 | plt.xlim(xm['lim'])
51 | if 'lim' in ym:
52 | plt.ylim(ym['lim'])
53 | plt.savefig(fn_out, bbox_inches='tight')
54 | plt.close()
55 |
56 |
57 | if __name__ == "__main__":
58 | parser = argparse.ArgumentParser()
59 | parser.add_argument(
60 | '--dataset',
61 | metavar="DATASET",
62 | default='glove-100-angular')
63 | parser.add_argument(
64 | '--count',
65 | default=10)
66 | parser.add_argument(
67 | '--definitions',
68 | metavar='FILE',
69 | help='load algorithm definitions from FILE',
70 | default='algos.yaml')
71 | parser.add_argument(
72 | '--limit',
73 | default=-1)
74 | parser.add_argument(
75 | '-o', '--output')
76 | parser.add_argument(
77 | '-x', '--x-axis',
78 | help='Which metric to use on the X-axis',
79 | choices=metrics.keys(),
80 | default="k-nn")
81 | parser.add_argument(
82 | '-y', '--y-axis',
83 | help='Which metric to use on the Y-axis',
84 | choices=metrics.keys(),
85 | default="qps")
86 | parser.add_argument(
87 | '-X', '--x-log',
88 | help='Draw the X-axis using a logarithmic scale',
89 | action='store_true')
90 | parser.add_argument(
91 | '-Y', '--y-log',
92 | help='Draw the Y-axis using a logarithmic scale',
93 | action='store_true')
94 | parser.add_argument(
95 | '--raw',
96 | help='Show raw results (not just Pareto frontier) in faded colours',
97 | action='store_true')
98 | parser.add_argument(
99 | '--batch',
100 | help='Plot runs in batch mode',
101 | action='store_true')
102 | parser.add_argument(
103 | '--recompute',
104 | help='Clears the cache and recomputes the metrics',
105 | action='store_true')
106 | args = parser.parse_args()
107 |
108 | if not args.output:
109 | args.output = 'results/%s.png' % get_algorithm_name(
110 | args.dataset, args.batch)
111 | print('writing output to %s' % args.output)
112 |
113 | dataset = get_dataset(args.dataset)
114 | count = int(args.count)
115 | unique_algorithms = get_unique_algorithms()
116 | results = load_all_results(args.dataset, count, True, args.batch)
117 | linestyles = create_linestyles(sorted(unique_algorithms))
118 | runs = compute_metrics(np.array(dataset["distances"]),
119 | results, args.x_axis, args.y_axis, args.recompute)
120 | if not runs:
121 | raise Exception('Nothing to plot')
122 |
123 | create_plot(runs, args.raw, args.x_log,
124 | args.y_log, args.x_axis, args.y_axis, args.output,
125 | linestyles, args.batch)
126 |
--------------------------------------------------------------------------------
/ann_benchmarks/algorithms/bruteforce.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | import numpy
3 | import sklearn.neighbors
4 | from ann_benchmarks.distance import metrics as pd
5 | from ann_benchmarks.algorithms.base import BaseANN
6 |
7 |
8 | class BruteForce(BaseANN):
9 | def __init__(self, metric):
10 | if metric not in ('angular', 'euclidean', 'hamming'):
11 | raise NotImplementedError(
12 | "BruteForce doesn't support metric %s" % metric)
13 | self._metric = metric
14 | self.name = 'BruteForce()'
15 |
16 | def fit(self, X):
17 | metric = {'angular': 'cosine', 'euclidean': 'l2',
18 | 'hamming': 'hamming'}[self._metric]
19 | self._nbrs = sklearn.neighbors.NearestNeighbors(
20 | algorithm='brute', metric=metric)
21 | self._nbrs.fit(X)
22 |
23 | def query(self, v, n):
24 | return list(self._nbrs.kneighbors(
25 | [v], return_distance=False, n_neighbors=n)[0])
26 |
27 | def query_with_distances(self, v, n):
28 | (distances, positions) = self._nbrs.kneighbors(
29 | [v], return_distance=True, n_neighbors=n)
30 | return zip(list(positions[0]), list(distances[0]))
31 |
32 |
33 | class BruteForceBLAS(BaseANN):
34 | """kNN search that uses a linear scan = brute force."""
35 |
36 | def __init__(self, metric, precision=numpy.float32):
37 | if metric not in ('angular', 'euclidean', 'hamming', 'jaccard'):
38 | raise NotImplementedError(
39 | "BruteForceBLAS doesn't support metric %s" % metric)
40 | elif metric == 'hamming' and precision != numpy.bool:
41 | raise NotImplementedError(
42 | "BruteForceBLAS doesn't support precision"
43 | " %s with Hamming distances" % precision)
44 | self._metric = metric
45 | self._precision = precision
46 | self.name = 'BruteForceBLAS()'
47 |
48 | def fit(self, X):
49 | """Initialize the search index."""
50 | if self._metric == 'angular':
51 | # precompute (squared) length of each vector
52 | lens = (X ** 2).sum(-1)
53 | # normalize index vectors to unit length
54 | X /= numpy.sqrt(lens)[..., numpy.newaxis]
55 | self.index = numpy.ascontiguousarray(X, dtype=self._precision)
56 | elif self._metric == 'hamming':
57 | # Regarding bitvectors as vectors in l_2 is faster for blas
58 | X = X.astype(numpy.float32)
59 | # precompute (squared) length of each vector
60 | lens = (X ** 2).sum(-1)
61 | self.index = numpy.ascontiguousarray(X, dtype=numpy.float32)
62 | self.lengths = numpy.ascontiguousarray(lens, dtype=numpy.float32)
63 | elif self._metric == 'euclidean':
64 | # precompute (squared) length of each vector
65 | lens = (X ** 2).sum(-1)
66 | self.index = numpy.ascontiguousarray(X, dtype=self._precision)
67 | self.lengths = numpy.ascontiguousarray(lens, dtype=self._precision)
68 | elif self._metric == 'jaccard':
69 | self.index = X
70 | else:
71 | # shouldn't get past the constructor!
72 | assert False, "invalid metric"
73 |
74 | def query(self, v, n):
75 | return [index for index, _ in self.query_with_distances(v, n)]
76 |
77 | def query_with_distances(self, v, n):
78 | """Find indices of `n` most similar vectors from the index to query
79 | vector `v`."""
80 |
81 | if self._metric != 'jaccard':
82 | # use same precision for query as for index
83 | v = numpy.ascontiguousarray(v, dtype=self.index.dtype)
84 |
85 | # HACK we ignore query length as that's a constant
86 | # not affecting the final ordering
87 | if self._metric == 'angular':
88 | # argmax_a cossim(a, b) = argmax_a dot(a, b) / |a||b| = argmin_a -dot(a, b) # noqa
89 | dists = -numpy.dot(self.index, v)
90 | elif self._metric == 'euclidean':
91 | # argmin_a (a - b)^2 = argmin_a a^2 - 2ab + b^2 = argmin_a a^2 - 2ab # noqa
92 | dists = self.lengths - 2 * numpy.dot(self.index, v)
93 | elif self._metric == 'hamming':
94 | # Just compute hamming distance using euclidean distance
95 | dists = self.lengths - 2 * numpy.dot(self.index, v)
96 | elif self._metric == 'jaccard':
97 | dists = [pd[self._metric]['distance'](v, e) for e in self.index]
98 | else:
99 | # shouldn't get past the constructor!
100 | assert False, "invalid metric"
101 | # partition-sort by distance, get `n` closest
102 | nearest_indices = numpy.argpartition(dists, n)[:n]
103 | indices = [idx for idx in nearest_indices if pd[self._metric]
104 | ["distance_valid"](dists[idx])]
105 |
106 | def fix(index):
107 | ep = self.index[index]
108 | ev = v
109 | return (index, pd[self._metric]['distance'](ep, ev))
110 | return map(fix, indices)
111 |
--------------------------------------------------------------------------------
/ann_benchmarks/plotting/metrics.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | import numpy as np
3 |
4 |
5 | def knn_threshold(data, count, epsilon):
6 | return data[count - 1] + epsilon
7 |
8 |
9 | def epsilon_threshold(data, count, epsilon):
10 | return data[count - 1] * (1 + epsilon)
11 |
12 |
13 | def get_recall_values(dataset_distances, run_distances, count, threshold,
14 | epsilon=1e-3):
15 | recalls = np.zeros(len(run_distances))
16 | for i in range(len(run_distances)):
17 | t = threshold(dataset_distances[i], count, epsilon)
18 | actual = 0
19 | for d in run_distances[i][:count]:
20 | if d <= t:
21 | actual += 1
22 | recalls[i] = actual
23 | return (np.mean(recalls) / float(count),
24 | np.std(recalls) / float(count),
25 | recalls)
26 |
27 |
28 | def knn(dataset_distances, run_distances, count, metrics, epsilon=1e-3):
29 | if 'knn' not in metrics:
30 | print('Computing knn metrics')
31 | knn_metrics = metrics.create_group('knn')
32 | mean, std, recalls = get_recall_values(dataset_distances,
33 | run_distances, count,
34 | knn_threshold, epsilon)
35 | knn_metrics.attrs['mean'] = mean
36 | knn_metrics.attrs['std'] = std
37 | knn_metrics['recalls'] = recalls
38 | else:
39 | print("Found cached result")
40 | return metrics['knn']
41 |
42 |
43 | def epsilon(dataset_distances, run_distances, count, metrics, epsilon=0.01):
44 | s = 'eps' + str(epsilon)
45 | if s not in metrics:
46 | print('Computing epsilon metrics')
47 | epsilon_metrics = metrics.create_group(s)
48 | mean, std, recalls = get_recall_values(dataset_distances,
49 | run_distances, count,
50 | epsilon_threshold, epsilon)
51 | epsilon_metrics.attrs['mean'] = mean
52 | epsilon_metrics.attrs['std'] = std
53 | epsilon_metrics['recalls'] = recalls
54 | else:
55 | print("Found cached result")
56 | return metrics[s]
57 |
58 |
59 | def rel(dataset_distances, run_distances, metrics):
60 | if 'rel' not in metrics.attrs:
61 | print('Computing rel metrics')
62 | total_closest_distance = 0.0
63 | total_candidate_distance = 0.0
64 | for true_distances, found_distances in zip(dataset_distances,
65 | run_distances):
66 | for rdist, cdist in zip(true_distances, found_distances):
67 | total_closest_distance += rdist
68 | total_candidate_distance += cdist
69 | if total_closest_distance < 0.01:
70 | metrics.attrs['rel'] = float("inf")
71 | else:
72 | metrics.attrs['rel'] = total_candidate_distance / \
73 | total_closest_distance
74 | else:
75 | print("Found cached result")
76 | return metrics.attrs['rel']
77 |
78 |
79 | def queries_per_second(queries, attrs):
80 | return 1.0 / attrs["best_search_time"]
81 |
82 |
83 | def index_size(queries, attrs):
84 | # TODO(erikbern): should replace this with peak memory usage or something
85 | return attrs.get("index_size", 0)
86 |
87 |
88 | def build_time(queries, attrs):
89 | return attrs["build_time"]
90 |
91 |
92 | def candidates(queries, attrs):
93 | return attrs["candidates"]
94 |
95 |
96 | def dist_computations(queries, attrs):
97 | return attrs.get("dist_comps", 0) / (attrs['run_count'] * len(queries))
98 |
99 |
100 | all_metrics = {
101 | "k-nn": {
102 | "description": "Recall",
103 | "function": lambda true_distances, run_distances, metrics, run_attrs: knn(true_distances, run_distances, run_attrs["count"], metrics).attrs['mean'], # noqa
104 | "worst": float("-inf"),
105 | "lim": [0.0, 1.03]
106 | },
107 | "epsilon": {
108 | "description": "Epsilon 0.01 Recall",
109 | "function": lambda true_distances, run_distances, metrics, run_attrs: epsilon(true_distances, run_distances, run_attrs["count"], metrics).attrs['mean'], # noqa
110 | "worst": float("-inf")
111 | },
112 | "largeepsilon": {
113 | "description": "Epsilon 0.1 Recall",
114 | "function": lambda true_distances, run_distances, metrics, run_attrs: epsilon(true_distances, run_distances, run_attrs["count"], metrics, 0.1).attrs['mean'], # noqa
115 | "worst": float("-inf")
116 | },
117 | "rel": {
118 | "description": "Relative Error",
119 | "function": lambda true_distances, run_distances, metrics, run_attrs: rel(true_distances, run_distances, metrics), # noqa
120 | "worst": float("inf")
121 | },
122 | "qps": {
123 | "description": "Queries per second (1/s)",
124 | "function": lambda true_distances, run_distances, metrics, run_attrs: queries_per_second(true_distances, run_attrs), # noqa
125 | "worst": float("-inf")
126 | },
127 | "distcomps": {
128 | "description": "Distance computations",
129 | "function": lambda true_distances, run_distances, metrics, run_attrs: dist_computations(true_distances, run_attrs), # noqa
130 | "worst": float("inf")
131 | },
132 | "build": {
133 | "description": "Build time (s)",
134 | "function": lambda true_distances, run_distances, metrics, run_attrs: build_time(true_distances, run_attrs), # noqa
135 | "worst": float("inf")
136 | },
137 | "candidates": {
138 | "description": "Candidates generated",
139 | "function": lambda true_distances, run_distances, metrics, run_attrs: candidates(true_distances, run_attrs), # noqa
140 | "worst": float("inf")
141 | },
142 | "indexsize": {
143 | "description": "Index size (kB)",
144 | "function": lambda true_distances, run_distances, metrics, run_attrs: index_size(true_distances, run_attrs), # noqa
145 | "worst": float("inf")
146 | },
147 | "queriessize": {
148 | "description": "Index size (kB)/Queries per second (s)",
149 | "function": lambda true_distances, run_distances, metrics, run_attrs: index_size(true_distances, run_attrs) / queries_per_second(true_distances, run_attrs), # noqa
150 | "worst": float("inf")
151 | }
152 | }
153 |
--------------------------------------------------------------------------------
/ann_benchmarks/algorithms/definitions.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | from os import sep as pathsep
3 | import collections
4 | import importlib
5 | import os
6 | import sys
7 | import traceback
8 | import yaml
9 | from enum import Enum
10 | from itertools import product
11 |
12 |
13 | Definition = collections.namedtuple(
14 | 'Definition',
15 | ['algorithm', 'constructor', 'module', 'docker_tag',
16 | 'arguments', 'query_argument_groups', 'disabled'])
17 |
18 |
19 | def get_algorithm_name(name, batch):
20 | if batch:
21 | return name + "-batch"
22 | return name
23 |
24 |
25 | def instantiate_algorithm(definition):
26 | print('Trying to instantiate %s.%s(%s)' %
27 | (definition.module, definition.constructor, definition.arguments))
28 | module = importlib.import_module(definition.module)
29 | constructor = getattr(module, definition.constructor)
30 | return constructor(*definition.arguments)
31 |
32 |
33 | class InstantiationStatus(Enum):
34 | AVAILABLE = 0
35 | NO_CONSTRUCTOR = 1
36 | NO_MODULE = 2
37 |
38 |
39 | def algorithm_status(definition):
40 | try:
41 | module = importlib.import_module(definition.module)
42 | if hasattr(module, definition.constructor):
43 | return InstantiationStatus.AVAILABLE
44 | else:
45 | return InstantiationStatus.NO_CONSTRUCTOR
46 | except ImportError:
47 | return InstantiationStatus.NO_MODULE
48 |
49 |
50 | def _generate_combinations(args):
51 | if isinstance(args, list):
52 | args = [el if isinstance(el, list) else [el] for el in args]
53 | return [list(x) for x in product(*args)]
54 | elif isinstance(args, dict):
55 | flat = []
56 | for k, v in args.items():
57 | if isinstance(v, list):
58 | flat.append([(k, el) for el in v])
59 | else:
60 | flat.append([(k, v)])
61 | return [dict(x) for x in product(*flat)]
62 | else:
63 | raise TypeError("No args handling exists for %s" % type(args).__name__)
64 |
65 |
66 | def _substitute_variables(arg, vs):
67 | if isinstance(arg, dict):
68 | return dict([(k, _substitute_variables(v, vs))
69 | for k, v in arg.items()])
70 | elif isinstance(arg, list):
71 | return [_substitute_variables(a, vs) for a in arg]
72 | elif isinstance(arg, str) and arg in vs:
73 | return vs[arg]
74 | else:
75 | return arg
76 |
77 |
78 | def _get_definitions(definition_file):
79 | with open(definition_file, "r") as f:
80 | return yaml.load(f, yaml.SafeLoader)
81 |
82 |
83 | def list_algorithms(definition_file):
84 | definitions = _get_definitions(definition_file)
85 |
86 | print('The following algorithms are supported...')
87 | for point in definitions:
88 | print('\t... for the point type "%s"...' % point)
89 | for metric in definitions[point]:
90 | print('\t\t... and the distance metric "%s":' % metric)
91 | for algorithm in definitions[point][metric]:
92 | print('\t\t\t%s' % algorithm)
93 |
94 |
95 | def get_unique_algorithms(definition_file):
96 | definitions = _get_definitions(definition_file)
97 | algos = set()
98 | for point in definitions:
99 | for metric in definitions[point]:
100 | for algorithm in definitions[point][metric]:
101 | algos.add(algorithm)
102 | return list(sorted(algos))
103 |
104 |
105 | def get_definitions(definition_file, dimension, point_type="float",
106 | distance_metric="euclidean", count=10):
107 | definitions = _get_definitions(definition_file)
108 |
109 | algorithm_definitions = {}
110 | if "any" in definitions[point_type]:
111 | algorithm_definitions.update(definitions[point_type]["any"])
112 | algorithm_definitions.update(definitions[point_type][distance_metric])
113 |
114 | definitions = []
115 | for (name, algo) in algorithm_definitions.items():
116 | for k in ['docker-tag', 'module', 'constructor']:
117 | if k not in algo:
118 | raise Exception(
119 | 'algorithm %s does not define a "%s" property' % (name, k))
120 |
121 | base_args = []
122 | if "base-args" in algo:
123 | base_args = algo["base-args"]
124 |
125 | for run_group in algo["run-groups"].values():
126 | if "arg-groups" in run_group:
127 | groups = []
128 | for arg_group in run_group["arg-groups"]:
129 | if isinstance(arg_group, dict):
130 | # Dictionaries need to be expanded into lists in order
131 | # for the subsequent call to _generate_combinations to
132 | # do the right thing
133 | groups.append(_generate_combinations(arg_group))
134 | else:
135 | groups.append(arg_group)
136 | args = _generate_combinations(groups)
137 | elif "args" in run_group:
138 | args = _generate_combinations(run_group["args"])
139 | else:
140 | assert False, "? what? %s" % run_group
141 |
142 | if "query-arg-groups" in run_group:
143 | groups = []
144 | for arg_group in run_group["query-arg-groups"]:
145 | if isinstance(arg_group, dict):
146 | groups.append(_generate_combinations(arg_group))
147 | else:
148 | groups.append(arg_group)
149 | query_args = _generate_combinations(groups)
150 | elif "query-args" in run_group:
151 | query_args = _generate_combinations(run_group["query-args"])
152 | else:
153 | query_args = []
154 |
155 | for arg_group in args:
156 | aargs = []
157 | aargs.extend(base_args)
158 | if isinstance(arg_group, list):
159 | aargs.extend(arg_group)
160 | else:
161 | aargs.append(arg_group)
162 |
163 | vs = {
164 | "@count": count,
165 | "@metric": distance_metric,
166 | "@dimension": dimension
167 | }
168 | aargs = [_substitute_variables(arg, vs) for arg in aargs]
169 | definitions.append(Definition(
170 | algorithm=name,
171 | docker_tag=algo['docker-tag'],
172 | module=algo['module'],
173 | constructor=algo['constructor'],
174 | arguments=aargs,
175 | query_argument_groups=query_args,
176 | disabled=algo.get('disabled', False)
177 | ))
178 |
179 | return definitions
180 |
--------------------------------------------------------------------------------
/protocol/bf-runner.py:
--------------------------------------------------------------------------------
1 | from __future__ import print_function
2 |
3 | import sys
4 | from enum import Enum
5 | from shlex import split
6 |
7 | from ann_benchmarks.data import type_info
8 | from ann_benchmarks.distance import metrics
9 | from ann_benchmarks.algorithms.bruteforce import BruteForce, BruteForceBLAS
10 |
11 |
12 | class QueryMode(Enum):
13 | NORMAL = 0,
14 | PREPARED = 1,
15 | BATCH = 2
16 |
17 |
18 | __true_print = print
19 |
20 |
21 | def print(*args, **kwargs):
22 | __true_print(*args, **kwargs)
23 | sys.stdout.flush()
24 |
25 |
26 | def next_line():
27 | for line in iter(sys.stdin.readline, ''):
28 | yield split(line.strip())
29 |
30 |
31 | if __name__ == '__main__':
32 | point_type = None
33 | distance = None
34 | query_mode = QueryMode.NORMAL
35 | fast = False
36 | query_parameters = False
37 | # Configuration mode
38 | for line in next_line():
39 | if not line:
40 | break
41 | elif len(line) == 2:
42 | var, val = line[0], line[1]
43 | if var == "point-type":
44 | if val in type_info:
45 | point_type = type_info[val]
46 | print("epbprtv0 ok")
47 | else:
48 | print("epbprtv0 fail")
49 | elif var == "distance":
50 | if val in metrics:
51 | distance = val
52 | print("epbprtv0 ok")
53 | else:
54 | print("epbprtv0 fail")
55 | elif var == "fast":
56 | fast = (val == "1")
57 | print("epbprtv0 ok")
58 | else:
59 | print("epbprtv0 fail")
60 | elif len(line) == 3 and line[0] == "frontend":
61 | var, val = line[1], line[2]
62 | if var == "prepared-queries":
63 | query_mode = \
64 | QueryMode.PREPARED if val == "1" else QueryMode.NORMAL
65 | print("epbprtv0 ok")
66 | elif var == "batch-queries":
67 | query_mode = \
68 | QueryMode.BATCH if val == "1" else QueryMode.NORMAL
69 | print("epbprtv0 ok")
70 | elif var == "query-parameters":
71 | query_parameters = (val == "1")
72 | print("epbprtv0 ok")
73 | else:
74 | print("epbprtv0 fail")
75 | else:
76 | print("epbprtv0 fail")
77 | if point_type and distance:
78 | print("epbprtv0 ok")
79 | else:
80 | print("epbprtv0 fail")
81 | sys.exit(1)
82 |
83 | obj = None
84 | if not fast:
85 | obj = BruteForce(distance)
86 | else:
87 | obj = BruteForceBLAS(distance)
88 |
89 | parser = point_type["parse_entry"]
90 | # Training mode
91 | points = []
92 | for line in next_line():
93 | if not line:
94 | break
95 | elif len(line) == 1:
96 | point = line[0]
97 | try:
98 | parsed = parser(point)
99 | print("epbprtv0 ok %d" % len(points))
100 | points.append(parsed)
101 | except ValueError:
102 | print("epbprtv0 fail")
103 | else:
104 | print("epbprtv0 fail len %d" % len(line))
105 | if "finish_entries" in point_type:
106 | points = point_type["finish_entries"](points)
107 | obj.fit(points)
108 | print("epbprtv0 ok %d" % len(points))
109 |
110 | def _query_parameters(line):
111 | if hasattr(obj, "set_query_arguments"):
112 | try:
113 | obj.set_query_arguments(*line[1:-1])
114 | print("epbprtv0 ok")
115 | except TypeError:
116 | print("epbprtv0 fail")
117 | else:
118 | print("epbprtv0 fail")
119 |
120 | if query_mode == QueryMode.NORMAL:
121 | # Query mode
122 | for line in next_line():
123 | if not line:
124 | break
125 | elif query_parameters and line[0] == "query-params" \
126 | and line[-1] == "set":
127 | _query_parameters(line)
128 | elif len(line) == 2:
129 | try:
130 | query_point, k = line[0], int(line[1])
131 | parsed = parser(query_point)
132 | results = obj.query(parsed, k)
133 | if results:
134 | print("epbprtv0 ok %d" % len(results))
135 | for index in results:
136 | print("epbprtv0 %d" % index)
137 | else:
138 | print("epbprtv0 fail")
139 | except ValueError:
140 | print("epbprtv0 fail")
141 | else:
142 | print("epbprtv0 fail")
143 | elif query_mode == QueryMode.PREPARED:
144 | # Prepared query mode
145 | parsed = None
146 | k = None
147 | for line in next_line():
148 | if not line:
149 | break
150 | elif query_parameters and line[0] == "query-params" \
151 | and line[-1] == "set":
152 | _query_parameters(line)
153 | elif line == ["query"]:
154 | if parsed and k:
155 | results = obj.query(parsed, k)
156 | if results:
157 | print("epbprtv0 ok %d" % len(results))
158 | for index in results:
159 | print("epbprtv0 %d" % index)
160 | else:
161 | print("epbprtv0 fail")
162 | else:
163 | print("epbprtv0 fail")
164 | elif len(line) == 2:
165 | try:
166 | parsed, k = parser(line[0]), int(line[1])
167 | print("epbprtv0 ok prepared true")
168 | except ValueError:
169 | print("epbprtv0 fail")
170 | else:
171 | print("epbprtv0 fail")
172 | elif query_mode == QueryMode.BATCH:
173 | # Batch query mode
174 | parsed = None
175 | k = None
176 | for line in next_line():
177 | if not line:
178 | break
179 | elif query_parameters and line[0] == "query-params" \
180 | and line[-1] == "set":
181 | _query_parameters(line)
182 | elif line == ["query"]:
183 | if parsed and k:
184 | results = obj.batch_query(parsed, k)
185 | print("epbprtv0 ok")
186 | for result in obj.get_batch_results():
187 | if result:
188 | print("epbprtv0 ok %d" % len(result))
189 | for index in result:
190 | print("epbprtv0 %d" % index)
191 | else:
192 | print("epbprtv0 fail")
193 | else:
194 | print("epbprtv0 fail")
195 | elif len(line) > 1:
196 | try:
197 | parsed, k = map(parser, line[0:-1]), int(line[-1])
198 | print("epbprtv0 ok")
199 | except ValueError as e:
200 | print("epbprtv0 fail" % e)
201 | else:
202 | print("epbprtv0 fail")
203 | pass
204 | print("epbprtv0 ok")
205 |
--------------------------------------------------------------------------------
/protocol/specification.md:
--------------------------------------------------------------------------------
1 | This document specifies a simple text-based protocol that can be used to benchmark algorithms that don't have a Python wrapper. A program that implements the algorithm side of this specification will be referred to in the rest of this document as a "front-end".
2 |
3 | This protocol is line-oriented; both sides should configure their input and output streams to be line-buffered. Front-ends receive messages by reading lines from standard input and send messages by writing lines to standard output.
4 |
5 | ## Modes
6 |
7 | A front-end begins in configuration mode. When configuration is complete, it transitions into training mode; when training data has been supplied, into query mode; and, when no more queries remain, it terminates. It isn't possible to return from one mode to an earlier mode without restarting the front-end.
8 |
9 | A front-end reads lines from standard input, tokenises them, and interprets them according to its current mode; responses are written as lines to standard output. To enable protocol responses to be distinguished from other messages that may appear on standard output, the first token of a line containing a response will always be `epbprtv0`; the second will be `ok` when a command succeeds, potentially followed by other tokens, and `fail` when it doesn't.
10 |
11 | (The obscure token `epbprtv0` is intended to uniquely identify this protocol, and is meant to suggest something like "**e**xternal **p**rogram **b**enchmarking **pr**o**t**ocol, **v**ersion **0**".)
12 |
13 | A front-end may choose to include extra tokens in its responses after the tokens required by this specification to communicate more information back to the caller.
14 |
15 | ## Tokenisation
16 |
17 | Both the front-end and `ann-benchmarks` perform *tokenisation* on the lines of text they send and receive. The rules for tokenisation are as follows:
18 |
19 | * A token is a sequence of characters separated by one or more whitespace characters.
20 |
21 | Input | Token 1 | Token 2 | Token 3
22 | ----- | ------- | ------- | -------
23 | abc | abc | |
24 | a bc | a | bc |
25 | a bc | a | bc |
26 | a b c | a | b | c
27 |
28 | * A sequence surrounded by single quote marks will be treated as part of a token, even if it contains whitespace or doesn't contain any other characters.
29 |
30 | Input | Token 1 | Token 2 | Token 3
31 | ----- | ------- | ------- | -------
32 | 'a b c' | a b c | |
33 | 'a b c'd | a b cd | |
34 | a '' b | a | *empty string* | b
35 |
36 | * A sequence surrounded by double quote marks will be treated as part of a token, even if it contains whitespace or doesn't contain any other characters.
37 |
38 | Input | Token 1 | Token 2 | Token 3
39 | ----- | ------- | ------- | -------
40 | "a b c" | a b c | |
41 | "a b c"d | a b cd | |
42 | a "" b | a | *empty string* | b
43 |
44 | * Outside of a quoted sequence, preceding a character with a backslash causes any special significance it may have to be ignored; the character is then said to have been "escaped".
45 |
46 | Input | Token 1 | Token 2 | Token 3
47 | ----- | ------- | ------- | -------
48 | \a \b \c | a | b | c
49 |
50 | An escaped whitespace character doesn't separate tokens:
51 |
52 | Input | Token 1 | Token 2
53 | ----- | ------- | -------
54 | a b\ c | a | b c |
55 | "a b c"\ d | a b c d |
56 |
57 | An escaped quote mark doesn't begin a sequence:
58 |
59 | Input | Token 1 | Token 2 | Token 3
60 | ----- | ------- | ------- | -------
61 | \'a b c\' | a | b | c |
62 | \"a b c\" | a | b | c |
63 |
64 | An escaped backslash doesn't escape the subsequent character:
65 |
66 | Input | Token 1 | Token 2
67 | ----- | ------- | -------
68 | a\\\\"b c" d | a\b c | d
69 |
70 | * In sequences begun by a double quote mark, only double quote marks and backslashes (and, for compatibility reasons, dollar signs) may be escaped; the backslash otherwise has no special significance.
71 |
72 | Input | Token 1 | Token 2 | Token 3
73 | ----- | ------- | ------- | -------
74 | "\a \b" \c | \a \b | c |
75 | "\\\\ \\" \\$ a" "\b" c | \ " $ a | \b | c
76 |
77 | * In sequences begun by a single quote mark, a backslash has no special significance.
78 |
79 | Input | Token 1 | Token 2
80 | ----- | ------- | -------
81 | 'a b' c | a b | c
82 | 'a b\\' c | a b\ | c
83 |
84 | Apart from the fact that newline characters can't be escaped, these rules should match the tokenisation rules of the POSIX shell.
85 |
86 | ## Commands
87 |
88 | Commands are sent to the front-end by `ann-benchmarks`. Each command consists of a single line of text; the front-end replies with one or more lines of text. Front-ends can't initiate communication; they can only reply to commands.
89 |
90 | This section specifies these commands, along with the possible responses a front-end might send.
91 |
92 | If a front-end receives a command that it doesn't understand in the current mode (or at all), it should respond with `epbprtv0 fail` and continue processing commands.
93 |
94 | ### Configuration mode
95 |
96 | In configuration mode, front-ends should respond to three different kinds of command:
97 |
98 | #### `VAR VAL` (two tokens)
99 |
100 | Set the value of the algorithm configuration option `VAR` to `VAL`.
101 |
102 | Responses:
103 |
104 | * `epbprtv0 ok`
105 |
106 | The value specified for the algorithm configuration option `VAR` was acceptable, and the option has been set.
107 |
108 | * `epbprtv0 fail`
109 |
110 | The value specified for the algorithm configuration option `VAR` wasn't acceptable. No change has been made to the value of this option.
111 |
112 | #### `frontend VAR VAL` (three tokens)
113 |
114 | Set the value of the front-end configuration option `VAR` to `VAL`. Front-end configuration options may cause the front-end to behave in a manner other than that described in this specification.
115 |
116 | Responses:
117 |
118 | * `epbprtv0 ok`
119 |
120 | The value specified for the front-end configuration option `VAR` was acceptable, and the option has been set.
121 |
122 | * `epbprtv0 fail`
123 |
124 | The value specified for the front-end configuration option `VAR` wasn't acceptable. No change has been made to the value of this option.
125 |
126 | #### *empty line* (zero tokens)
127 |
128 | Finish configuration mode and enter training mode.
129 |
130 | Responses:
131 |
132 | * `epbprtv0 ok`
133 |
134 | Training mode has been entered.
135 |
136 | * `epbprtv0 fail`
137 |
138 | One or more configuration options required by the algorithm weren't specified, and so the query process has terminated.
139 |
140 | ### Training mode
141 |
142 | In training mode, front-ends should respond to two different kinds of command:
143 |
144 | #### `ENTRY` (one token)
145 |
146 | Interpret `ENTRY` as an item of training data.
147 |
148 | Responses:
149 |
150 | * `epbprtv0 ok`
151 |
152 | `ENTRY` was added as the next item of training data. The index values returned in query mode refer to the first item added as `0`, the second as `1`, and so on.
153 |
154 | * `epbprtv0 fail`
155 |
156 | Either `ENTRY` couldn't be interpreted as an item of training data, or the training data wasn't accepted.
157 |
158 | #### *empty line* (zero tokens)
159 |
160 | Finish training mode and enter query mode.
161 |
162 | Responses:
163 |
164 | * `epbprtv0 ok COUNT1 [fail COUNT2]`
165 |
166 | `COUNT1` (potentially zero) entries were successfully interpreted and added to the data structure. (`COUNT2` entries couldn't be interpreted or couldn't be added for other reasons.)
167 |
168 | ### Query mode
169 |
170 | In query mode, front-ends should respond to two different kinds of command:
171 |
172 | #### `ENTRY N` (two tokens)
173 |
174 | Return the indices of at most `N` (greater than or equal to 1) close matches for `ENTRY`.
175 |
176 | Responses:
177 |
178 | * `epbprtv0 ok R`
179 |
180 | `R` (greater than zero and less than or equal to `N`) close matches were found. Each of the next `R` lines, when tokenised, will consist of the token `epbprtv0` followed by a token specifying the index of a close match. (The first line should identify the *closest* close match, and the `R`-th should identify the furthest away.)
181 |
182 | * `epbprtv0 fail`
183 |
184 | No close matches were found.
185 |
186 | #### *empty line* (zero tokens)
187 |
188 | Finish query mode and terminate the front-end.
189 |
190 | Responses:
191 |
192 | * `epbprtv0 ok`
193 |
194 | The front-end has terminated.
195 |
--------------------------------------------------------------------------------
/ann_benchmarks/algorithms/subprocess.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | from os.path import basename
3 | import shlex
4 | from types import MethodType
5 | import psutil
6 | import subprocess
7 | from ann_benchmarks.data import \
8 | bit_unparse_entry, int_unparse_entry, float_unparse_entry
9 | from ann_benchmarks.algorithms.base import BaseANN
10 |
11 |
12 | class SubprocessStoppedError(Exception):
13 | def __init__(self, code):
14 | super(Exception, self).__init__(code)
15 | self.code = code
16 |
17 |
18 | class Subprocess(BaseANN):
19 | def _raw_line(self):
20 | return shlex.split(
21 | self._get_program_handle().stdout.readline().strip())
22 |
23 | def _line(self):
24 | line = self._raw_line()
25 | # print("<- %s" % (" ".join(line)))
26 | while len(line) < 1 or line[0] != "epbprtv0":
27 | line = self._raw_line()
28 | return line[1:]
29 |
30 | @staticmethod
31 | def _quote(token):
32 | return "'" + str(token).replace("'", "'\\'") + "'"
33 |
34 | def _write(self, string):
35 | # print("-> %s" % string)
36 | self._get_program_handle().stdin.write(string + "\n")
37 |
38 | # Called immediately before transitioning from query mode to training mode
39 | def _configuration_hook(self):
40 | pass
41 |
42 | def _get_program_handle(self):
43 | if self._program:
44 | self._program.poll()
45 | if self._program.returncode:
46 | raise SubprocessStoppedError(self._program.returncode)
47 | else:
48 | self._program = subprocess.Popen(
49 | self._args,
50 | bufsize=1, # line buffering
51 | stdin=subprocess.PIPE,
52 | stdout=subprocess.PIPE,
53 | universal_newlines=True)
54 |
55 | for key, value in iter(self._params.items()):
56 | self._write("%s %s" %
57 | (Subprocess._quote(key), Subprocess._quote(value)))
58 | assert self._line()[0] == "ok", """\
59 | assigning value '%s' to option '%s' failed""" % (value, key)
60 | self._configuration_hook()
61 |
62 | self._write("")
63 | assert self._line()[0] == "ok", """\
64 | transitioning to training mode failed"""
65 | return self._program
66 |
67 | def __init__(self, args, encoder, params):
68 | self.name = "Subprocess(program = %s, %s)" % \
69 | (basename(args[0]), str(params))
70 | self._program = None
71 | self._args = args
72 | self._encoder = encoder
73 | self._params = params
74 |
75 | def get_memory_usage(self):
76 | if not self._program:
77 | self._get_program_handle()
78 | return psutil.Process(pid=self._program.pid).memory_info().rss / 1024
79 |
80 | def fit(self, X):
81 | for entry in X:
82 | d = Subprocess._quote(self._encoder(entry))
83 | self._write(d)
84 | assert self._line()[0] == "ok", """\
85 | encoded training point '%s' was rejected""" % d
86 | self._write("")
87 | assert self._line()[0] == "ok", """\
88 | transitioning to query mode failed"""
89 |
90 | def query(self, v, n):
91 | d = Subprocess._quote(self._encoder(v))
92 | self._write("%s %d" % (d, n))
93 | return self._handle_query_response()
94 |
95 | def _handle_query_response(self):
96 | status = self._line()
97 | if status[0] == "ok":
98 | count = int(status[1])
99 | return self._collect_query_response_lines(count)
100 | else:
101 | assert status[0] == "fail", """\
102 | query neither succeeded nor failed"""
103 | return []
104 |
105 | def _collect_query_response_lines(self, count):
106 | results = []
107 | i = 0
108 | while i < count:
109 | line = self._line()
110 | results.append(int(line[0]))
111 | i += 1
112 | return results
113 |
114 | def done(self):
115 | if self._program:
116 | self._program.poll()
117 | if not self._program.returncode:
118 | self._program.terminate()
119 |
120 |
121 | class PreparedSubprocess(Subprocess):
122 | def __init__(self, args, encoder, params):
123 | super(PreparedSubprocess, self).__init__(args, encoder, params)
124 | self._result_count = None
125 |
126 | def _configuration_hook(self):
127 | self._write("frontend prepared-queries 1")
128 | assert self._line()[0] == "ok", """\
129 | enabling prepared queries mode failed"""
130 |
131 | def query(self, v, n):
132 | self.prepare_query(v, n)
133 | self.run_prepared_query()
134 | return self.get_prepared_query_results()
135 |
136 | def prepare_query(self, v, n):
137 | d = Subprocess._quote(self._encoder(v))
138 | self._write("%s %d" % (d, n))
139 | assert self._line()[0] == "ok", """\
140 | preparing the query '%s' failed""" % d
141 |
142 | def run_prepared_query(self):
143 | self._write("query")
144 | status = self._line()
145 | if status[0] == "ok":
146 | self._result_count = int(status[1])
147 | else:
148 | assert status[0] == "fail", """\
149 | query neither succeeded nor failed"""
150 | self._result_count = 0
151 |
152 | def get_prepared_query_results(self):
153 | if self._result_count:
154 | try:
155 | return self._collect_query_response_lines(self._result_count)
156 | finally:
157 | self._result_count = 0
158 | else:
159 | return []
160 |
161 |
162 | class BatchSubprocess(Subprocess):
163 | def __init__(self, args, encoder, params):
164 | super(BatchSubprocess, self).__init__(args, encoder, params)
165 | self._qp_count = None
166 |
167 | def _configuration_hook(self):
168 | self._write("frontend batch-queries 1")
169 | assert self._line()[0] == "ok", """\
170 | enabling batch queries mode failed"""
171 |
172 | def query(self, v, n):
173 | self.prepare_batch_query([v], n)
174 | self.run_batch_query()
175 | return self.get_batch_results()[0]
176 |
177 | def prepare_batch_query(self, X, n):
178 | d = " ".join(map(lambda p: Subprocess._quote(self._encoder(p)), X))
179 | self._qp_count = len(X)
180 | self._write("%s %d" % (d, n))
181 | assert self._line()[0] == "ok", """\
182 | preparing the batch query '%s' failed""" % d
183 |
184 | def run_batch_query(self):
185 | self._write("query")
186 | status = self._line()
187 | assert status[0] == "ok", """\
188 | batch query failed completely"""
189 |
190 | def get_batch_results(self):
191 | results = []
192 | i = 0
193 | while i < self._qp_count:
194 | # print("%d/%d" % (i, self._qp_count))
195 | status = self._line()
196 | if status[0] == "ok":
197 | rc = int(status[1])
198 | results.append(self._collect_query_response_lines(rc))
199 | else:
200 | results.append([])
201 | i += 1
202 | return results
203 |
204 |
205 | def BitSubprocess(args, params):
206 | return Subprocess(args, bit_unparse_entry, params)
207 |
208 |
209 | def BitSubprocessPrepared(args, params):
210 | return PreparedSubprocess(args, bit_unparse_entry, params)
211 |
212 |
213 | def FloatSubprocess(args, params):
214 | return Subprocess(args, float_unparse_entry, params)
215 |
216 |
217 | def FloatSubprocessPrepared(args, params):
218 | return PreparedSubprocess(args, float_unparse_entry, params)
219 |
220 |
221 | def FloatSubprocessBatch(args, params):
222 | return BatchSubprocess(args, float_unparse_entry, params)
223 |
224 |
225 | def IntSubprocess(args, params):
226 | return Subprocess(args, int_unparse_entry, params)
227 |
228 |
229 | def QueryParamWrapper(constructor, args, params):
230 | r = constructor(args, params)
231 |
232 | def _do(self, original=r._configuration_hook):
233 | original()
234 | self._write("frontend query-parameters 1")
235 | assert self._line()[0] == "ok", """\
236 | enabling query parameter support failed"""
237 | r._configuration_hook = MethodType(_do, r)
238 |
239 | def _sqa(self, *args):
240 | self._write("query-params %s set" %
241 | (" ".join(map(Subprocess._quote, args))))
242 | assert self._line()[0] == "ok", """\
243 | reconfiguring query parameters failed"""
244 | print(args)
245 | r.set_query_arguments = MethodType(_sqa, r)
246 | return r
247 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | Benchmarking nearest neighbors
2 | ==============================
3 |
4 | [](https://travis-ci.org/erikbern/ann-benchmarks)
5 |
6 | Doing fast searching of nearest neighbors in high dimensional spaces is an increasingly important problem, but so far there has not been a lot of empirical attempts at comparing approaches in an objective way.
7 |
8 | This project contains some tools to benchmark various implementations of approximate nearest neighbor (ANN) search for different metrics. We have pregenerated datasets (in HDF5) formats and we also have Docker containers for each algorithm. There's a [test suite](https://travis-ci.org/erikbern/ann-benchmarks) that makes sure every algorithm works.
9 |
10 | Evaluated
11 | =========
12 |
13 | * [Annoy](https://github.com/spotify/annoy)
14 | * [FLANN](http://www.cs.ubc.ca/research/flann/)
15 | * [scikit-learn](http://scikit-learn.org/stable/modules/neighbors.html): LSHForest, KDTree, BallTree
16 | * [PANNS](https://github.com/ryanrhymes/panns)
17 | * [NearPy](http://pixelogik.github.io/NearPy/)
18 | * [KGraph](https://github.com/aaalgo/kgraph)
19 | * [NMSLIB (Non-Metric Space Library)](https://github.com/nmslib/nmslib): SWGraph, HNSW, BallTree, MPLSH
20 | * [hnswlib (a part of nmslib project)](https://github.com/nmslib/hnsw)
21 | * [RPForest](https://github.com/lyst/rpforest)
22 | * [FAISS](https://github.com/facebookresearch/faiss.git)
23 | * [DolphinnPy](https://github.com/ipsarros/DolphinnPy)
24 | * [Datasketch](https://github.com/ekzhu/datasketch)
25 | * [PyNNDescent](https://github.com/lmcinnes/pynndescent)
26 | * [MRPT](https://github.com/teemupitkanen/mrpt)
27 | * [NGT](https://github.com/yahoojapan/NGT): ONNG, PANNG
28 | * [SPTAG](https://github.com/microsoft/SPTAG)
29 | * [PUFFINN](https://github.com/puffinn/puffinn)
30 | * [N2](https://github.com/kakao/n2)
31 | * [ScaNN](https://github.com/google-research/google-research/tree/master/scann)
32 | * [Elastiknn](https://github.com/alexklibisz/elastiknn)
33 |
34 | Data sets
35 | =========
36 |
37 | We have a number of precomputed data sets for this. All data sets are pre-split into train/test and come with ground truth data in the form of the top 100 neighbors. We store them in a HDF5 format:
38 |
39 | | Dataset | Dimensions | Train size | Test size | Neighbors | Distance | Download |
40 | | ----------------------------------------------------------------- | ---------: | ---------: | --------: | --------: | --------- | -------------------------------------------------------------------------- |
41 | | [DEEP1B](http://sites.skoltech.ru/compvision/noimi/) | 96 | 9,990,000 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/deep-image-96-angular.hdf5) (3.6GB)
42 | | [Fashion-MNIST](https://github.com/zalandoresearch/fashion-mnist) | 784 | 60,000 | 10,000 | 100 | Euclidean | [HDF5](http://ann-benchmarks.com/fashion-mnist-784-euclidean.hdf5) (217MB) |
43 | | [GIST](http://corpus-texmex.irisa.fr/) | 960 | 1,000,000 | 1,000 | 100 | Euclidean | [HDF5](http://ann-benchmarks.com/gist-960-euclidean.hdf5) (3.6GB) |
44 | | [GloVe](http://nlp.stanford.edu/projects/glove/) | 25 | 1,183,514 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/glove-25-angular.hdf5) (121MB) |
45 | | GloVe | 50 | 1,183,514 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/glove-50-angular.hdf5) (235MB) |
46 | | GloVe | 100 | 1,183,514 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/glove-100-angular.hdf5) (463MB) |
47 | | GloVe | 200 | 1,183,514 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/glove-200-angular.hdf5) (918MB) |
48 | | [Kosarak](http://fimi.uantwerpen.be/data/) | 27983 | 74,962 | 500 | 100 | Jaccard | [HDF5](http://ann-benchmarks.com/kosarak-jaccard.hdf5) (2.0GB) |
49 | | [MNIST](http://yann.lecun.com/exdb/mnist/) | 784 | 60,000 | 10,000 | 100 | Euclidean | [HDF5](http://ann-benchmarks.com/mnist-784-euclidean.hdf5) (217MB) |
50 | | [NYTimes](https://archive.ics.uci.edu/ml/datasets/bag+of+words) | 256 | 290,000 | 10,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/nytimes-256-angular.hdf5) (301MB) |
51 | | [SIFT](https://corpus-texmex.irisa.fr/) | 128 | 1,000,000 | 10,000 | 100 | Euclidean | [HDF5](http://ann-benchmarks.com/sift-128-euclidean.hdf5) (501MB) |
52 | | [Last.fm](https://github.com/erikbern/ann-benchmarks/pull/91) | 65 | 292,385 | 50,000 | 100 | Angular | [HDF5](http://ann-benchmarks.com/lastfm-64-dot.hdf5) (135MB) |
53 |
54 | Results
55 | =======
56 |
57 | These are all as of 2020-07-12, running all benchmarks on a c5.4xlarge machine on AWS with `--parallelism` set to 3:
58 |
59 | glove-100-angular
60 | -----------------
61 |
62 | 
63 |
64 | sift-128-euclidean
65 | ------------------
66 |
67 | 
68 |
69 | fashion-mnist-784-euclidean
70 | ---------------------------
71 |
72 | 
73 |
74 | lastfm-64-dot
75 | ------------------
76 |
77 | 
78 |
79 | nytimes-256-angular
80 | -------------------
81 |
82 | 
83 |
84 | glove-25-angular
85 | ----------------
86 |
87 | 
88 |
89 | Install
90 | =======
91 |
92 | The only prerequisite is Python (tested with 3.6) and Docker.
93 |
94 | 1. Clone the repo.
95 | 2. Run `pip install -r requirements.txt`.
96 | 3. Run `python install.py` to build all the libraries inside Docker containers (this can take a while, like 10-30 minutes).
97 |
98 | Running
99 | =======
100 |
101 | 1. Run `python run.py` (this can take an extremely long time, potentially days)
102 | 2. Run `python plot.py` or `python create_website.py` to plot results.
103 |
104 | You can customize the algorithms and datasets if you want to:
105 |
106 | * Check that `algos.yaml` contains the parameter settings that you want to test
107 | * To run experiments on SIFT, invoke `python run.py --dataset glove-100-angular`. See `python run.py --help` for more information on possible settings. Note that experiments can take a long time.
108 | * To process the results, either use `python plot.py --dataset glove-100-angular` or `python create_website.py`. An example call: `python create_website.py --plottype recall/time --latex --scatter --outputdir website/`.
109 |
110 | Including your algorithm
111 | ========================
112 |
113 | 1. Add your algorithm into `ann_benchmarks/algorithms` by providing a small Python wrapper.
114 | 2. Add a Dockerfile in `install/` for it
115 | 3. Add it to `algos.yaml`
116 | 4. Add it to `.travis.yml`
117 |
118 | Principles
119 | ==========
120 |
121 | * Everyone is welcome to submit pull requests with tweaks and changes to how each library is being used.
122 | * In particular: if you are the author of any of these libraries, and you think the benchmark can be improved, consider making the improvement and submitting a pull request.
123 | * This is meant to be an ongoing project and represent the current state.
124 | * Make everything easy to replicate, including installing and preparing the datasets.
125 | * Try many different values of parameters for each library and ignore the points that are not on the precision-performance frontier.
126 | * High-dimensional datasets with approximately 100-1000 dimensions. This is challenging but also realistic. Not more than 1000 dimensions because those problems should probably be solved by doing dimensionality reduction separately.
127 | * Single queries are used by default. ANN-Benchmarks enforces that only one CPU is saturated during experimentation, i.e., no multi-threading. A batch mode is available that provides all queries to the implementations at once. Add the flag `--batch` to `run.py` and `plot.py` to enable batch mode.
128 | * Avoid extremely costly index building (more than several hours).
129 | * Focus on datasets that fit in RAM. Out of core ANN could be the topic of a later comparison.
130 | * We mainly support CPU-based ANN algorithms. GPU support exists for FAISS, but it has to be compiled with GPU support locally and experiments must be run using the flags `--local --batch`.
131 | * Do proper train/test set of index data and query points.
132 | * Note that set similarity was supported in the past. This might hopefully be added back soon.
133 |
134 |
135 | Authors
136 | =======
137 |
138 | Built by [Erik Bernhardsson](https://erikbern.com) with significant contributions from [Martin Aumüller](http://itu.dk/people/maau/) and [Alexander Faithfull](https://github.com/ale-f).
139 |
140 | Related Publication
141 | ==================
142 |
143 | The following publication details design principles behind the benchmarking framework:
144 |
145 | - M. Aumüller, E. Bernhardsson, A. Faithfull:
146 | [ANN-Benchmarks: A Benchmarking Tool for Approximate Nearest Neighbor Algorithms](https://arxiv.org/abs/1807.05614). Information Systems 2019. DOI: [10.1016/j.is.2019.02.006](https://doi.org/10.1016/j.is.2019.02.006)
147 |
--------------------------------------------------------------------------------
/ann_benchmarks/main.py:
--------------------------------------------------------------------------------
1 | from __future__ import absolute_import
2 | import argparse
3 | import logging
4 | import logging.config
5 |
6 | import docker
7 | import multiprocessing.pool
8 | import os
9 | import psutil
10 | import random
11 | import shutil
12 | import sys
13 | import traceback
14 |
15 | from ann_benchmarks.datasets import get_dataset, DATASETS
16 | from ann_benchmarks.constants import INDEX_DIR
17 | from ann_benchmarks.algorithms.definitions import (get_definitions,
18 | list_algorithms,
19 | algorithm_status,
20 | InstantiationStatus)
21 | from ann_benchmarks.results import get_result_filename
22 | from ann_benchmarks.runner import run, run_docker
23 |
24 |
25 | def positive_int(s):
26 | i = None
27 | try:
28 | i = int(s)
29 | except ValueError:
30 | pass
31 | if not i or i < 1:
32 | raise argparse.ArgumentTypeError("%r is not a positive integer" % s)
33 | return i
34 |
35 |
36 | def run_worker(cpu, args, queue):
37 | while not queue.empty():
38 | definition = queue.get()
39 | if args.local:
40 | run(definition, args.dataset, args.count, args.runs, args.batch)
41 | else:
42 | memory_margin = 500e6 # reserve some extra memory for misc stuff
43 | mem_limit = int((psutil.virtual_memory().available - memory_margin) / args.parallelism)
44 | run_docker(definition, args.dataset, args.count,
45 | args.runs, args.timeout, args.batch, str(cpu), mem_limit)
46 |
47 |
48 | def main():
49 | parser = argparse.ArgumentParser(
50 | formatter_class=argparse.ArgumentDefaultsHelpFormatter)
51 | parser.add_argument(
52 | '--dataset',
53 | metavar='NAME',
54 | help='the dataset to load training points from',
55 | default='glove-100-angular',
56 | choices=DATASETS.keys())
57 | parser.add_argument(
58 | "-k", "--count",
59 | default=10,
60 | type=positive_int,
61 | help="the number of near neighbours to search for")
62 | parser.add_argument(
63 | '--definitions',
64 | metavar='FILE',
65 | help='load algorithm definitions from FILE',
66 | default='algos.yaml')
67 | parser.add_argument(
68 | '--algorithm',
69 | metavar='NAME',
70 | help='run only the named algorithm',
71 | default=None)
72 | parser.add_argument(
73 | '--docker-tag',
74 | metavar='NAME',
75 | help='run only algorithms in a particular docker image',
76 | default=None)
77 | parser.add_argument(
78 | '--list-algorithms',
79 | help='print the names of all known algorithms and exit',
80 | action='store_true')
81 | parser.add_argument(
82 | '--force',
83 | help='re-run algorithms even if their results already exist',
84 | action='store_true')
85 | parser.add_argument(
86 | '--runs',
87 | metavar='COUNT',
88 | type=positive_int,
89 | help='run each algorithm instance %(metavar)s times and use only'
90 | ' the best result',
91 | default=5)
92 | parser.add_argument(
93 | '--timeout',
94 | type=int,
95 | help='Timeout (in seconds) for each individual algorithm run, or -1'
96 | 'if no timeout should be set',
97 | default=2 * 3600)
98 | parser.add_argument(
99 | '--local',
100 | action='store_true',
101 | help='If set, then will run everything locally (inside the same '
102 | 'process) rather than using Docker')
103 | parser.add_argument(
104 | '--batch',
105 | action='store_true',
106 | help='If set, algorithms get all queries at once')
107 | parser.add_argument(
108 | '--max-n-algorithms',
109 | type=int,
110 | help='Max number of algorithms to run (just used for testing)',
111 | default=-1)
112 | parser.add_argument(
113 | '--run-disabled',
114 | help='run algorithms that are disabled in algos.yml',
115 | action='store_true')
116 | parser.add_argument(
117 | '--parallelism',
118 | type=positive_int,
119 | help='Number of Docker containers in parallel',
120 | default=1)
121 |
122 | args = parser.parse_args()
123 | if args.timeout == -1:
124 | args.timeout = None
125 |
126 | if args.list_algorithms:
127 | list_algorithms(args.definitions)
128 | sys.exit(0)
129 |
130 | logging.config.fileConfig("logging.conf")
131 | logger = logging.getLogger("annb")
132 |
133 | # Nmslib specific code
134 | # Remove old indices stored on disk
135 | if os.path.exists(INDEX_DIR):
136 | shutil.rmtree(INDEX_DIR)
137 |
138 | dataset = get_dataset(args.dataset)
139 | dimension = len(dataset['train'][0]) # TODO(erikbern): ugly
140 | point_type = dataset.attrs.get('point_type', 'float')
141 | distance = dataset.attrs['distance']
142 | definitions = get_definitions(
143 | args.definitions, dimension, point_type, distance, args.count)
144 |
145 | # Filter out, from the loaded definitions, all those query argument groups
146 | # that correspond to experiments that have already been run. (This might
147 | # mean removing a definition altogether, so we can't just use a list
148 | # comprehension.)
149 | filtered_definitions = []
150 | for definition in definitions:
151 | query_argument_groups = definition.query_argument_groups
152 | if not query_argument_groups:
153 | query_argument_groups = [[]]
154 | not_yet_run = []
155 | for query_arguments in query_argument_groups:
156 | fn = get_result_filename(args.dataset,
157 | args.count, definition,
158 | query_arguments, args.batch)
159 | if args.force or not os.path.exists(fn):
160 | not_yet_run.append(query_arguments)
161 | if not_yet_run:
162 | if definition.query_argument_groups:
163 | definition = definition._replace(
164 | query_argument_groups=not_yet_run)
165 | filtered_definitions.append(definition)
166 | definitions = filtered_definitions
167 |
168 | random.shuffle(definitions)
169 |
170 | if args.algorithm:
171 | logger.info(f'running only {args.algorithm}')
172 | definitions = [d for d in definitions if d.algorithm == args.algorithm]
173 |
174 | if not args.local:
175 | # See which Docker images we have available
176 | docker_client = docker.from_env()
177 | docker_tags = set()
178 | for image in docker_client.images.list():
179 | for tag in image.tags:
180 | tag = tag.split(':')[0]
181 | docker_tags.add(tag)
182 |
183 | if args.docker_tag:
184 | logger.info(f'running only {args.docker_tag}')
185 | definitions = [
186 | d for d in definitions if d.docker_tag == args.docker_tag]
187 |
188 | if set(d.docker_tag for d in definitions).difference(docker_tags):
189 | logger.info(f'not all docker images available, only: {set(docker_tags)}')
190 | logger.info(f'missing docker images: '
191 | f'{str(set(d.docker_tag for d in definitions).difference(docker_tags))}')
192 | definitions = [
193 | d for d in definitions if d.docker_tag in docker_tags]
194 | else:
195 | def _test(df):
196 | status = algorithm_status(df)
197 | # If the module was loaded but doesn't actually have a constructor
198 | # of the right name, then the definition is broken
199 | if status == InstantiationStatus.NO_CONSTRUCTOR:
200 | raise Exception("%s.%s(%s): error: the module '%s' does not"
201 | " expose the named constructor" % (
202 | df.module, df.constructor,
203 | df.arguments, df.module))
204 |
205 | if status == InstantiationStatus.NO_MODULE:
206 | # If the module couldn't be loaded (presumably because
207 | # of a missing dependency), print a warning and remove
208 | # this definition from the list of things to be run
209 | logging.warning("%s.%s(%s): the module '%s' could not be "
210 | "loaded; skipping" % (df.module, df.constructor,
211 | df.arguments, df.module))
212 | return False
213 | else:
214 | return True
215 | definitions = [d for d in definitions if _test(d)]
216 |
217 | if not args.run_disabled:
218 | if len([d for d in definitions if d.disabled]):
219 | logger.info(f'Not running disabled algorithms {[d for d in definitions if d.disabled]}')
220 | definitions = [d for d in definitions if not d.disabled]
221 |
222 | if args.max_n_algorithms >= 0:
223 | definitions = definitions[:args.max_n_algorithms]
224 |
225 | if len(definitions) == 0:
226 | raise Exception('Nothing to run')
227 | else:
228 | logger.info(f'Order: {definitions}')
229 |
230 | if args.parallelism > multiprocessing.cpu_count() - 1:
231 | raise Exception('Parallelism larger than %d! (CPU count minus one)' % (multiprocessing.cpu_count() - 1))
232 |
233 | # Multiprocessing magic to farm this out to all CPUs
234 | queue = multiprocessing.Queue()
235 | for definition in definitions:
236 | queue.put(definition)
237 | workers = [multiprocessing.Process(target=run_worker, args=(i+1, args, queue))
238 | for i in range(args.parallelism)]
239 | [worker.start() for worker in workers]
240 | [worker.join() for worker in workers]
241 |
242 | # TODO: need to figure out cleanup handling here
243 |
--------------------------------------------------------------------------------
/ann_benchmarks/runner.py:
--------------------------------------------------------------------------------
1 | import argparse
2 | import json
3 | import logging
4 | import os
5 | import threading
6 | import time
7 | import traceback
8 |
9 | import colors
10 | import docker
11 | import numpy
12 | import psutil
13 |
14 | from ann_benchmarks.algorithms.definitions import (Definition,
15 | instantiate_algorithm,
16 | get_algorithm_name)
17 | from ann_benchmarks.datasets import get_dataset, DATASETS
18 | from ann_benchmarks.distance import metrics, dataset_transform
19 | from ann_benchmarks.results import store_results
20 |
21 |
22 | def run_individual_query(algo, X_train, X_test, distance, count, run_count,
23 | batch):
24 | prepared_queries = \
25 | (batch and hasattr(algo, "prepare_batch_query")) or \
26 | ((not batch) and hasattr(algo, "prepare_query"))
27 |
28 | best_search_time = float('inf')
29 | for i in range(run_count):
30 | print('Run %d/%d...' % (i + 1, run_count))
31 | # a bit dumb but can't be a scalar since of Python's scoping rules
32 | n_items_processed = [0]
33 |
34 | def single_query(v):
35 | if prepared_queries:
36 | algo.prepare_query(v, count)
37 | start = time.time()
38 | algo.run_prepared_query()
39 | total = (time.time() - start)
40 | candidates = algo.get_prepared_query_results()
41 | else:
42 | start = time.time()
43 | candidates = algo.query(v, count)
44 | total = (time.time() - start)
45 | candidates = [(int(idx), float(metrics[distance]['distance'](v, X_train[idx]))) # noqa
46 | for idx in candidates]
47 | n_items_processed[0] += 1
48 | if n_items_processed[0] % 1000 == 0:
49 | print('Processed %d/%d queries...' % (n_items_processed[0], len(X_test)))
50 | if len(candidates) > count:
51 | print('warning: algorithm %s returned %d results, but count'
52 | ' is only %d)' % (algo, len(candidates), count))
53 | return (total, candidates)
54 |
55 | def batch_query(X):
56 | if prepared_queries:
57 | algo.prepare_batch_query(X, count)
58 | start = time.time()
59 | algo.run_batch_query()
60 | total = (time.time() - start)
61 | else:
62 | start = time.time()
63 | algo.batch_query(X, count)
64 | total = (time.time() - start)
65 | results = algo.get_batch_results()
66 | candidates = [[(int(idx), float(metrics[distance]['distance'](v, X_train[idx]))) # noqa
67 | for idx in single_results]
68 | for v, single_results in zip(X, results)]
69 | return [(total / float(len(X)), v) for v in candidates]
70 |
71 | if batch:
72 | results = batch_query(X_test)
73 | else:
74 | results = [single_query(x) for x in X_test]
75 |
76 | total_time = sum(time for time, _ in results)
77 | total_candidates = sum(len(candidates) for _, candidates in results)
78 | search_time = total_time / len(X_test)
79 | avg_candidates = total_candidates / len(X_test)
80 | best_search_time = min(best_search_time, search_time)
81 |
82 | verbose = hasattr(algo, "query_verbose")
83 | attrs = {
84 | "batch_mode": batch,
85 | "best_search_time": best_search_time,
86 | "candidates": avg_candidates,
87 | "expect_extra": verbose,
88 | "name": str(algo),
89 | "run_count": run_count,
90 | "distance": distance,
91 | "count": int(count)
92 | }
93 | additional = algo.get_additional()
94 | for k in additional:
95 | attrs[k] = additional[k]
96 | return (attrs, results)
97 |
98 |
99 | def run(definition, dataset, count, run_count, batch):
100 | algo = instantiate_algorithm(definition)
101 | assert not definition.query_argument_groups \
102 | or hasattr(algo, "set_query_arguments"), """\
103 | error: query argument groups have been specified for %s.%s(%s), but the \
104 | algorithm instantiated from it does not implement the set_query_arguments \
105 | function""" % (definition.module, definition.constructor, definition.arguments)
106 |
107 | D = get_dataset(dataset)
108 | X_train = numpy.array(D['train'])
109 | X_test = numpy.array(D['test'])
110 | distance = D.attrs['distance']
111 | print('got a train set of size (%d * %d)' % X_train.shape)
112 | print('got %d queries' % len(X_test))
113 |
114 | X_train = dataset_transform[distance](X_train)
115 | X_test = dataset_transform[distance](X_test)
116 |
117 | try:
118 | prepared_queries = False
119 | if hasattr(algo, "supports_prepared_queries"):
120 | prepared_queries = algo.supports_prepared_queries()
121 |
122 | t0 = time.time()
123 | memory_usage_before = algo.get_memory_usage()
124 | algo.fit(X_train)
125 | build_time = time.time() - t0
126 | index_size = algo.get_memory_usage() - memory_usage_before
127 | print('Built index in', build_time)
128 | print('Index size: ', index_size)
129 |
130 | query_argument_groups = definition.query_argument_groups
131 | # Make sure that algorithms with no query argument groups still get run
132 | # once by providing them with a single, empty, harmless group
133 | if not query_argument_groups:
134 | query_argument_groups = [[]]
135 |
136 | for pos, query_arguments in enumerate(query_argument_groups, 1):
137 | print("Running query argument group %d of %d..." %
138 | (pos, len(query_argument_groups)))
139 | if query_arguments:
140 | algo.set_query_arguments(*query_arguments)
141 | descriptor, results = run_individual_query(
142 | algo, X_train, X_test, distance, count, run_count, batch)
143 | descriptor["build_time"] = build_time
144 | descriptor["index_size"] = index_size
145 | descriptor["algo"] = get_algorithm_name(
146 | definition.algorithm, batch)
147 | descriptor["dataset"] = dataset
148 | store_results(dataset, count, definition,
149 | query_arguments, descriptor, results, batch)
150 | finally:
151 | algo.done()
152 |
153 |
154 | def run_from_cmdline():
155 | parser = argparse.ArgumentParser()
156 | parser.add_argument(
157 | '--dataset',
158 | choices=DATASETS.keys(),
159 | required=True)
160 | parser.add_argument(
161 | '--algorithm',
162 | required=True)
163 | parser.add_argument(
164 | '--module',
165 | required=True)
166 | parser.add_argument(
167 | '--constructor',
168 | required=True)
169 | parser.add_argument(
170 | '--count',
171 | required=True,
172 | type=int)
173 | parser.add_argument(
174 | '--runs',
175 | required=True,
176 | type=int)
177 | parser.add_argument(
178 | '--batch',
179 | action='store_true')
180 | parser.add_argument(
181 | 'build')
182 | parser.add_argument(
183 | 'queries',
184 | nargs='*',
185 | default=[])
186 | args = parser.parse_args()
187 | algo_args = json.loads(args.build)
188 | query_args = [json.loads(q) for q in args.queries]
189 |
190 | definition = Definition(
191 | algorithm=args.algorithm,
192 | docker_tag=None, # not needed
193 | module=args.module,
194 | constructor=args.constructor,
195 | arguments=algo_args,
196 | query_argument_groups=query_args,
197 | disabled=False
198 | )
199 | run(definition, args.dataset, args.count, args.runs, args.batch)
200 |
201 |
202 | def run_docker(definition, dataset, count, runs, timeout, batch, cpu_limit,
203 | mem_limit=None):
204 | cmd = ['--dataset', dataset,
205 | '--algorithm', definition.algorithm,
206 | '--module', definition.module,
207 | '--constructor', definition.constructor,
208 | '--runs', str(runs),
209 | '--count', str(count)]
210 | if batch:
211 | cmd += ['--batch']
212 | cmd.append(json.dumps(definition.arguments))
213 | cmd += [json.dumps(qag) for qag in definition.query_argument_groups]
214 |
215 | client = docker.from_env()
216 | if mem_limit is None:
217 | mem_limit = psutil.virtual_memory().available
218 |
219 | container = client.containers.run(
220 | definition.docker_tag,
221 | cmd,
222 | volumes={
223 | os.path.abspath('ann_benchmarks'):
224 | {'bind': '/home/app/ann_benchmarks', 'mode': 'ro'},
225 | os.path.abspath('data'):
226 | {'bind': '/home/app/data', 'mode': 'ro'},
227 | os.path.abspath('results'):
228 | {'bind': '/home/app/results', 'mode': 'rw'},
229 | },
230 | cpuset_cpus=cpu_limit,
231 | mem_limit=mem_limit,
232 | detach=True)
233 | logger = logging.getLogger(f"annb.{container.short_id}")
234 |
235 | logger.info('Created container %s: CPU limit %s, mem limit %s, timeout %d, command %s' % \
236 | (container.short_id, cpu_limit, mem_limit, timeout, cmd))
237 |
238 | def stream_logs():
239 | for line in container.logs(stream=True):
240 | logger.info(colors.color(line.decode().rstrip(), fg='blue'))
241 |
242 | t = threading.Thread(target=stream_logs, daemon=True)
243 | t.start()
244 |
245 | try:
246 | exit_code = container.wait(timeout=timeout)
247 |
248 | # Exit if exit code
249 | if exit_code not in [0, None]:
250 | logger.error(colors.color(container.logs().decode(), fg='red'))
251 | logger.error('Child process for container %s raised exception %d' % (container.short_id, exit_code))
252 | except:
253 | logger.error('Container.wait for container %s failed with exception' % container.short_id)
254 | traceback.print_exc()
255 | finally:
256 | pass
257 | # container.remove(force=True)
258 |
--------------------------------------------------------------------------------
/create_website.py:
--------------------------------------------------------------------------------
1 | import matplotlib as mpl
2 | mpl.use('Agg') # noqa
3 | import argparse
4 | import os
5 | import json
6 | import pickle
7 | import yaml
8 | import numpy
9 | import hashlib
10 | from jinja2 import Environment, FileSystemLoader
11 |
12 | from ann_benchmarks import results
13 | from ann_benchmarks.algorithms.definitions import get_algorithm_name
14 | from ann_benchmarks.datasets import get_dataset
15 | from ann_benchmarks.plotting.plot_variants import (all_plot_variants
16 | as plot_variants)
17 | from ann_benchmarks.plotting.metrics import all_metrics as metrics
18 | from ann_benchmarks.plotting.utils import (get_plot_label, compute_metrics,
19 | compute_all_metrics,
20 | create_pointset,
21 | create_linestyles)
22 | import plot
23 |
24 | colors = [
25 | "rgba(166,206,227,1)",
26 | "rgba(31,120,180,1)",
27 | "rgba(178,223,138,1)",
28 | "rgba(51,160,44,1)",
29 | "rgba(251,154,153,1)",
30 | "rgba(227,26,28,1)",
31 | "rgba(253,191,111,1)",
32 | "rgba(255,127,0,1)",
33 | "rgba(202,178,214,1)"
34 | ]
35 |
36 | point_styles = {
37 | "o": "circle",
38 | "<": "triangle",
39 | "*": "star",
40 | "x": "cross",
41 | "+": "rect",
42 | }
43 |
44 |
45 | def convert_color(color):
46 | r, g, b, a = color
47 | return "rgba(%(r)d, %(g)d, %(b)d, %(a)d)" % {
48 | "r": r * 255, "g": g * 255, "b": b * 255, "a": a}
49 |
50 |
51 | def convert_linestyle(ls):
52 | new_ls = {}
53 | for algo in ls.keys():
54 | algostyle = ls[algo]
55 | new_ls[algo] = (convert_color(algostyle[0]),
56 | convert_color(algostyle[1]),
57 | algostyle[2], point_styles[algostyle[3]])
58 | return new_ls
59 |
60 |
61 | def get_run_desc(properties):
62 | return "%(dataset)s_%(count)d_%(distance)s" % properties
63 |
64 |
65 | def get_dataset_from_desc(desc):
66 | return desc.split("_")[0]
67 |
68 |
69 | def get_count_from_desc(desc):
70 | return desc.split("_")[1]
71 |
72 |
73 | def get_distance_from_desc(desc):
74 | return desc.split("_")[2]
75 |
76 |
77 | def get_dataset_label(desc):
78 | return "{} (k = {})".format(get_dataset_from_desc(desc),
79 | get_count_from_desc(desc))
80 |
81 |
82 | def directory_path(s):
83 | if not os.path.isdir(s):
84 | raise argparse.ArgumentTypeError("'%s' is not a directory" % s)
85 | return s + "/"
86 |
87 |
88 | def prepare_data(data, xn, yn):
89 | """Change format from (algo, instance, dict) to (algo, instance, x, y)."""
90 | res = []
91 | for algo, algo_name, result in data:
92 | res.append((algo, algo_name, result[xn], result[yn]))
93 | return res
94 |
95 |
96 | parser = argparse.ArgumentParser()
97 | parser.add_argument(
98 | '--plottype',
99 | help='Generate only the plots specified',
100 | nargs='*',
101 | choices=plot_variants.keys(),
102 | default=plot_variants.keys())
103 | parser.add_argument(
104 | '--outputdir',
105 | help='Select output directory',
106 | default='.',
107 | type=directory_path,
108 | action='store')
109 | parser.add_argument(
110 | '--latex',
111 | help='generates latex code for each plot',
112 | action='store_true')
113 | parser.add_argument(
114 | '--scatter',
115 | help='create scatterplot for data',
116 | action='store_true')
117 | parser.add_argument(
118 | '--recompute',
119 | help='Clears the cache and recomputes the metrics',
120 | action='store_true')
121 | args = parser.parse_args()
122 |
123 |
124 | def get_lines(all_data, xn, yn, render_all_points):
125 | """ For each algorithm run on a dataset, obtain its performance
126 | curve coords."""
127 | plot_data = []
128 | for algo in sorted(all_data.keys(), key=lambda x: x.lower()):
129 | xs, ys, ls, axs, ays, als = \
130 | create_pointset(prepare_data(all_data[algo], xn, yn), xn, yn)
131 | if render_all_points:
132 | xs, ys, ls = axs, ays, als
133 | plot_data.append({"name": algo, "coords": zip(xs, ys), "labels": ls,
134 | "scatter": render_all_points})
135 | return plot_data
136 |
137 |
138 | def create_plot(all_data, xn, yn, linestyle, j2_env, additional_label="",
139 | plottype="line"):
140 | xm, ym = (metrics[xn], metrics[yn])
141 | render_all_points = plottype == "bubble"
142 | plot_data = get_lines(all_data, xn, yn, render_all_points)
143 | latex_code = j2_env.get_template("latex.template").\
144 | render(plot_data=plot_data, caption=get_plot_label(xm, ym),
145 | xlabel=xm["description"], ylabel=ym["description"])
146 | plot_data = get_lines(all_data, xn, yn, render_all_points)
147 | button_label = hashlib.sha224((get_plot_label(xm, ym) + additional_label)
148 | .encode("utf-8")).hexdigest()
149 | return j2_env.get_template("chartjs.template").\
150 | render(args=args, latex_code=latex_code, button_label=button_label,
151 | data_points=plot_data,
152 | xlabel=xm["description"], ylabel=ym["description"],
153 | plottype=plottype, plot_label=get_plot_label(xm, ym),
154 | label=additional_label, linestyle=linestyle,
155 | render_all_points=render_all_points)
156 |
157 |
158 | def build_detail_site(data, label_func, j2_env, linestyles, batch=False):
159 | for (name, runs) in data.items():
160 | print("Building '%s'" % name)
161 | all_runs = runs.keys()
162 | label = label_func(name)
163 | data = {"normal": [], "scatter": []}
164 |
165 | for plottype in args.plottype:
166 | xn, yn = plot_variants[plottype]
167 | data["normal"].append(create_plot(
168 | runs, xn, yn, convert_linestyle(linestyles), j2_env))
169 | if args.scatter:
170 | data["scatter"].append(
171 | create_plot(runs, xn, yn, convert_linestyle(linestyles),
172 | j2_env, "Scatterplot ", "bubble"))
173 |
174 | # create png plot for summary page
175 | data_for_plot = {}
176 | for k in runs.keys():
177 | data_for_plot[k] = prepare_data(runs[k], 'k-nn', 'qps')
178 | plot.create_plot(
179 | data_for_plot, False,
180 | False, True, 'k-nn', 'qps',
181 | args.outputdir + get_algorithm_name(name, batch) + ".png",
182 | linestyles, batch)
183 | output_path = "".join([args.outputdir,
184 | get_algorithm_name(name, batch),
185 | ".html"])
186 | with open(output_path, "w") as text_file:
187 | text_file.write(j2_env.get_template("detail_page.html").
188 | render(title=label, plot_data=data,
189 | args=args, batch=batch))
190 |
191 |
192 | def build_index_site(datasets, algorithms, j2_env, file_name):
193 | dataset_data = {'batch': [], 'non-batch': []}
194 | for mode in ['batch', 'non-batch']:
195 | distance_measures = sorted(
196 | set([get_distance_from_desc(e) for e in datasets[mode].keys()]))
197 | sorted_datasets = sorted(
198 | set([get_dataset_from_desc(e) for e in datasets[mode].keys()]))
199 |
200 | for dm in distance_measures:
201 | d = {"name": dm.capitalize(), "entries": []}
202 | for ds in sorted_datasets:
203 | matching_datasets = [e for e in datasets[mode].keys()
204 | if get_dataset_from_desc(e) == ds and # noqa
205 | get_distance_from_desc(e) == dm]
206 | sorted_matches = sorted(
207 | matching_datasets,
208 | key=lambda e: int(get_count_from_desc(e)))
209 | for idd in sorted_matches:
210 | d["entries"].append(
211 | {"name": idd, "desc": get_dataset_label(idd)})
212 | dataset_data[mode].append(d)
213 |
214 | with open(args.outputdir + "index.html", "w") as text_file:
215 | text_file.write(j2_env.get_template("summary.html").
216 | render(title="ANN-Benchmarks",
217 | dataset_with_distances=dataset_data,
218 | algorithms=algorithms,
219 | label_func=get_algorithm_name))
220 |
221 |
222 | def load_all_results():
223 | """Read all result files and compute all metrics"""
224 | all_runs_by_dataset = {'batch': {}, 'non-batch': {}}
225 | all_runs_by_algorithm = {'batch': {}, 'non-batch': {}}
226 | cached_true_dist = []
227 | old_sdn = None
228 | for properties, f in results.load_all_results():
229 | sdn = get_run_desc(properties)
230 | if sdn != old_sdn:
231 | dataset = get_dataset(properties["dataset"])
232 | cached_true_dist = list(dataset["distances"])
233 | old_sdn = sdn
234 | algo = properties["algo"]
235 | ms = compute_all_metrics(
236 | cached_true_dist, f, properties, args.recompute)
237 | algo_ds = get_dataset_label(sdn)
238 | idx = "non-batch"
239 | if properties["batch_mode"]:
240 | idx = "batch"
241 | all_runs_by_algorithm[idx].setdefault(
242 | algo, {}).setdefault(algo_ds, []).append(ms)
243 | all_runs_by_dataset[idx].setdefault(
244 | sdn, {}).setdefault(algo, []).append(ms)
245 |
246 | return (all_runs_by_dataset, all_runs_by_algorithm)
247 |
248 |
249 | j2_env = Environment(loader=FileSystemLoader("./templates/"), trim_blocks=True)
250 | j2_env.globals.update(zip=zip, len=len)
251 | runs_by_ds, runs_by_algo = load_all_results()
252 | dataset_names = [get_dataset_label(x) for x in list(
253 | runs_by_ds['batch'].keys()) + list(runs_by_ds['non-batch'].keys())]
254 | algorithm_names = list(runs_by_algo['batch'].keys(
255 | )) + list(runs_by_algo['non-batch'].keys())
256 |
257 | linestyles = {**create_linestyles(dataset_names),
258 | **create_linestyles(algorithm_names)}
259 |
260 | build_detail_site(
261 | runs_by_ds['non-batch'],
262 | lambda label: get_dataset_label(label), j2_env, linestyles, False)
263 |
264 | build_detail_site(
265 | runs_by_ds['batch'],
266 | lambda label: get_dataset_label(label), j2_env, linestyles, True)
267 |
268 | build_detail_site(
269 | runs_by_algo['non-batch'],
270 | lambda x: x, j2_env, linestyles, False)
271 |
272 | build_detail_site(
273 | runs_by_algo['batch'], lambda x: x, j2_env, linestyles, True)
274 |
275 | build_index_site(runs_by_ds, runs_by_algo, j2_env, "index.html")
276 |
--------------------------------------------------------------------------------
89 |
90 |
91 |