From 3f3d18d2e7821356a8ad9105a236a6230a720528 Mon Sep 17 00:00:00 2001 From: Junjie Qi Date: Tue, 21 Jan 2025 11:45:07 -0800 Subject: [PATCH] add test to cover GPU (#4130) Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4130 same as title Reviewed By: asadoughi Differential Revision: D68388863 fbshipit-source-id: 4ebb38d8454bf95733c918950d7d8d3b22e00d5d --- tests/test_contrib.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_contrib.py b/tests/test_contrib.py index fb778afb81..ca5d2bcca7 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -35,7 +35,7 @@ class TestComputeGT(unittest.TestCase): - def do_test_compute_GT(self, metric=faiss.METRIC_L2): + def do_test_compute_GT(self, metric=faiss.METRIC_L2, ngpu=0): d = 64 xt, xb, xq = get_dataset_2(d, 0, 10000, 100) @@ -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, ngpu=0) + xq, matrix_iterator(xb, 1000), 10, metric, ngpu=ngpu) np.testing.assert_array_equal(Iref, Inew) # decimal = 4 required when run on GPU @@ -62,6 +62,12 @@ def test_compute_GT(self): def test_compute_GT_ip(self): self.do_test_compute_GT(faiss.METRIC_INNER_PRODUCT) + def test_compute_GT_gpu(self): + self.do_test_compute_GT(ngpu=-1) + + def test_compute_GT_ip_gpu(self): + self.do_test_compute_GT(faiss.METRIC_INNER_PRODUCT, ngpu=-1) + class TestDatasets(unittest.TestCase): """here we test only the synthetic dataset. Datasets that require