def test_IMI_2(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset_2(d, nt, nb, nq) d = xt.shape[1] gt_index = faiss.IndexFlatL2(d) gt_index.add(xb) D, gt_nns = gt_index.search(xq, 1) ############# redo including training nbits = 5 ai0 = faiss.IndexFlatL2(int(d / 2)) ai1 = faiss.IndexFlatL2(int(d / 2)) coarse_quantizer = faiss.MultiIndexQuantizer2(d, nbits, ai0, ai1) index = faiss.IndexIVFPQ(coarse_quantizer, d, (1 << nbits)**2, 8, 8) index.quantizer_trains_alone = 1 index.train(xt) index.add(xb) index.nprobe = 100 D, nns = index.search(xq, 10) n_ok = (nns == gt_nns).sum() # should return the same result self.assertGreater(n_ok, 165)
def test_IMI(self): d = 32 nb = 1000 nt = 1500 nq = 200 (xt, xb, xq) = get_dataset_2(d, nt, nb, nq) d = xt.shape[1] gt_index = faiss.IndexFlatL2(d) gt_index.add(xb) D, gt_nns = gt_index.search(xq, 1) nbits = 5 coarse_quantizer = faiss.MultiIndexQuantizer(d, 2, nbits) index = faiss.IndexIVFPQ(coarse_quantizer, d, (1 << nbits)**2, 8, 8) index.quantizer_trains_alone = 1 index.train(xt) index.add(xb) index.nprobe = 100 D, nns = index.search(xq, 10) n_ok = (nns == gt_nns).sum() # Should return 166 on mac, and 170 on linux. self.assertGreater(n_ok, 165) ############# replace with explicit assignment indexes nbits = 5 pq = coarse_quantizer.pq centroids = faiss.vector_to_array(pq.centroids) centroids = centroids.reshape(pq.M, pq.ksub, pq.dsub) ai0 = faiss.IndexFlatL2(pq.dsub) ai0.add(centroids[0]) ai1 = faiss.IndexFlatL2(pq.dsub) ai1.add(centroids[1]) coarse_quantizer_2 = faiss.MultiIndexQuantizer2(d, nbits, ai0, ai1) coarse_quantizer_2.pq = pq coarse_quantizer_2.is_trained = True index.quantizer = coarse_quantizer_2 index.reset() index.add(xb) D, nns = index.search(xq, 10) n_ok = (nns == gt_nns).sum() # should return the same result self.assertGreater(n_ok, 165)