def test_OPQ(self): M = 4 ev = Randu10kUnbalanced() d = ev.d index = faiss.IndexPQ(d, M, 8) res = ev.launch('PQ', index) e_pq = ev.evalres(res) index_pq = faiss.IndexPQ(d, M, 8) opq_matrix = faiss.OPQMatrix(d, M) # opq_matrix.verbose = true opq_matrix.niter = 10 opq_matrix.niter_pq = 4 index = faiss.IndexPreTransform(opq_matrix, index_pq) res = ev.launch('OPQ', index) e_opq = ev.evalres(res) print('e_pq=%s' % e_pq) print('e_opq=%s' % e_opq) # verify that OPQ better than PQ for r in 1, 10, 100: assert(e_opq[r] > e_pq[r])
def test_faiss_to_nanopq(self): D, M, Ks = 32, 4, 256 Nt, Nb, Nq = 2000, 10000, 100 nbits = int(np.log2(Ks)) assert nbits == 8 Xt = np.random.rand(Nt, D).astype(np.float32) Xb = np.random.rand(Nb, D).astype(np.float32) Xq = np.random.rand(Nq, D).astype(np.float32) pq_faiss = faiss.IndexPQ(D, M, nbits) pq_faiss.train(x=Xt) pq_faiss.add(x=Xb) pq_nanopq, Cb_faiss = nanopq.faiss_to_nanopq(pq_faiss=pq_faiss) self.assertEqual(pq_nanopq.codewords.shape, (M, Ks, int(D / M))) # Encoded results should be same Cb_nanopq = pq_nanopq.encode(vecs=Xb) self.assertTrue(np.array_equal(Cb_nanopq, Cb_faiss)) # Search result should be same topk = 100 _, ids1 = pq_faiss.search(x=Xq, k=topk) ids2 = np.array( [ np.argsort(pq_nanopq.dtable(query=xq).adist(codes=Cb_nanopq))[:topk] for xq in Xq ] ) self.assertTrue(np.array_equal(ids1, ids2))
def __init__(self, vector_sz, n_subquantizers=0, n_bits=8): if n_subquantizers > 0: self.index = faiss.IndexPQ(vector_sz, n_subquantizers, n_bits, faiss.METRIC_INNER_PRODUCT) else: self.index = faiss.IndexFlatIP(vector_sz) self.index_id_to_db_id = np.empty((0), dtype=np.int64)
def test_faiss_nanopq_compare_accuracy(self): D, M, Ks = 32, 4, 256 Nt, Nb, Nq = 20000, 10000, 100 nbits = int(np.log2(Ks)) assert nbits == 8 Xt = np.random.rand(Nt, D).astype(np.float32) Xb = np.random.rand(Nb, D).astype(np.float32) Xq = np.random.rand(Nq, D).astype(np.float32) pq_faiss = faiss.IndexPQ(D, M, nbits) pq_faiss.train(x=Xt) Cb_faiss = pq_faiss.pq.compute_codes(Xb) Xb_faiss_ = pq_faiss.pq.decode(Cb_faiss) pq_nanopq = nanopq.PQ(M=M, Ks=Ks) pq_nanopq.fit(vecs=Xt) Cb_nanopq = pq_nanopq.encode(vecs=Xb) Xb_nanopq_ = pq_nanopq.decode(codes=Cb_nanopq) # Reconstruction error should be almost identical avg_relative_error_faiss = ((Xb - Xb_faiss_) ** 2).sum() / (Xb ** 2).sum() avg_relative_error_nanopq = ((Xb - Xb_nanopq_) ** 2).sum() / (Xb ** 2).sum() diff_rel = ( avg_relative_error_faiss - avg_relative_error_nanopq ) / avg_relative_error_faiss diff_rel = np.sqrt(diff_rel ** 2) print("avg_rel_error_faiss:", avg_relative_error_faiss) print("avg_rel_error_nanopq:", avg_relative_error_nanopq) print("diff rel:", diff_rel) self.assertLess(diff_rel, 0.01)
def test_IndexPQ_ip(self): q = faiss.IndexPQ(d, M, nbits_per_index, faiss.METRIC_INNER_PRODUCT) res = ev.launch('FLAT / PQ IP', q) e = ev.evalres(res) # should give 0.070 0.230 0.260 #(same result as regular PQ on normalized distances) assert e[10] > 0.2
def nanopq_to_faiss(pq_nanopq): """Convert a :class:`nanopq.PQ` instance to `faiss.IndexPQ <https://github.com/facebookresearch/faiss/blob/master/IndexPQ.h>`_. To use this function, `faiss module needs to be installed <https://github.com/facebookresearch/faiss/blob/master/INSTALL.md>`_. Args: pq_nanopq (nanopq.PQ): An input PQ instance. Returns: faiss.IndexPQ: A converted PQ instance, with the same codewords to the input. """ assert isinstance(pq_nanopq, PQ), "Error. pq_nanopq must be nanopq.pq" assert pq_nanopq.codewords is not None, "Error. pq_nanopq.codewords must have been set beforehand" D = pq_nanopq.Ds * pq_nanopq.M nbits = {np.uint8: 8, np.uint16: 16, np.uint32: 32}[pq_nanopq.code_dtype] pq_faiss = faiss.IndexPQ(D, pq_nanopq.M, nbits) for m in range(pq_nanopq.M): # Prepare std::vector<float> codewords_cpp_m = faiss.FloatVector() # Flatten m-th codewords from (Ks, Ds) to (Ks * Ds, ), then copy them to cpp faiss.copy_array_to_vector(pq_nanopq.codewords[m].reshape(-1), codewords_cpp_m) # Set the codeword to ProductQuantizer in IndexPQ pq_faiss.pq.set_params(centroids=codewords_cpp_m.data(), m=m) pq_faiss.is_trained = True return pq_faiss
def pq_train(self, x, m, n_bits): d = x.shape[1] # Create the index self.pq = faiss.IndexPQ(d, m, n_bits) # Training self.pq.train(x)
def fit(self, X): print("start to fit, X.shape[1]=%d, summary=%s" % (X.shape[1], self.__str__())) index = faiss.IndexPQ(X.shape[1], self._n_M, self._n_bits, faiss.METRIC_L2) index.train(X) index.add(X) self.index = index
def test_polysemous_OOM(self): """ this used to cause OOM when training polysemous with large nb bits""" d = 32 xt, xb, xq = get_dataset_2(d, 10000, 0, 0) index = faiss.IndexPQ(d, M, 13) index.do_polysemous_training = True index.pq.cp.niter = 0 index.polysemous_training.max_memory = 128 * 1024 * 1024 self.assertRaises(RuntimeError, index.train, xt)
def test_id_remap_idmap(self): # reference: index without remapping index = faiss.IndexPQ(d, 8, 8) k = 10 index.train(xt) index.add(xb) _Dref, Iref = index.search(xq, k) # try a remapping ids = np.arange(nb)[::-1].copy() sub_index = faiss.IndexPQ(d, 8, 8) index2 = faiss.IndexIDMap(sub_index) index2.train(xt) index2.add_with_ids(xb, ids) _D, I = index2.search(xq, k) assert np.all(I == nb - 1 - Iref)
def test_encode_to_binary(self): d = 256 nt = 256 nb = 1500 nq = 500 (xt, xb, xq) = make_binary_dataset(d, nt, nb, nq) pq = faiss.ProductQuantizer(d, int(d / 8), 8) centroids = binary_to_float( np.tile(np.arange(256), int(d / 8)).astype('uint8').reshape(-1, 1)) faiss.copy_array_to_vector(centroids.ravel(), pq.centroids) pq.is_trained = True codes = pq.compute_codes(binary_to_float(xb)) assert np.all(codes == xb) indexpq = faiss.IndexPQ(d, int(d / 8), 8) indexpq.pq = pq indexpq.is_trained = True indexpq.add(binary_to_float(xb)) D, I = indexpq.search(binary_to_float(xq), 3) for i in range(nq): for j, dj in zip(I[i], D[i]): ref_dis = binary_dis(xq[i], xb[j]) assert 4 * ref_dis == dj nlist = 32 quantizer = faiss.IndexFlatL2(d) # pretext class for training iflat = faiss.IndexIVFFlat(quantizer, d, nlist) iflat.train(binary_to_float(xt)) indexivfpq = faiss.IndexIVFPQ(quantizer, d, nlist, int(d / 8), 8) indexivfpq.pq = pq indexivfpq.is_trained = True indexivfpq.by_residual = False indexivfpq.add(binary_to_float(xb)) indexivfpq.nprobe = 4 D, I = indexivfpq.search(binary_to_float(xq), 3) for i in range(nq): for j, dj in zip(I[i], D[i]): ref_dis = binary_dis(xq[i], xb[j]) assert 4 * ref_dis == dj
def test_id_remap_idmap(self): # reference: index without remapping index = faiss.IndexPQ(d, 8, 8) k = 10 index.train(xt) index.add(xb) _Dref, Iref = index.search(xq, k) # try a remapping ids = np.arange(nb)[::-1].copy() sub_index = faiss.IndexPQ(d, 8, 8) index2 = faiss.IndexIDMap(sub_index) index2.train(xt) index2.add_with_ids(xb, ids) # false = do not add 1 to the returned ids (this is done by # default to accommodate lua indexing) _D, I = index2.search(xq, k) assert np.all(I == nb - 1 - Iref)
def test_IndexPQ_refined(self): q = faiss.IndexPQ(d, M, nbits_per_index) res = ev.launch('PQ non-refined', q) e = ev.evalres(res) q.reset() rq = faiss.IndexRefineFlat(q) res = ev.launch('PQ refined', rq) e2 = ev.evalres(res) assert e2[10] >= e[10] rq.k_factor = 4 res = ev.launch('PQ refined*4', rq) e3 = ev.evalres(res) assert e3[10] >= e2[10]
def test_rand_vector(self): """ test if the smooth_vectors function is reasonably compressible with a small PQ """ x = faiss.rand_smooth_vectors(1300, 32) xt = x[:1000] xb = x[1000:1200] xq = x[1200:] _, gt = faiss.knn(xq, xb, 10) index = faiss.IndexPQ(32, 4, 4) index.train(xt) index.add(xb) D, I = index.search(xq, 10) ninter = faiss.eval_intersection(I, gt) # 445 for SyntheticDataset self.assertGreater(ninter, 420) self.assertLess(ninter, 460)
def qnnd(features_train, features_test, k, m, c): # approximated L2 distance with product quantization # FAISS library # input: features_train, features_test # output: anomaly score and test runtime #k = 1 d = features_train.shape[1] # building index index = faiss.IndexPQ(d, m, c) index.train(features_train) index.add(features_train) # searching nearest neighbour start = time.time() Dq, Iq = index.search(features_test, k) test_runtime = time.time() - start anomalyScores = Dq[:, 0] return anomalyScores, test_runtime
def main(): xd = fvecs_read("D:\data\gist\gist_base.fvecs") xq = fvecs_read("D:\data\gist\gist_query.fvecs") xt = fvecs_read("D:\data\gist\gist_learn.fvecs") index = faiss.IndexFlatL2(xd.shape[1]) index.add(sanitize(xd)) pqidx = faiss.IndexPQ(xd.shape[1], int(xd.shape[1] / 2), 8) pqidx.train(sanitize(xt)) centroids = faiss.vector_to_array(pqidx.pq.centroids) print(centroids.shape) centroids = centroids.reshape(int(xd.shape[1] / 2), 256, 2) DEFAULT_write("gist_vector.bin", xd) DEFAULT_write("gist_query.bin", xq) codebooks_write("gist_codebook.bin", centroids)
def test_polysemous(self): index = faiss.IndexPQ(d, M, nbits_per_index) index.do_polysemous_training = True # reduce nb iterations to speed up training for the test index.polysemous_training.n_iter = 50000 index.polysemous_training.n_redo = 1 res = ev.launch('normal PQ', index) e_baseline = ev.evalres(res) index.search_type = faiss.IndexPQ.ST_polysemous index.polysemous_ht = int(M / 16. * 58) stats = faiss.cvar.indexPQ_stats stats.reset() res = ev.launch('Polysemous ht=%d' % index.polysemous_ht, index) e_polysemous = ev.evalres(res) print(e_baseline, e_polysemous, index.polysemous_ht) print(stats.n_hamming_pass, stats.ncode) # The randu dataset is difficult, so we are not too picky on # the results. Here we assert that we have < 10 % loss when # computing full PQ on fewer than 20% of the data. assert stats.n_hamming_pass < stats.ncode / 5
def test_IndexPQ(self): q = faiss.IndexPQ(d, M, nbits_per_index) res = ev.launch('FLAT / PQ L2', q) e = ev.evalres(res) # should give 0.070 0.230 0.260 assert e[10] > 0.2
def __init__(self, d, nlist=100): self.index = faiss.IndexPQ(d, 16, 8) self.index.nprobe = 10 self.index2 = faiss.IndexIDMap(self.index)
with open(path) as f_in, open(os.path.join(args.output, 'docid'), 'w') as f_out: for line in f_in: info = json.loads(line) docid = info['id'] vector = info['vector'] f_out.write(f'{docid}\n') vectors.append(vector) vectors = np.array(vectors, dtype='float32') print(vectors.shape) if args.hnsw and args.pq: index = faiss.IndexHNSWPQ(args.dim, args.pq_m, args.M) index.hnsw.efConstruction = args.efC index.metric_type = faiss.METRIC_INNER_PRODUCT elif args.hnsw: index = faiss.IndexHNSWFlat(args.dim, args.M, faiss.METRIC_INNER_PRODUCT) index.hnsw.efConstruction = args.efC elif args.pq: index = faiss.IndexPQ(args.dim, args.M, 8, faiss.METRIC_INNER_PRODUCT) else: index = faiss.IndexFlatIP(args.dim) index.verbose = True if args.pq: index.train(vectors) index.add(vectors) faiss.write_index(index, os.path.join(args.output, 'index'))
codes = rs.randint(1 << 31, size=nb).astype('uint64') print("code_size=%d nv=%d %.2f bits" % (codec.code_size, codec.nv, np.log2(codec.nv))) assert codec.code_size == 8 dis = np.empty((nq, k), dtype='float32') labels = np.empty((nq, k), dtype='int64') t0 = time.time() codec.find_nn(codes, xq) t1 = time.time() print("time for code_size=%d nq=%d nb=%d: %.3f s (%.3f ms/query)" % (codec.code_size, nq, nb, t1 - t0, (t1 - t0) * 1000 / nq)) index = faiss.IndexPQ(dim, 8, 8) xb = rs.randn(nb, dim).astype('float32') print("train") index.train(xb) print("add") index.add(xb) t0 = time.time() index.search(xq, 1) t1 = time.time() print("time for IndexPQ code_size=%d nq=%d nb=%d: %.3f s (%.3f ms/query)" % (index.pq.code_size, nq, nb, t1 - t0, (t1 - t0) * 1000 / nq))
#!/usr/bin/env python2 from __future__ import print_function import time import numpy as np import faiss from datasets import load_sift1M, evaluate print("load data") xb, xq, xt, gt = load_sift1M() nq, d = xq.shape # index with 16 subquantizers, 8 bit each index = faiss.IndexPQ(d, 16, 8) index.do_polysemous_training = True index.verbose = True print("train") index.train(xt) print("add vectors to index") index.add(xb) nt = 1 faiss.omp_set_num_threads(1) print("PQ baseline", end=' ')
for line in f_in: info = json.loads(line) docid = info['id'] vector = info['vector'] f_out.write(f'{docid}\n') vectors.append(vector) vectors = np.array(vectors, dtype='float32') print(vectors.shape) if args.hnsw and args.pq: index = faiss.IndexHNSWPQ(args.dim, args.pq_m, args.M) index.hnsw.efConstruction = args.efC index.metric_type = faiss.METRIC_INNER_PRODUCT elif args.hnsw: index = faiss.IndexHNSWFlat(args.dim, args.M, faiss.METRIC_INNER_PRODUCT) index.hnsw.efConstruction = args.efC elif args.pq: index = faiss.IndexPQ(args.dim, args.pq_m, args.pq_nbits, faiss.METRIC_INNER_PRODUCT) else: index = faiss.IndexFlatIP(args.dim) index.verbose = True if args.pq: index.train(vectors) index.add(vectors) print(index.ntotal) faiss.write_index(index, os.path.join(args.output, 'index'))
def PQCpu(config): print("PQCpu, ", config) d = config['dimension'] # dimension nb = config['db_size'] # database size nq = config['query_num'] # nb of queries topk = config['top_k'] m = config['sub_quantizers'] nbits = config['bits_per_code'] search_repeat = 10 index_list = [] create_ave_duration = 0 search_ave_duration = 0 if config['test_batch_write'] == True: batch_write_ave_duration = 0 batch_write_num = config['write_batch_num'] batch_write_time = int(nb / config['write_batch_num']) print("batch_write_time = ", batch_write_num) for i in range(config['db_num']): index_pq = faiss.IndexPQ(d, m, nbits) batch_write_ave_one_lib = 0 np.random.seed(66666) xb = np.random.random((10000, d)).astype('float32') xb[:, 0] += np.arange(10000) / 1000. if index_pq.is_trained == False: index_pq.train(xb) for j in range(batch_write_time): np.random.seed(i * batch_write_time + j) xb = np.random.random((batch_write_num, d)).astype('float32') xb[:, 0] += np.arange(batch_write_num) / 1000. begin_time = time.time() index_pq.add(xb) duration = time.time() - begin_time batch_write_ave_one_lib += duration batch_write_ave_duration += duration print("batch_write_ave_one_lib = ", (batch_write_ave_one_lib / batch_write_time) * 1000 * 1000, " us") index_list.append(index_pq) print("batch_write_ave_duration = ", (batch_write_ave_duration / len(index_list) / batch_write_time) * 1000 * 1000, " us") return index_list for i in range(config['db_num']): np.random.seed(i) # make reproducible xb = np.random.random((nb, d)).astype('float32') xb[:, 0] += np.arange(nb) / 1000. begin_time = time.time() index_pq = faiss.IndexPQ(d, m, nbits) assert not index_pq.is_trained index_pq.train(xb) assert index_pq.is_trained index_pq.add(xb) # add vectors to the index duration = time.time() - begin_time create_ave_duration += duration index_list.append(index_pq) print("craete ave duration = ", create_ave_duration / len(index_list), " s") if len(index_list) == 0: return index_list for i in range(len(index_list)): for j in range(search_repeat): np.random.seed(i * search_repeat + j + config['db_num']) xq = np.random.random((nq, d)).astype('float32') xq[:, 0] += np.arange(nq) / 1000. begin_time = time.time() index_list[i].search(xq, topk) # actual search duration = time.time() - begin_time search_ave_duration += duration print("search index aver time = ", search_ave_duration / len(index_list) / search_repeat, " s") return index_list
# Copyright (c) Facebook, Inc. and its affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. from __future__ import print_function import faiss from datasets import load_sift1M, evaluate xb, xq, xt, gt = load_sift1M() nq, d = xq.shape k = 32 for nbits in 4, 6, 8, 10, 12: index = faiss.IndexPQ(d, 8, nbits) index.train(xt) index.add(xb) t, r = evaluate(index, xq, gt, k) print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1])) del index
from recall_data import recall_data # 基本参数 d = 300 # 向量维数 data_size = 10000 # 数据库大小 k = 50 M = 30 nbits = 8 # 生成测试数据 numpy.random.seed(13) data = numpy.random.random(size=(data_size, d)).astype('float32') test_data = recall_data # 创建索引模型并添加向量 index = faiss.IndexPQ(d, M, nbits) # 利用点积作为索引 # 训练数据 start_time = time.time() assert not index.is_trained index.train(data) assert index.is_trained print "Train Index Used %.2f sec." % (time.time() - start_time) # print(index.is_trained) # 该索引是否训练过 # print(index.ntotal) # 索引容量 start_time = time.time() index.add(data) # 将数据添加进索引 print "Add vector Used %.2f sec." % (time.time() - start_time) start_time = time.time()
def evaluate(dataset, poly, width): xt = fvecs_read(dataset+"/learn.fvecs") xb = fvecs_read(dataset+"/base.fvecs") xq = fvecs_read(dataset+"/query.fvecs") #xt = xt[:32000] #xb = xb[:32000] #xq = xb[:320] nq, d = xq.shape #print xq.shape gt = ivecs_read(dataset+"/groundtruth.ivecs") # index with 16 subquantizers, 8 bit each index = faiss.IndexPQ(d, width, 8) if(poly): index.search_type = faiss.IndexPQ.ST_polysemous index.do_polysemous_training = True else: index.do_polysemous_training = False index.search_type = faiss.IndexPQ.ST_PQ index.verbose=False index.train(xt) index.add(xb) nt = 1 faiss.omp_set_num_threads(1) if(poly and width == 8): htList =[0,17,19,21,23,25] elif(poly and width == 16): htList =[0,43,45,47,49,51,53,55,57,59] elif(poly and width == 32): htList =[0,78,81,84,87,90,93,96,99,102,105,108,111,114,117] else: htList = [0] for ht in htList: index.polysemous_ht = ht t0 = time.time() D, I = index.search(xq, 100) t1 = time.time() recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / float(nq) recall_at_100 = (I[:, :100] == gt[:, :1]).sum() / float(nq) if(poly): print "Poly ", else: print "PQ ", print width,"x8 tau=", ht," ", print "\t %7.3f ms per query, R@1 %.4f R@100 %.4f" % ( (t1 - t0) * 1000.0 / nq * nt, recall_at_1, recall_at_100)
logging.info('d = %d, k = %d', d, args.k) stn = time.time() if args.norm: feats_index = feats_index / np.linalg.norm( feats_index, axis=1, keepdims=True) feats_query = feats_query / np.linalg.norm( feats_query, axis=1, keepdims=True) logging.info('Feats normalization took %.2f', (time.time() - stn)) nbits = 2**3 if args.index_choice == 'flat': index = faiss.IndexFlatL2(d) elif args.index_choice == 'pq': index = faiss.IndexPQ(d, args.M, nbits) stt = time.time() index.train(feats_index) logging.info('Train took %.2f', (time.time() - stt)) sti = time.time() index.add(feats_index) logging.info('Index took %.2f', (time.time() - sti)) sts = time.time() D, I = index.search( feats_query[:args.limit] if args.limit > 0 else feats_query, args.k) logging.info('Search took %.2f', (time.time() - sts)) results = np.empty_like(I, dtype='S8') for i in range(I.shape[0]):