def test_indices_ivfsq(self):
res = faiss.StandardGpuResources()
d = 128
nb = 5000
nlist = 10
qtype = faiss.ScalarQuantizer.QT_4bit
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.GpuIndexIVFScalarQuantizerConfig()
idx = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, qtype,
faiss.METRIC_L2, True, 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.GpuIndexIVFScalarQuantizer(res, d, nlist, qtype,
faiss.METRIC_L2, True, 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 make_indices_copy_from_cpu(nlist, d, qtype, by_residual, metric, clamp):
to_train = make_t(10000, d, clamp)
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
qtype, metric, by_residual)
idx_cpu.train(to_train)
idx_cpu.add(to_train)
res = faiss.StandardGpuResources()
res.noTempMemory()
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, idx_cpu)
return idx_cpu, idx_gpu
def test_serialize(self):
res = faiss.StandardGpuResources()
d = 32
k = 10
train = make_t(10000, d)
add = make_t(10000, d)
query = make_t(10, d)
# Construct various GPU index types
indexes = []
# Flat
indexes.append(faiss.GpuIndexFlatL2(res, d))
# IVF
nlist = 5
# IVFFlat
indexes.append(faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2))
# IVFSQ
indexes.append(faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, faiss.ScalarQuantizer.QT_fp16))
# IVFPQ
indexes.append(faiss.GpuIndexIVFPQ(res, d, nlist, 4, 8, faiss.METRIC_L2))
for index in indexes:
index.train(train)
index.add(add)
orig_d, orig_i = index.search(query, k)
ser = faiss.serialize_index(faiss.index_gpu_to_cpu(index))
cpu_index = faiss.deserialize_index(ser)
gpu_index_restore = faiss.index_cpu_to_gpu(res, 0, cpu_index)
restore_d, restore_i = gpu_index_restore.search(query, k)
self.assertTrue(np.array_equal(orig_d, restore_d))
self.assertTrue(np.array_equal(orig_i, restore_i))
# Make sure the index is in a state where we can add to it
# without error
gpu_index_restore.add(query)
