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)