def test_constructor(self): d = 32 ds = datasets.SyntheticDataset(d, 2000, 5000, 200) index = faiss.index_factory(d, f'PQ{d//2}x4np') index.train(ds.get_train()) index.add(ds.get_database()) Dref, Iref = index.search(ds.get_queries(), 10) nq = Iref.shape[0] index2 = faiss.IndexPQFastScan(d, d // 2, 4) index2.train(ds.get_train()) index2.add(ds.get_database()) Dnew, Inew = index2.search(ds.get_queries(), 10) recall_at_1 = (Iref[:, 0] == Inew[:, 0]).sum() / nq self.assertGreater(recall_at_1, 0.99) data = faiss.serialize_index(index2) index3 = faiss.deserialize_index(data) self.assertEqual(index2.implem, index3.implem) D3, I3 = index3.search(ds.get_queries(), 10) np.testing.assert_array_equal(D3, Dnew) np.testing.assert_array_equal(I3, Inew)
def do_test_rounding(self, implem=4, metric=faiss.METRIC_L2): ds = datasets.SyntheticDataset(32, 2000, 5000, 200) index = faiss.index_factory(32, 'PQ16x4', metric) index.train(ds.get_train()) index.add(ds.get_database()) Dref, Iref = index.search(ds.get_queries(), 10) nq = Iref.shape[0] index2 = faiss.IndexPQFastScan(index) # simply repro normal search index2.implem = 2 D2, I2 = index2.search(ds.get_queries(), 10) np.testing.assert_array_equal(I2, Iref) np.testing.assert_array_equal(D2, Dref) # rounded LUT with correction index2.implem = implem D4, I4 = index2.search(ds.get_queries(), 10) # check accuracy of indexes recalls = {} for rank in 1, 10: recalls[rank] = (Iref[:, :1] == I4[:, :rank]).sum() / nq min_r1 = 0.98 if metric == faiss.METRIC_INNER_PRODUCT else 0.99 self.assertGreater(recalls[1], min_r1) self.assertGreater(recalls[10], 0.995) # check accuracy of distances # err3 = ((D3 - D2) ** 2).sum() err4 = ((D4 - D2) ** 2).sum() nf = (D2 ** 2).sum() self.assertLess(err4, nf * 1e-4)
def test_equiv_pq(self): ds = datasets.SyntheticDataset(32, 2000, 200, 4) index = faiss.index_factory(32, "IVF1,PQ16x4np") index.by_residual = False # force coarse quantizer index.quantizer.add(np.zeros((1, 32), dtype='float32')) index.train(ds.get_train()) index.add(ds.get_database()) Dref, Iref = index.search(ds.get_queries(), 4) index_pq = faiss.index_factory(32, "PQ16x4np") index_pq.pq = index.pq index_pq.is_trained = True index_pq.codes = faiss. downcast_InvertedLists( index.invlists).codes.at(0) index_pq.ntotal = index.ntotal Dnew, Inew = index_pq.search(ds.get_queries(), 4) np.testing.assert_array_equal(Iref, Inew) np.testing.assert_array_equal(Dref, Dnew) index_pq2 = faiss.IndexPQFastScan(index_pq) index_pq2.implem = 12 Dref, Iref = index_pq2.search(ds.get_queries(), 4) index2 = faiss.IndexIVFPQFastScan(index) index2.implem = 12 Dnew, Inew = index2.search(ds.get_queries(), 4) np.testing.assert_array_equal(Iref, Inew) np.testing.assert_array_equal(Dref, Dnew)
def do_test_add(self, d, bbs): ds = datasets.SyntheticDataset(d, 2000, 5000, 200) index = faiss.index_factory(d, f'PQ{d//2}x4np') index.train(ds.get_train()) xb = ds.get_database() index.add(xb[:1235]) index2 = faiss.IndexPQFastScan(index, bbs) index2.add(xb[1235:]) new_codes = faiss.AlignedTable_to_array(index2.codes) index.add(xb[1235:]) index3 = faiss.IndexPQFastScan(index, bbs) ref_codes = faiss.AlignedTable_to_array(index3.codes) self.assertEqual(index3.ntotal, index2.ntotal) np.testing.assert_array_equal(ref_codes, new_codes)
def test_pqfastscan(self): ds = datasets.SyntheticDataset(20, 1000, 1000, 0) index = faiss.index_factory(20, 'PQ5x4') index.train(ds.get_train()) index.add(ds.get_database()) recons = index.reconstruct_n(0, index.ntotal) index2 = faiss.IndexPQFastScan(index) recons2 = index2.reconstruct_n(0, index.ntotal) np.testing.assert_array_equal(recons, recons2)
def get_index(self, d, metric): if (d, metric) not in self.cache: ds = datasets.SyntheticDataset(d, 1000, 2000, 200) target_size = d // 2 index = faiss.index_factory(d, 'PQ%dx4' % target_size, metric) index.train(ds.get_train()) index.add(ds.get_database()) index2 = faiss.IndexPQFastScan(index) # uint8 LUT but no SIMD index2.implem = 4 Dref, Iref = index2.search(ds.get_queries(), 10) self.cache[(d, metric)] = (ds, index, Dref, Iref) return self.cache[(d, metric)]
def build_fast_scan_index(self, index, params): qbs, bbs = params index2 = faiss.IndexPQFastScan(index, bbs) index2.qbs = qbs index2.implem = 15 return index2
def build_fast_scan_index(self, index, qbs): index2 = faiss.IndexPQFastScan(index) index2.qbs = qbs index2.implem = 13 return index2
def build_fast_scan_index(self, index, params): index2 = faiss.IndexPQFastScan(index) index2.implem = 5 return index2