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
