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)
