Moved the query loop of the runner.run function into a new function
This separates training from querying, which might be useful in the event
that support is added for reusing trained data structures. Cough
@@ -17,25 +17,7 @@ from ann_benchmarks.distance import metrics
from ann_benchmarks.results import store_results
-def run(definition, dataset, count, run_count=3, force_single=False, use_batch_query=False):
- algo = instantiate_algorithm(definition)
-
- D = get_dataset(dataset)
- X_train = numpy.array(D['train'])
- X_test = numpy.array(D['test'])
- distance = D.attrs['distance']
- print('got a train set of size (%d * %d)' % X_train.shape)
- print('got %d queries' % len(X_test))
-
- try:
- t0 = time.time()
- index_size_before = algo.get_index_size("self")
- algo.fit(X_train)
- build_time = time.time() - t0
- index_size = algo.get_index_size("self") - index_size_before
- print('Built index in', build_time)
- print('Index size: ', index_size)
-
+def run_individual_query(algo, X_train, X_test, distance, count, run_count=3, force_single=False, use_batch_query=False):
best_search_time = float('inf')
for i in range(run_count):
print('Run %d/%d...' % (i+1, run_count))
@@ -80,16 +62,40 @@ def run(definition, dataset, count, run_count=3, force_single=False, use_batch_q
verbose = hasattr(algo, "query_verbose")
attrs = {
"batch_mode": use_batch_query,
- "build_time": build_time,
"best_search_time": best_search_time,
"candidates": avg_candidates,
"expect_extra": verbose,
- "index_size": index_size,
"name": algo.name,
"run_count": run_count,
"run_alone": force_single,
}
- store_results(dataset, count, definition, attrs, results)
+ return (attrs, results)
+
+
+def run(definition, dataset, count, run_count=3, force_single=False, use_batch_query=False):
+ algo = instantiate_algorithm(definition)
+
+ D = get_dataset(dataset)
+ X_train = numpy.array(D['train'])
+ X_test = numpy.array(D['test'])
+ distance = D.attrs['distance']
+ print('got a train set of size (%d * %d)' % X_train.shape)
+ print('got %d queries' % len(X_test))
+
+ try:
+ t0 = time.time()
+ index_size_before = algo.get_index_size("self")
+ algo.fit(X_train)
+ build_time = time.time() - t0
+ index_size = algo.get_index_size("self") - index_size_before
+ print('Built index in', build_time)
+ print('Index size: ', index_size)
+
+ descriptor, results = run_individual_query(algo, X_train, X_test,
+ distance, count, run_count, force_single, use_batch_query)
+ descriptor["build_time"] = build_time
+ descriptor["index_size"] = index_size
+ store_results(dataset, count, definition, descriptor, results)
finally:
algo.done()
