make calling conventions uniform between faiss.knn and faiss.knn_gpu

Summary: The order of xb an xq was different between `faiss.knn` and `faiss.knn_gpu`. Also the metric argument was called distance_type. This diff fixes both. Hopefully not too much external code depends on it.

Reviewed By: wickedfoo

Differential Revision: D26222853

fbshipit-source-id: b43e143d64d9ecbbdf541734895c13847cf2696c

contrib/torch_utils.py

@@ -3,6 +3,24 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+"""
+
+This is a set of function wrappers that override the default numpy versions.
+
+Interoperability functions for pytorch and Faiss: Importing this will allow
+pytorch Tensors (CPU or GPU) to be used as arguments to Faiss indexes and
+other functions. Torch GPU tensors can only be used with Faiss GPU indexes.
+If this is imported with a package that supports Faiss GPU, the necessary
+stream synchronization with the current pytorch stream will be automatically
+performed.
+
+Numpy ndarrays can continue to be used in the Faiss python interface after
+importing this file. All arguments must be uniformly either numpy ndarrays
+or Torch tensors; no mixing is allowed.
+
+"""
+
+
 import faiss
 import torch
 import contextlib
@@ -472,10 +490,10 @@ for symbol in dir(faiss_module):
             handle_torch_Index(the_class)
 
 # allows torch tensor usage with bfKnn
-def torch_replacement_knn_gpu(res, xb, xq, k, D=None, I=None, metric=faiss.METRIC_L2):
+def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2):
     if type(xb) is np.ndarray:
         # Forward to faiss __init__.py base method
-        return faiss.knn_gpu_numpy(res, xb, xq, k, D, I, metric)
+        return faiss.knn_gpu_numpy(res, xq, xb, k, D, I, metric)
 
     nb, d = xb.size()
     if xb.is_contiguous():

faiss/gpu/test/test_contrib.py

@@ -40,7 +40,7 @@ class TestComputeGT(unittest.TestCase):
         ds = datasets.SyntheticDataset(32, 0, 1000, 10)
         xq = ds.get_queries()
         xb = ds.get_database()
-        D, I = faiss.knn(xq, xb, 10, distance_type=metric)
+        D, I = faiss.knn(xq, xb, 10, metric=metric)
         threshold = float(D[:, -1].mean())
 
         index = faiss.IndexFlat(32, metric)

faiss/gpu/test/torch_test_contrib_gpu.py

@@ -252,7 +252,7 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
                     else:
                         xb_c = xb_np
 
-                    D, I = faiss.knn_gpu(res, xb_c, xq_c, k)
+                    D, I = faiss.knn_gpu(res, xq_c, xb_c, k)
 
                     self.assertTrue(torch.equal(torch.from_numpy(I), gt_I))
                     self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4)
@@ -278,7 +278,7 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
                             xb_c = to_column_major_torch(xb)
                             assert not xb_c.is_contiguous()
 
-                        D, I = faiss.knn_gpu(res, xb_c, xq_c, k)
+                        D, I = faiss.knn_gpu(res, xq_c, xb_c, k)
 
                         self.assertTrue(torch.equal(I.cpu(), gt_I))
                         self.assertLess((D.cpu() - gt_D).abs().max(), 1e-4)
@@ -286,7 +286,7 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
                         # test on subset
                         try:
                             # This internally uses the current pytorch stream
-                            D, I = faiss.knn_gpu(res, xb_c, xq_c[6:8], k)
+                            D, I = faiss.knn_gpu(res, xq_c[6:8], xb_c, k)
                         except TypeError:
                             if not xq_row_major:
                                 # then it is expected
@@ -320,7 +320,7 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
         D = torch.zeros(nq, k, device=xb_c.device, dtype=torch.float32)
         I = torch.zeros(nq, k, device=xb_c.device, dtype=torch.int32)
 
-        faiss.knn_gpu(res, xb_c, xq_c, k, D, I)
+        faiss.knn_gpu(res, xq_c, xb_c, k, D, I)
 
         self.assertTrue(torch.equal(I.long().cpu(), gt_I))
         self.assertLess((D.float().cpu() - gt_D).abs().max(), 1.5e-3)
