From 3f2ebf4b1c1400eb85550f23e23e7b35d7609f8d Mon Sep 17 00:00:00 2001 From: Matthijs Douze Date: Thu, 25 Feb 2021 11:37:44 -0800 Subject: [PATCH] Add preassigned functions to contrib Summary: Adds the preassigned add and search python wrappers to contrib. Adds the preassigned search for the binary case (was missing before). Also adds a real test for that functionality. Reviewed By: beauby Differential Revision: D26560021 fbshipit-source-id: 330b715a9ed0073cfdadbfbcb1c23b10bed963a5 --- contrib/README.md | 3 + contrib/ivf_tools.py | 91 ++++++++++++++++++++++++++++++ faiss/IndexBinaryIVF.cpp | 18 +++++- faiss/IndexBinaryIVF.h | 8 +++ tests/test_contrib.py | 116 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 233 insertions(+), 3 deletions(-) create mode 100644 contrib/ivf_tools.py diff --git a/contrib/README.md b/contrib/README.md index 117546359..a31f85a47 100644 --- a/contrib/README.md +++ b/contrib/README.md @@ -44,6 +44,9 @@ Tested in `tests/test_contrib_torch.py` (CPU) and `gpu/test/test_contrib_torch_g Functions to inspect C++ objects wrapped by SWIG. Most often this just means reading fields and converting them to the proper python array. +### ivf_tools.py + +A few functions to override the coarse quantizer in IVF, providing additional flexibility for assignment. ### datasets.py diff --git a/contrib/ivf_tools.py b/contrib/ivf_tools.py new file mode 100644 index 000000000..8d1e2be6f --- /dev/null +++ b/contrib/ivf_tools.py @@ -0,0 +1,91 @@ +# Copyright (c) Facebook, Inc. and its affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import numpy as np +import faiss + +def add_preassigned(index_ivf, x, a, ids=None): + """ + Add elements to an IVF index, where the assignment is already computed + """ + n, d = x.shape + assert a.shape == (n, ) + if isinstance(index_ivf, faiss.IndexBinaryIVF): + d *= 8 + assert d == index_ivf.d + if ids is not None: + assert ids.shape == (n, ) + ids = faiss.swig_ptr(ids) + index_ivf.add_core( + n, faiss.swig_ptr(x), ids, faiss.swig_ptr(a) + ) + + +def search_preassigned(index_ivf, xq, k, list_nos, coarse_dis=None): + """ + Perform a search in the IVF index, with predefined lists to search into + """ + n, d = xq.shape + if isinstance(index_ivf, faiss.IndexBinaryIVF): + d *= 8 + dis_type = "int32" + else: + dis_type = "float32" + + assert d == index_ivf.d + assert list_nos.shape == (n, index_ivf.nprobe) + + # the coarse distances are used in IVFPQ with L2 distance and by_residual=True + # otherwise we provide dummy coarse_dis + if coarse_dis is None: + coarse_dis = np.zeros((n, index_ivf.nprobe), dtype=dis_type) + else: + assert coarse_dis.shape == (n, index_ivf.nprobe) + + D = np.empty((n, k), dtype=dis_type) + I = np.empty((n, k), dtype='int64') + + sp = faiss.swig_ptr + index_ivf.search_preassigned( + n, sp(xq), k, + sp(list_nos), sp(coarse_dis), sp(D), sp(I), False) + return D, I + + +def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None): + """ + Perform a range search in the IVF index, with predefined lists to search into + """ + n, d = x.shape + if isinstance(index_ivf, faiss.IndexBinaryIVF): + d *= 8 + dis_type = "int32" + else: + dis_type = "float32" + + # the coarse distances are used in IVFPQ with L2 distance and by_residual=True + # otherwise we provide dummy coarse_dis + if coarse_dis is None: + coarse_dis = np.empty((n, index_ivf.nprobe), dtype=dis_type) + else: + assert coarse_dis.shape == (n, index_ivf.nprobe) + + assert d == index_ivf.d + assert list_nos.shape == (n, index_ivf.nprobe) + + res = faiss.RangeSearchResult(n) + sp = faiss.swig_ptr + + index_ivf.range_search_preassigned( + n, sp(x), radius, + sp(list_nos), sp(coarse_dis), + res + ) + # get pointers and copy them + lims = faiss.rev_swig_ptr(res.lims, n + 1).copy() + num_results = int(lims[-1]) + dist = faiss.rev_swig_ptr(res.distances, num_results).copy() + indices = faiss.rev_swig_ptr(res.labels, num_results).copy() + return lims, dist, indices diff --git a/faiss/IndexBinaryIVF.cpp b/faiss/IndexBinaryIVF.cpp index e91b1735f..9d1258a50 100644 --- a/faiss/IndexBinaryIVF.cpp +++ b/faiss/IndexBinaryIVF.cpp @@ -656,6 +656,19 @@ void IndexBinaryIVF::range_search( t0 = getmillisecs(); invlists->prefetch_lists(idx.get(), n * nprobe); + range_search_preassigned(n, x, radius, idx.get(), coarse_dis.get(), res); + + indexIVF_stats.search_time += getmillisecs() - t0; +} + +void IndexBinaryIVF::range_search_preassigned( + idx_t n, + const uint8_t* x, + int radius, + const idx_t* assign, + const int32_t* centroid_dis, + RangeSearchResult* res) const { + const size_t nprobe = std::min(nlist, this->nprobe); bool store_pairs = false; size_t nlistv = 0, ndis = 0; @@ -671,7 +684,7 @@ void IndexBinaryIVF::range_search( all_pres[omp_get_thread_num()] = &pres; auto scan_list_func = [&](size_t i, size_t ik, RangeQueryResult& qres) { - idx_t key = idx[i * nprobe + ik]; /* select the list */ + idx_t key = assign[i * nprobe + ik]; /* select the list */ if (key < 0) return; FAISS_THROW_IF_NOT_FMT( @@ -688,7 +701,7 @@ void IndexBinaryIVF::range_search( InvertedLists::ScopedCodes scodes(invlists, key); InvertedLists::ScopedIds ids(invlists, key); - scanner->set_list(key, coarse_dis[i * nprobe + ik]); + scanner->set_list(key, assign[i * nprobe + ik]); nlistv++; ndis += list_size; scanner->scan_codes_range( @@ -711,7 +724,6 @@ void IndexBinaryIVF::range_search( indexIVF_stats.nq += n; indexIVF_stats.nlist += nlistv; indexIVF_stats.ndis += ndis; - indexIVF_stats.search_time += getmillisecs() - t0; } IndexBinaryIVF::~IndexBinaryIVF() { diff --git a/faiss/IndexBinaryIVF.h b/faiss/IndexBinaryIVF.h index bf6ce0fb6..1de4c82cb 100644 --- a/faiss/IndexBinaryIVF.h +++ b/faiss/IndexBinaryIVF.h @@ -131,6 +131,14 @@ struct IndexBinaryIVF : IndexBinary { int radius, RangeSearchResult* result) const override; + void range_search_preassigned( + idx_t n, + const uint8_t* x, + int radius, + const idx_t* assign, + const int32_t* centroid_dis, + RangeSearchResult* result) const; + void reconstruct(idx_t key, uint8_t* recons) const override; /** Reconstruct a subset of the indexed vectors. diff --git a/tests/test_contrib.py b/tests/test_contrib.py index 91f84decb..f9b646504 100644 --- a/tests/test_contrib.py +++ b/tests/test_contrib.py @@ -11,6 +11,7 @@ import platform from faiss.contrib import datasets from faiss.contrib import inspect_tools from faiss.contrib import evaluation +from faiss.contrib import ivf_tools from common import get_dataset_2 try: @@ -265,3 +266,118 @@ class TestRangeEval(unittest.TestCase): np.testing.assert_array_almost_equal(ref_precisions, precisions) np.testing.assert_array_almost_equal(ref_recalls, recalls) + + +class TestPreassigned(unittest.TestCase): + + def test_float(self): + ds = datasets.SyntheticDataset(128, 2000, 2000, 200) + + d = ds.d + xt = ds.get_train() + xq = ds.get_queries() + xb = ds.get_database() + + # define alternative quantizer on the 20 first dims of vectors + km = faiss.Kmeans(20, 50) + km.train(xt[:, :20].copy()) + alt_quantizer = km.index + + index = faiss.index_factory(d, "IVF50,PQ16np") + index.by_residual = False + + # (optional) fake coarse quantizer + fake_centroids = np.zeros((index.nlist, index.d), dtype="float32") + index.quantizer.add(fake_centroids) + + # train the PQ part + index.train(xt) + + # add elements xb + a = alt_quantizer.search(xb[:, :20].copy(), 1)[1].ravel() + ivf_tools.add_preassigned(index, xb, a) + + # search elements xq, increase nprobe, check 4 first results w/ groundtruth + prev_inter_perf = 0 + for nprobe in 1, 10, 20: + + index.nprobe = nprobe + a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1] + D, I = ivf_tools.search_preassigned(index, xq, 4, a) + inter_perf = (I == ds.get_groundtruth()[:, :4]).sum() / I.size + self.assertTrue(inter_perf >= prev_inter_perf) + prev_inter_perf = inter_perf + + # test range search + + index.nprobe = 20 + + a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1] + + # just to find a reasonable radius + D, I = ivf_tools.search_preassigned(index, xq, 4, a) + radius = D.max() * 1.01 + + lims, DR, IR = ivf_tools.range_search_preassigned(index, xq, radius, a) + + # with that radius the k-NN results are a subset of the range search results + for q in range(len(xq)): + l0, l1 = lims[q], lims[q + 1] + self.assertTrue(set(I[q]) <= set(IR[l0:l1])) + + def test_binary(self): + ds = datasets.SyntheticDataset(128, 2000, 2000, 200) + + d = ds.d + xt = ds.get_train() + xq = ds.get_queries() + xb = ds.get_database() + + # define alternative quantizer on the 20 first dims of vectors (will be in float) + km = faiss.Kmeans(20, 50) + km.train(xt[:, :20].copy()) + alt_quantizer = km.index + + binarizer = faiss.index_factory(d, "ITQ,LSHt") + binarizer.train(xt) + + xb_bin = binarizer.sa_encode(xb) + xq_bin = binarizer.sa_encode(xq) + + index = faiss.index_binary_factory(d, "BIVF200") + + fake_centroids = np.zeros((index.nlist, index.d // 8), dtype="uint8") + index.quantizer.add(fake_centroids) + index.is_trained = True + + # add elements xb + a = alt_quantizer.search(xb[:, :20].copy(), 1)[1].ravel() + ivf_tools.add_preassigned(index, xb_bin, a) + + # search elements xq, increase nprobe, check 4 first results w/ groundtruth + prev_inter_perf = 0 + for nprobe in 1, 10, 20: + + index.nprobe = nprobe + a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1] + D, I = ivf_tools.search_preassigned(index, xq_bin, 4, a) + inter_perf = (I == ds.get_groundtruth()[:, :4]).sum() / I.size + self.assertTrue(inter_perf >= prev_inter_perf) + prev_inter_perf = inter_perf + + # test range search + + index.nprobe = 20 + + a = alt_quantizer.search(xq[:, :20].copy(), index.nprobe)[1] + + # just to find a reasonable radius + D, I = ivf_tools.search_preassigned(index, xq_bin, 4, a) + radius = int(D.max() + 1) + + lims, DR, IR = ivf_tools.range_search_preassigned(index, xq_bin, radius, a) + + # with that radius the k-NN results are a subset of the range search results + for q in range(len(xq)): + l0, l1 = lims[q], lims[q + 1] + self.assertTrue(set(I[q]) <= set(IR[l0:l1]))