add knn search

Author: Alexsandruss

cuml/knn_clsf.py

@@ -12,6 +12,9 @@
 parser = argparse.ArgumentParser(
     description='cuML kNN classifier benchmark')
 
+parser.add_argument('--task', default='classification', type=str,
+                    choices=('search', 'classification'),
+                    help='kNN task: search or classification')
 parser.add_argument('--n-neighbors', default=5, type=int,
                     help='Number of neighbors to use')
 parser.add_argument('--weights', type=str, default='uniform',
@@ -33,21 +36,31 @@
                                 metric=params.metric)
 
 # Measure time and accuracy on fitting
-train_time, _ = measure_function_time(knn_clsf.fit, X_train, y_train,
-    params=params)
-y_pred = knn_clsf.predict(X_train)
-train_acc = 100 * accuracy_score(y_pred, y_train)
+train_time, _ = measure_function_time(knn_clsf.fit, X_train, y_train, params=params)
+if args.task == 'classification':
+    y_pred = knn_clsf.predict(X_train)
+    train_acc = 100 * accuracy_score(y_pred, y_train)
 
 # Measure time and accuracy on prediction
-predict_time, yp = measure_function_time(knn_clsf.predict, X_test, params=params)
-test_acc = 100 * accuracy_score(yp, y_test)
+if args.task == 'classification':
+    predict_time, yp = measure_function_time(knn_clsf.predict, X_test, params=params)
+    test_acc = 100 * accuracy_score(yp, y_test)
+else:
+    predict_time, _ = measure_function_time(knn_clsf.kneighbors, X_test, params=params)
 
 columns = ('batch', 'arch', 'prefix', 'function', 'threads', 'dtype', 'size',
            'n_neighbors', 'n_classes', 'time')
 
-print_output(library='cuml', algorithm='knn_classification',
-             stages=['training', 'prediction'], columns=columns, params=params,
-             functions=['knn_clsf.fit', 'knn_clsf.predict'],
-             times=[train_time, predict_time],
-             accuracies=[train_acc, test_acc], accuracy_type='accuracy[%]',
-             data=[X_train, X_test], alg_instance=knn_clsf)
+if args.task == 'classification':
+    print_output(library='cuml', algorithm='knn_classification',
+                 stages=['training', 'prediction'], columns=columns, params=params,
+                 functions=['knn_clsf.fit', 'knn_clsf.predict'],
+                 times=[train_time, predict_time],
+                 accuracies=[train_acc, test_acc], accuracy_type='accuracy[%]',
+                 data=[X_train, X_test], alg_instance=knn_clsf)
+else:
+    print_output(library='cuml', algorithm='knn_search',
+                 stages=['training', 'search'], columns=columns, params=params,
+                 functions=['knn_clsf.fit', 'knn_clsf.kneighbors'],
+                 times=[train_time, predict_time],
+                 data=[X_train, X_test], alg_instance=knn_clsf)

runner.py

@@ -288,8 +288,8 @@ class GenerationArgs:
                     if args.output_format == 'json':
                         try:
                             json_result['results'].extend(json.loads(stdout))
-                        except json.JSONDecodeError:
-                            pass
+                        except json.JSONDecodeError as decoding_exception:
+                            stderr += str(decoding_exception) + '\n'
                     elif args.output_format == 'csv':
                         csv_result += stdout + '\n'
                     if stderr != '':

sklearn/knn_clsf.py

@@ -12,6 +12,9 @@
 parser = argparse.ArgumentParser(
     description='scikit-learn kNN classifier benchmark')
 
+parser.add_argument('--task', default='classification', type=str,
+                    choices=('search', 'classification'),
+                    help='kNN task: search or classification')
 parser.add_argument('--n-neighbors', default=5, type=int,
                     help='Number of neighbors to use')
 parser.add_argument('--weights', type=str, default='uniform',
@@ -34,21 +37,31 @@
                                 metric=params.metric)
 
 # Measure time and accuracy on fitting
-train_time, _ = measure_function_time(knn_clsf.fit, X_train, y_train,
-    params=params)
-y_pred = knn_clsf.predict(X_train)
-train_acc = 100 * accuracy_score(y_pred, y_train)
+train_time, _ = measure_function_time(knn_clsf.fit, X_train, y_train, params=params)
+if args.task == 'classification':
+    y_pred = knn_clsf.predict(X_train)
+    train_acc = 100 * accuracy_score(y_pred, y_train)
 
 # Measure time and accuracy on prediction
-predict_time, yp = measure_function_time(knn_clsf.predict, X_test, params=params)
-test_acc = 100 * accuracy_score(yp, y_test)
+if args.task == 'classification':
+    predict_time, yp = measure_function_time(knn_clsf.predict, X_test, params=params)
+    test_acc = 100 * accuracy_score(yp, y_test)
+else:
+    predict_time, _ = measure_function_time(knn_clsf.kneighbors, X_test, params=params)
 
 columns = ('batch', 'arch', 'prefix', 'function', 'threads', 'dtype', 'size',
            'n_neighbors', 'n_classes', 'time')
 
-print_output(library='sklearn', algorithm='knn_classification',
-             stages=['training', 'prediction'], columns=columns, params=params,
-             functions=['knn_clsf.fit', 'knn_clsf.predict'],
-             times=[train_time, predict_time],
-             accuracies=[train_acc, test_acc], accuracy_type='accuracy[%]',
-             data=[X_train, X_test], alg_instance=knn_clsf)
+if args.task == 'classification':
+    print_output(library='sklearn', algorithm='knn_classification',
+                 stages=['training', 'prediction'], columns=columns, params=params,
+                 functions=['knn_clsf.fit', 'knn_clsf.predict'],
+                 times=[train_time, predict_time],
+                 accuracies=[train_acc, test_acc], accuracy_type='accuracy[%]',
+                 data=[X_train, X_test], alg_instance=knn_clsf)
+else:
+    print_output(library='sklearn', algorithm='knn_search',
+                 stages=['training', 'search'], columns=columns, params=params,
+                 functions=['knn_clsf.fit', 'knn_clsf.kneighbors'],
+                 times=[train_time, predict_time],
+                 data=[X_train, X_test], alg_instance=knn_clsf)