@@ -332,7 +332,7 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
         xb_c = xb.half().numpy()
         xq_c = xq.half().numpy()
 
-        faiss.knn_gpu(res, xb_c, xq_c, k, D, I)
+        faiss.knn_gpu(res, xq_c, xb_c, k, D, I)
 
         self.assertTrue(torch.equal(torch.from_numpy(I).long(), gt_I))
         self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1.5e-3)

faiss/python/__init__.py

@@ -548,7 +548,7 @@ def index_cpu_to_gpus_list(index, co=None, gpus=None, ngpu=-1):
     return index_gpu
 
 # allows numpy ndarray usage with bfKnn
-def knn_gpu(res, xb, xq, k, D=None, I=None, metric=METRIC_L2):
+def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2):
     nq, d = xq.shape
     if xq.flags.c_contiguous:
         xq_row_major = True
@@ -903,7 +903,7 @@ def range_search_with_parameters(index, x, radius, params=None, output_stats=Fal
 # KNN function
 ######################################################
 
-def knn(xq, xb, k, distance_type=METRIC_L2):
+def knn(xq, xb, k, metric=METRIC_L2):
     """ wrapper around the faiss knn functions without index """
     nq, d = xq.shape
     nb, d2 = xb.shape
@@ -912,7 +912,7 @@ def knn(xq, xb, k, distance_type=METRIC_L2):
     I = np.empty((nq, k), dtype='int64')
     D = np.empty((nq, k), dtype='float32')
 
-    if distance_type == METRIC_L2:
+    if metric == METRIC_L2:
         heaps = float_maxheap_array_t()
         heaps.k = k
         heaps.nh = nq
@@ -922,7 +922,7 @@ def knn(xq, xb, k, distance_type=METRIC_L2):
             swig_ptr(xq), swig_ptr(xb),
             d, nq, nb, heaps
         )
-    elif distance_type == METRIC_INNER_PRODUCT:
+    elif metric == METRIC_INNER_PRODUCT:
         heaps = float_minheap_array_t()
         heaps.k = k
         heaps.nh = nq
@@ -932,6 +932,8 @@ def knn(xq, xb, k, distance_type=METRIC_L2):
             swig_ptr(xq), swig_ptr(xb),
             d, nq, nb, heaps
         )
+    else:
+        raise NotImplementedError("only L2 and INNER_PRODUCT are supported")
     return D, I
 
 

tests/test_contrib.py

@@ -101,7 +101,7 @@ class TestExhaustiveSearch(unittest.TestCase):
         index.add(xb)
         Dref, Iref = index.search(xq, 10)
 
-        Dnew, Inew = knn(xq, xb, 10, distance_type=faiss.METRIC_INNER_PRODUCT)
+        Dnew, Inew = knn(xq, xb, 10, metric=faiss.METRIC_INNER_PRODUCT)
 
         assert np.all(Inew == Iref)
         assert np.allclose(Dref, Dnew)
@@ -110,7 +110,7 @@ class TestExhaustiveSearch(unittest.TestCase):
         ds = datasets.SyntheticDataset(32, 0, 1000, 10)
         xq = ds.get_queries()
         xb = ds.get_database()
-        D, I = faiss.knn(xq, xb, 10, distance_type=metric)
+        D, I = faiss.knn(xq, xb, 10, metric=metric)
         threshold = float(D[:, -1].mean())
 
         index = faiss.IndexFlat(32, metric)
@@ -136,7 +136,7 @@ class TestExhaustiveSearch(unittest.TestCase):
         ds = datasets.SyntheticDataset(32, 0, 1000, 1000)
         xq = ds.get_queries()
         xb = ds.get_database()
-        D, I = faiss.knn(xq, xb, 10, distance_type=metric)
+        D, I = faiss.knn(xq, xb, 10, metric=metric)
         threshold = float(D[:, -1].mean())
         print(threshold)