From 2220dc0d49ee4be76da8440cb5ac91e0ff277503 Mon Sep 17 00:00:00 2001 From: Di-Is Date: Fri, 10 Jan 2025 21:49:05 +0900 Subject: [PATCH 1/2] fix: CPU test does not use GPU --- contrib/exhaustive_search.py | 9 ++++++--- tests/test_contrib.py | 2 +- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/contrib/exhaustive_search.py b/contrib/exhaustive_search.py index 808453e6c1..004217d23e 100644 --- a/contrib/exhaustive_search.py +++ b/contrib/exhaustive_search.py @@ -11,7 +11,7 @@ LOG = logging.getLogger(__name__) -def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2): +def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, ngpu=-1): """Computes the exact KNN search results for a dataset that possibly does not fit in RAM but for which we have an iterator that returns it block by block. @@ -23,8 +23,11 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2): rh = faiss.ResultHeap(nq, k, keep_max=keep_max) index = faiss.IndexFlat(d, metric_type) - if faiss.get_num_gpus(): - LOG.info('running on %d GPUs' % faiss.get_num_gpus()) + if ngpu == -1: + ngpu = faiss.get_num_gpus() + + if ngpu: + LOG.info('running on %d GPUs' % ngpu) index = faiss.index_cpu_to_all_gpus(index) # compute ground-truth by blocks, and add to heaps diff --git a/tests/test_contrib.py b/tests/test_contrib.py index ba185f92b2..fb778afb81 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -50,7 +50,7 @@ def matrix_iterator(xb, bs): yield xb[i0:i0 + bs] Dnew, Inew = knn_ground_truth( - xq, matrix_iterator(xb, 1000), 10, metric) + xq, matrix_iterator(xb, 1000), 10, metric, ngpu=0) np.testing.assert_array_equal(Iref, Inew) # decimal = 4 required when run on GPU From f1a7d45eccd37b70ab3c8507a50fa4090bde0f91 Mon Sep 17 00:00:00 2001 From: Di-Is Date: Tue, 14 Jan 2025 19:43:20 +0900 Subject: [PATCH 2/2] change: align GPU handling in knn_ground_truth to other functions --- contrib/exhaustive_search.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/contrib/exhaustive_search.py b/contrib/exhaustive_search.py index 004217d23e..257427777d 100644 --- a/contrib/exhaustive_search.py +++ b/contrib/exhaustive_search.py @@ -11,7 +11,7 @@ LOG = logging.getLogger(__name__) -def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, ngpu=-1): +def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, shard=False, ngpu=-1): """Computes the exact KNN search results for a dataset that possibly does not fit in RAM but for which we have an iterator that returns it block by block. @@ -28,7 +28,9 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, ngpu=-1): if ngpu: LOG.info('running on %d GPUs' % ngpu) - index = faiss.index_cpu_to_all_gpus(index) + co = faiss.GpuMultipleClonerOptions() + co.shard = shard + index = faiss.index_cpu_to_all_gpus(index, co=co, ngpu=ngpu) # compute ground-truth by blocks, and add to heaps i0 = 0