def test_range_search(self): d = 4 nt = 100 nq = 10 nb = 50 (xt, xb, xq) = get_dataset(d, nb, nt, nq) index = faiss.IndexFlatL2(d) index.add(xb) Dref, Iref = index.search(xq, 5) thresh = 0.1 # *squared* distance lims, D, I = index.range_search(xq, thresh) for i in range(nq): Iline = I[lims[i]:lims[i + 1]] Dline = D[lims[i]:lims[i + 1]] for j, dis in zip(Iref[i], Dref[i]): if dis < thresh: li, = np.where(Iline == j) self.assertTrue(li.size == 1) idx = li[0] self.assertGreaterEqual(1e-4, abs(Dline[idx] - dis))
def test_4variants(self): d = 32 nt = 2500 nq = 400 nb = 5000 (xt, xb, xq) = get_dataset(d, nb, nt, nq) index_gt = faiss.IndexFlatL2(d) index_gt.add(xb) D_ref, I_ref = index_gt.search(xq, 10) nok = {} for qname in "QT_4bit QT_4bit_uniform QT_8bit QT_8bit_uniform QT_fp16".split( ): qtype = getattr(faiss.ScalarQuantizer, qname) index = faiss.IndexScalarQuantizer(d, qtype, faiss.METRIC_L2) index.train(xt) index.add(xb) D, I = index.search(xq, 10) nok[qname] = (I[:, 0] == I_ref[:, 0]).sum() print(nok, nq) self.assertGreaterEqual(nok['QT_8bit'], nq * 0.9) self.assertGreaterEqual(nok['QT_8bit'], nok['QT_4bit']) self.assertGreaterEqual(nok['QT_8bit'], nok['QT_8bit_uniform']) self.assertGreaterEqual(nok['QT_4bit'], nok['QT_4bit_uniform']) self.assertGreaterEqual(nok['QT_fp16'], nok['QT_8bit'])
def test_IndexFlat(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset(d, nb, nt, nq) index = faiss.IndexFlatL2(d) index.add(xb) self.run_search_and_reconstruct(index, xb, xq, eps=0.0)
def test_IndexTransform(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset(d, nb, nt, nq) index = faiss.index_factory(d, "L2norm,PCA8,IVF32,PQ8np") faiss.ParameterSpace().set_index_parameter(index, "nprobe", 4) index.train(xt) index.add(xb) self.run_search_and_reconstruct(index, xb, xq)
def test_MultiIndex(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset(d, nb, nt, nq) index = faiss.index_factory(d, "IMI2x5,PQ8np") faiss.ParameterSpace().set_index_parameter(index, "nprobe", 4) index.train(xt) index.add(xb) self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
def test_IndexIVFPQ(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset(d, nb, nt, nq) quantizer = faiss.IndexFlatL2(d) index = faiss.IndexIVFPQ(quantizer, d, 32, 8, 8) index.cp.min_points_per_centroid = 5 # quiet warning index.nprobe = 4 index.train(xt) index.add(xb) self.run_search_and_reconstruct(index, xb, xq, eps=1.0)
def test_search_k1(self): # verify codepath for k = 1 and k > 1 d = 64 nb = 0 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset(d, nb, nt, nq) miq = faiss.MultiIndexQuantizer(d, 2, 6) miq.train(xt) D1, I1 = miq.search(xq, 1) D5, I5 = miq.search(xq, 5) self.assertEqual(np.abs(I1[:, :1] - I5[:, :1]).max(), 0) self.assertEqual(np.abs(D1[:, :1] - D5[:, :1]).max(), 0)