def test_indices_ivfpq(self): res = faiss.StandardGpuResources() d = 128 nb = 5000 nlist = 10 M = 4 nbits = 8 rs = np.random.RandomState(567) xb = rs.rand(nb, d).astype('float32') xb_indices_base = np.arange(nb, dtype=np.int64) # Force values to not be representable in int32 xb_indices = (xb_indices_base + 4294967296).astype('int64') config = faiss.GpuIndexIVFPQConfig() idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits, faiss.METRIC_L2, config) idx.train(xb) idx.add_with_ids(xb, xb_indices) _, I = idx.search(xb[10:20], 5) self.assertTrue(np.array_equal(xb_indices[10:20], I[:, 0])) # Store values using 32-bit indices instead config.indicesOptions = faiss.INDICES_32_BIT idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits, faiss.METRIC_L2, config) idx.train(xb) idx.add_with_ids(xb, xb_indices) _, I = idx.search(xb[10:20], 5) # This will strip the high bit self.assertTrue(np.array_equal(xb_indices_base[10:20], I[:, 0]))
def test_copy_to_gpu(self): res = faiss.StandardGpuResources() for bits_per_code in [4, 5, 6, 8]: d = 128 nb = 10000 nq = 20 rs = np.random.RandomState(567) xb = rs.rand(nb, d).astype('float32') xq = rs.rand(nq, d).astype('float32') nlist = int(math.sqrt(nb)) sub_q = 16 bits_per_code = 8 nprobe = 4 config = faiss.GpuIndexIVFPQConfig() config.interleavedLayout = True idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config) q = faiss.IndexFlatL2(d) idx_cpu = faiss.IndexIVFPQ(q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2) idx_cpu.train(xb) idx_cpu.add(xb) idx_gpu.copyFrom(idx_cpu) idx_gpu.nprobe = nprobe idx_cpu.nprobe = nprobe # Try without precomputed codes d_g, i_g = idx_gpu.search(xq, 10) d_c, i_c = idx_cpu.search(xq, 10) self.assertGreaterEqual((i_g == i_c).sum(), i_g.size * 0.9) self.assertTrue(np.allclose(d_g, d_c)) # Try with precomputed codes (different kernel) idx_gpu.setPrecomputedCodes(True) d_g, i_g = idx_gpu.search(xq, 10) d_c, i_c = idx_cpu.search(xq, 10) self.assertGreaterEqual((i_g == i_c).sum(), i_g.size * 0.9) self.assertTrue(np.allclose(d_g, d_c))
def test_indices_ivfpq(self): res = faiss.StandardGpuResources() d = 128 nb = 5000 nlist = 10 M = 4 nbits = 8 rs = np.random.RandomState(567) xb = rs.rand(nb, d).astype('float32') xb_indices_base = np.arange(nb, dtype=np.int64) # Force values to not be representable in int32 xb_indices = (xb_indices_base + 4294967296).astype('int64') config = faiss.GpuIndexIVFPQConfig() idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits, faiss.METRIC_L2, config) idx.train(xb) idx.add_with_ids(xb, xb_indices) # invalid k (should be > 0) k = -5 idx.setNumProbes(3) self.assertRaises(AssertionError, idx.search, xb[10:20], k) # invalid nprobe (should be > 0) self.assertRaises(RuntimeError, idx.setNumProbes, 0) self.assertRaises(RuntimeError, idx.setNumProbes, -3) k = 5 idx.nprobe = -3 self.assertRaises(RuntimeError, idx.search, xb[10:20], k) # valid params k = 5 idx.setNumProbes(3) _, I = idx.search(xb[10:20], k) self.assertTrue(np.array_equal(xb_indices[10:20], I[:, 0]))
def test_IndexIVFPQ(self): (xt, xb, xq) = self.get_dataset() d = xt.shape[1] dev_no = 0 usePrecomputed = True res = faiss.StandardGpuResources() flat_config = faiss.GpuIndexFlatConfig() flat_config.device = dev_no gt_index = faiss.GpuIndexFlatL2(res, d, flat_config) gt_index.add(xb) D, gt_nns = gt_index.search(xq, 1) coarse_quantizer = faiss.IndexFlatL2(d) ncentroids = int(np.sqrt(xb.shape[0])) * 4 index = faiss.IndexIVFPQ(coarse_quantizer, d, ncentroids, 32, 8) # add implemented on GPU but not train index.train(xt) ivfpq_config = faiss.GpuIndexIVFPQConfig() ivfpq_config.device = dev_no ivfpq_config.usePrecomputedTables = usePrecomputed gpuIndex = faiss.GpuIndexIVFPQ(res, index, ivfpq_config) gpuIndex.setNumProbes(64) index.add(xb) D, nns = index.search(xq, 10) n_ok = (nns == gt_nns).sum() nq = xq.shape[0] print ncentroids, n_ok, nq self.assertGreater(n_ok, nq * 0.2)