def test_equiv_sh(self): """ make sure that the IVFSpectralHash sa_encode function gives the same result as the concatenated RQ + LSH index sa_encode """ ds = SyntheticDataset(32, 500, 100, 0) index1 = faiss.index_factory(ds.d, "RQ1x4,Refine(ITQ16,LSH)") index1.train(ds.get_train()) # reproduce this in an IndexIVFSpectralHash coarse_quantizer = faiss.IndexFlat(ds.d) rq = faiss.downcast_index(index1.base_index).rq centroids = get_additive_quantizer_codebooks(rq)[0] coarse_quantizer.add(centroids) encoder = faiss.downcast_index(index1.refine_index) # larger than the magnitude of the vectors # negative because otherwise the bits are flipped period = -100000.0 index2 = faiss.IndexIVFSpectralHash(coarse_quantizer, ds.d, coarse_quantizer.ntotal, encoder.sa_code_size() * 8, period) # replace with the vt of the encoder. Binarization is performed by # the IndexIVFSpectralHash itself index2.replace_vt(encoder) codes1 = index1.sa_encode(ds.get_database()) codes2 = index2.sa_encode(ds.get_database()) np.testing.assert_array_equal(codes1, codes2)
def test_precomp(self): ds = datasets.SyntheticDataset(32, 1000, 1000, 0) # make sure it work with varying nb of bits nbits = faiss.UInt64Vector() nbits.push_back(5) nbits.push_back(6) nbits.push_back(7) rq = faiss.ResidualQuantizer(ds.d, nbits) rq.train_type = faiss.ResidualQuantizer.Train_default rq.train(ds.get_train()) codebooks = get_additive_quantizer_codebooks(rq) precomp = precomp_codebooks(codebooks) codebook_cross_prods_ref, cent_norms_ref = precomp # check C++ precomp tables codebook_cross_prods_ref = np.hstack([ np.vstack(c) for c in codebook_cross_prods_ref]) rq.compute_codebook_tables() codebook_cross_prods = faiss.vector_to_array( rq.codebook_cross_products) codebook_cross_prods = codebook_cross_prods.reshape( rq.total_codebook_size, rq.total_codebook_size) cent_norms = faiss.vector_to_array(rq.cent_norms) np.testing.assert_array_almost_equal( codebook_cross_prods, codebook_cross_prods_ref, decimal=5) np.testing.assert_array_almost_equal( np.hstack(cent_norms_ref), cent_norms, decimal=5) # validate that the python tab-based encoding works xb = ds.get_database() ref_codes, _, _ = beam_search_encoding_ref(codebooks, xb, 7) new_codes, _ = beam_search_encoding_tab(codebooks, xb, 7, precomp) np.testing.assert_array_equal(ref_codes, new_codes) # validate the C++ beam_search_encode_step_tab function beam_search_encoding_tab(codebooks, xb, 7, precomp, implem="ref cpp") # check implem w/ residuals n = ref_codes.shape[0] sp = faiss.swig_ptr ref_codes_packed = np.zeros((n, rq.code_size), dtype='uint8') ref_codes_int32 = ref_codes.astype('int32') rq.pack_codes( n, sp(ref_codes_int32), sp(ref_codes_packed), rq.M * ref_codes.shape[1] ) rq.max_beam_size = 7 codes_ref_residuals = rq.compute_codes(xb) np.testing.assert_array_equal(ref_codes_packed, codes_ref_residuals) rq.use_beam_LUT = 1 codes_new = rq.compute_codes(xb) np.testing.assert_array_equal(codes_ref_residuals, codes_new)
def test_equiv_rq(self): """ make sure it is equivalent to search a RQ and to search an IVF with RCQ + RQ with the same codebooks. """ ds = datasets.SyntheticDataset(32, 3000, 1000, 50) # make a flat RQ iflat = faiss.IndexResidualQuantizer(ds.d, 5, 4) iflat.rq.train_type = faiss.ResidualQuantizer.Train_default iflat.train(ds.get_train()) iflat.add(ds.get_database()) # ref search result Dref, Iref = iflat.search(ds.get_queries(), 10) # get its codebooks + encoded version of the dataset codebooks = get_additive_quantizer_codebooks(iflat.rq) codes = faiss.vector_to_array(iflat.codes).reshape(-1, iflat.code_size) # make an IVF with 2x4 + 3x4 = 5x4 bits ivf = faiss.index_factory(ds.d, "IVF256(RCQ2x4),RQ3x4") # initialize the codebooks rcq = faiss.downcast_index(ivf.quantizer) faiss.copy_array_to_vector( np.vstack(codebooks[:rcq.rq.M]).ravel(), rcq.rq.codebooks ) rcq.rq.is_trained = True # translation of AdditiveCoarseQuantizer::train rcq.ntotal = 1 << rcq.rq.tot_bits rcq.centroid_norms.resize(rcq.ntotal) rcq.rq.compute_centroid_norms(rcq.centroid_norms.data()) rcq.is_trained = True faiss.copy_array_to_vector( np.vstack(codebooks[rcq.rq.M:]).ravel(), ivf.rq.codebooks ) ivf.rq.is_trained = True ivf.is_trained = True # add the codes (this works because 2x4 is a multiple of 8 bits) ivf.add_sa_codes(codes) # perform exhaustive search ivf.nprobe = ivf.nlist Dnew, Inew = ivf.search(ds.get_queries(), 10) np.testing.assert_array_equal(Iref, Inew) np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5)
def test_training_with_limited_mem(self): """ make sure a different batch size gives the same result""" ds = datasets.SyntheticDataset(32, 3000, 1000, 0) xt = ds.get_train() rq0 = faiss.ResidualQuantizer(ds.d, 4, 6) rq0.train_type = faiss.ResidualQuantizer.Train_default rq0.max_beam_size = 5 # rq0.verbose = True rq0.train(xt) cb0 = get_additive_quantizer_codebooks(rq0) rq1 = faiss.ResidualQuantizer(ds.d, 4, 6) rq1.train_type = faiss.ResidualQuantizer.Train_default rq1.max_beam_size = 5 rq1.max_mem_distances rq1.max_mem_distances = 3000 * ds.d * 4 * 3 # rq1.verbose = True rq1.train(xt) cb1 = get_additive_quantizer_codebooks(rq1) for c0, c1 in zip(cb0, cb1): self.assertTrue(np.all(c0 == c1))