| # Copyright (c) Meta Platforms, Inc. and affiliates. |
| # |
| # This source code is licensed under the MIT license found in the |
| # LICENSE file in the root directory of this source tree. |
| |
| """ test byte codecs """ |
| |
| from __future__ import print_function |
| import numpy as np |
| import unittest |
| import faiss |
| |
| from common_faiss_tests import get_dataset_2 |
| from faiss.contrib.datasets import SyntheticDataset |
| from faiss.contrib.inspect_tools import get_additive_quantizer_codebooks |
| |
| |
| class TestEncodeDecode(unittest.TestCase): |
| |
| def do_encode_twice(self, factory_key): |
| d = 96 |
| nb = 1000 |
| nq = 0 |
| nt = 2000 |
| |
| xt, x, _ = get_dataset_2(d, nt, nb, nq) |
| |
| assert x.size > 0 |
| |
| codec = faiss.index_factory(d, factory_key) |
| |
| codec.train(xt) |
| |
| codes = codec.sa_encode(x) |
| x2 = codec.sa_decode(codes) |
| |
| codes2 = codec.sa_encode(x2) |
| |
| if 'IVF' in factory_key or 'RQ' in factory_key: |
| # some rows are not reconstructed exactly because they |
| # flip into another quantization cell |
| nrowdiff = (codes != codes2).any(axis=1).sum() |
| self.assertTrue(nrowdiff < 10) |
| else: |
| self.assertTrue(np.all(codes == codes2)) |
| |
| x3 = codec.sa_decode(codes2) |
| |
| if 'IVF' in factory_key or 'RQ' in factory_key: |
| diffs = np.abs(x2 - x3).sum(axis=1) |
| avg = np.abs(x2).sum(axis=1).mean() |
| diffs.sort() |
| assert diffs[-10] < avg * 1e-5 |
| else: |
| self.assertTrue(np.allclose(x2, x3)) |
| |
| def test_SQ8(self): |
| self.do_encode_twice('SQ8') |
| |
| def test_IVFSQ8(self): |
| self.do_encode_twice('IVF256,SQ8') |
| |
| def test_PCAIVFSQ8(self): |
| self.do_encode_twice('PCAR32,IVF256,SQ8') |
| |
| def test_PQ6x8(self): |
| self.do_encode_twice('PQ6np') |
| |
| def test_PQ6x6(self): |
| self.do_encode_twice('PQ6x6np') |
| |
| def test_IVFPQ6x8np(self): |
| self.do_encode_twice('IVF512,PQ6np') |
| |
| def test_LSH(self): |
| self.do_encode_twice('LSHrt') |
| |
| def test_RQ6x8(self): |
| self.do_encode_twice('RQ6x8') |
| |
| |
| class TestIndexEquiv(unittest.TestCase): |
| |
| def do_test(self, key1, key2): |
| d = 96 |
| nb = 1000 |
| nq = 0 |
| nt = 2000 |
| |
| xt, x, _ = get_dataset_2(d, nt, nb, nq) |
| |
| codec_ref = faiss.index_factory(d, key1) |
| codec_ref.train(xt) |
| |
| code_ref = codec_ref.sa_encode(x) |
| x_recons_ref = codec_ref.sa_decode(code_ref) |
| |
| codec_new = faiss.index_factory(d, key2) |
| codec_new.pq = codec_ref.pq |
| |
| # replace quantizer, avoiding mem leak |
| oldq = codec_new.q1.quantizer |
| oldq.this.own() |
| codec_new.q1.own_fields = False |
| codec_new.q1.quantizer = codec_ref.quantizer |
| codec_new.is_trained = True |
| |
| code_new = codec_new.sa_encode(x) |
| x_recons_new = codec_new.sa_decode(code_new) |
| |
| self.assertTrue(np.all(code_new == code_ref)) |
| self.assertTrue(np.all(x_recons_new == x_recons_ref)) |
| |
| codec_new_2 = faiss.deserialize_index( |
| faiss.serialize_index(codec_new)) |
| |
| code_new = codec_new_2.sa_encode(x) |
| x_recons_new = codec_new_2.sa_decode(code_new) |
| |
| self.assertTrue(np.all(code_new == code_ref)) |
| self.assertTrue(np.all(x_recons_new == x_recons_ref)) |
| |
| def test_IVFPQ(self): |
| self.do_test("IVF512,PQ6np", "Residual512,PQ6") |
| |
| def test_IMI(self): |
| self.do_test("IMI2x5,PQ6np", "Residual2x5,PQ6") |
| |
| |
| class TestAccuracy(unittest.TestCase): |
| """ comparative accuracy of a few types of indexes """ |
| |
| def compare_accuracy(self, lowac, highac, max_errs=(1e10, 1e10)): |
| d = 96 |
| nb = 1000 |
| nq = 0 |
| nt = 2000 |
| |
| xt, x, _ = get_dataset_2(d, nt, nb, nq) |
| |
| errs = [] |
| |
| for factory_string in lowac, highac: |
| |
| codec = faiss.index_factory(d, factory_string) |
| print('sa codec: code size %d' % codec.sa_code_size()) |
| codec.train(xt) |
| |
| codes = codec.sa_encode(x) |
| x2 = codec.sa_decode(codes) |
| |
| err = ((x - x2) ** 2).sum() |
| errs.append(err) |
| |
| self.assertGreater(errs[0], errs[1]) |
| |
| self.assertGreater(max_errs[0], errs[0]) |
| self.assertGreater(max_errs[1], errs[1]) |
| |
| # just a small IndexLattice I/O test |
| if 'Lattice' in highac: |
| codec2 = faiss.deserialize_index( |
| faiss.serialize_index(codec)) |
| codes = codec2.sa_encode(x) |
| x3 = codec2.sa_decode(codes) |
| self.assertTrue(np.all(x2 == x3)) |
| |
| def test_SQ(self): |
| self.compare_accuracy('SQ4', 'SQ8') |
| |
| def test_SQ2(self): |
| self.compare_accuracy('SQ6', 'SQ8') |
| |
| def test_SQ3(self): |
| self.compare_accuracy('SQ8', 'SQfp16') |
| |
| def test_SQ4(self): |
| self.compare_accuracy('SQ8', 'SQbf16') |
| |
| def test_PQ(self): |
| self.compare_accuracy('PQ6x8np', 'PQ8x8np') |
| |
| def test_PQ2(self): |
| self.compare_accuracy('PQ8x6np', 'PQ8x8np') |
| |
| def test_IVFvsPQ(self): |
| self.compare_accuracy('PQ8np', 'IVF256,PQ8np') |
| |
| def test_Lattice(self): |
| # measured low/high: 20946.244, 5277.483 |
| self.compare_accuracy('ZnLattice3x10_4', |
| 'ZnLattice3x20_4', |
| (22000, 5400)) |
| |
| def test_Lattice2(self): |
| # here the difference is actually tiny |
| # measured errs: [16403.072, 15967.735] |
| self.compare_accuracy('ZnLattice3x12_1', |
| 'ZnLattice3x12_7', |
| (18000, 16000)) |
| |
| |
| swig_ptr = faiss.swig_ptr |
| |
| |
| class LatticeTest(unittest.TestCase): |
| """ Low-level lattice tests """ |
| |
| def test_repeats(self): |
| rs = np.random.RandomState(123) |
| dim = 32 |
| for _i in range(1000): |
| vec = np.floor((rs.rand(dim) ** 7) * 3).astype('float32') |
| vecs = vec.copy() |
| vecs.sort() |
| repeats = faiss.Repeats(dim, swig_ptr(vecs)) |
| code = repeats.encode(swig_ptr(vec)) |
| vec2 = np.zeros(dim, dtype='float32') |
| repeats.decode(code, swig_ptr(vec2)) |
| assert np.all(vec == vec2) |
| |
| def test_ZnSphereCodec_encode_centroid(self): |
| dim = 8 |
| r2 = 5 |
| ref_codec = faiss.ZnSphereCodec(dim, r2) |
| codec = faiss.ZnSphereCodecRec(dim, r2) |
| assert ref_codec.nv == codec.nv |
| s = set() |
| for i in range(ref_codec.nv): |
| c = np.zeros(dim, dtype='float32') |
| ref_codec.decode(i, swig_ptr(c)) |
| code = codec.encode_centroid(swig_ptr(c)) |
| assert 0 <= code < codec.nv |
| s.add(code) |
| assert len(s) == codec.nv |
| |
| def test_ZnSphereCodecRec(self): |
| dim = 16 |
| r2 = 6 |
| codec = faiss.ZnSphereCodecRec(dim, r2) |
| for i in range(codec.nv): |
| c = np.zeros(dim, dtype='float32') |
| codec.decode(i, swig_ptr(c)) |
| code = codec.encode_centroid(swig_ptr(c)) |
| assert code == i |
| |
| def run_ZnSphereCodecAlt(self, dim, r2): |
| # dim = 32 |
| # r2 = 14 |
| codec = faiss.ZnSphereCodecAlt(dim, r2) |
| rs = np.random.RandomState(123) |
| n = 100 |
| codes = rs.randint(codec.nv, size=n, dtype='uint64') |
| x = np.empty((n, dim), dtype='float32') |
| codec.decode_multi(n, swig_ptr(codes), swig_ptr(x)) |
| codes2 = np.empty(n, dtype='uint64') |
| codec.encode_multi(n, swig_ptr(x), swig_ptr(codes2)) |
| |
| assert np.all(codes == codes2) |
| |
| def test_ZnSphereCodecAlt32(self): |
| self.run_ZnSphereCodecAlt(32, 14) |
| |
| def test_ZnSphereCodecAlt24(self): |
| self.run_ZnSphereCodecAlt(24, 14) |
| |
| def test_lattice_index(self): |
| index = faiss.index_factory(96, "ZnLattice3x10_4") |
| rs = np.random.RandomState(123) |
| xq = rs.randn(10, 96).astype('float32') |
| xb = rs.randn(20, 96).astype('float32') |
| index.train(xb) |
| index.add(xb) |
| D, I = index.search(xq, 5) |
| for i in range(10): |
| recons = index.reconstruct_batch(I[i, :]) |
| ref_dis = ((recons - xq[i]) ** 2).sum(1) |
| np.testing.assert_allclose(D[i, :], ref_dis, atol=1e-4) |
| |
| |
| class TestBitstring(unittest.TestCase): |
| |
| def test_rw(self): |
| """ Low-level bit string tests """ |
| rs = np.random.RandomState(1234) |
| nbyte = 1000 |
| sz = 0 |
| |
| bs = np.ones(nbyte, dtype='uint8') |
| bw = faiss.BitstringWriter(swig_ptr(bs), nbyte) |
| |
| if False: |
| ctrl = [(7, 0x35), (13, 0x1d74)] |
| for nbit, x in ctrl: |
| bw.write(x, nbit) |
| else: |
| ctrl = [] |
| while True: |
| nbit = int(1 + 62 * rs.rand() ** 4) |
| if sz + nbit > nbyte * 8: |
| break |
| x = int(rs.randint(1 << nbit, dtype='int64')) |
| bw.write(x, nbit) |
| ctrl.append((nbit, x)) |
| sz += nbit |
| |
| bignum = 0 |
| sz = 0 |
| for nbit, x in ctrl: |
| bignum |= x << sz |
| sz += nbit |
| |
| for i in range(nbyte): |
| self.assertTrue(((bignum >> (i * 8)) & 255) == bs[i]) |
| |
| br = faiss.BitstringReader(swig_ptr(bs), nbyte) |
| |
| for nbit, xref in ctrl: |
| xnew = br.read(nbit) |
| self.assertTrue(xnew == xref) |
| |
| def test_arrays(self): |
| nbit = 5 |
| M = 10 |
| n = 20 |
| rs = np.random.RandomState(123) |
| a = rs.randint(1<<nbit, size=(n, M), dtype='int32') |
| b = faiss.pack_bitstrings(a, nbit) |
| c = faiss.unpack_bitstrings(b, M, nbit) |
| np.testing.assert_array_equal(a, c) |
| |
| def test_arrays_variable_size(self): |
| nbits = [10, 5, 3, 12, 6, 7, 4] |
| n = 20 |
| rs = np.random.RandomState(123) |
| a = rs.randint(1<<16, size=(n, len(nbits)), dtype='int32') |
| a &= (1 << np.array(nbits)) - 1 |
| b = faiss.pack_bitstrings(a, nbits) |
| c = faiss.unpack_bitstrings(b, nbits) |
| np.testing.assert_array_equal(a, c) |
| |
| |
| class TestIVFTransfer(unittest.TestCase): |
| |
| def test_transfer(self): |
| |
| ds = SyntheticDataset(32, 2000, 200, 100) |
| index = faiss.index_factory(ds.d, "IVF20,SQ8") |
| index.train(ds.get_train()) |
| index.add(ds.get_database()) |
| Dref, Iref = index.search(ds.get_queries(), 10) |
| index.reset() |
| |
| codes = index.sa_encode(ds.get_database()) |
| index.add_sa_codes(codes) |
| |
| Dnew, Inew = index.search(ds.get_queries(), 10) |
| |
| np.testing.assert_array_equal(Iref, Inew) |
| np.testing.assert_array_equal(Dref, Dnew) |
| |
| |
| class TestIDMap(unittest.TestCase): |
| def test_idmap(self): |
| ds = SyntheticDataset(32, 2000, 200, 100) |
| ids = np.random.randint(10000, size=ds.nb, dtype='int64') |
| index = faiss.index_factory(ds.d, "IDMap2,PQ8x2") |
| index.train(ds.get_train()) |
| index.add_with_ids(ds.get_database(), ids) |
| Dref, Iref = index.search(ds.get_queries(), 10) |
| |
| index.reset() |
| |
| index.train(ds.get_train()) |
| codes = index.index.sa_encode(ds.get_database()) |
| index.add_sa_codes(codes, ids) |
| Dnew, Inew = index.search(ds.get_queries(), 10) |
| |
| np.testing.assert_array_equal(Iref, Inew) |
| np.testing.assert_array_equal(Dref, Dnew) |
| |
| |
| |
| class TestRefine(unittest.TestCase): |
| |
| def test_refine(self): |
| """ Make sure that IndexRefine can function as a standalone codec """ |
| |
| ds = SyntheticDataset(32, 500, 100, 0) |
| index = faiss.index_factory(ds.d, "RQ2x5,Refine(ITQ,LSHt)") |
| index.train(ds.get_train()) |
| |
| index1 = index.base_index |
| index2 = index.refine_index |
| |
| codes12 = index.sa_encode(ds.get_database()) |
| codes1 = index1.sa_encode(ds.get_database()) |
| codes2 = index2.sa_encode(ds.get_database()) |
| |
| np.testing.assert_array_equal( |
| codes12, |
| np.hstack((codes1, codes2)) |
| ) |
| |
| def test_equiv_rcq_rq(self): |
| """ make sure that the codes generated by the standalone codec are the same |
| between an |
| IndexRefine with ResidualQuantizer |
| and |
| IVF with ResidualCoarseQuantizer |
| both are the centroid id concatenated with the code. |
| """ |
| ds = SyntheticDataset(16, 400, 100, 0) |
| index1 = faiss.index_factory(ds.d, "RQ2x3,Refine(Flat)") |
| index1.train(ds.get_train()) |
| irq = faiss.downcast_index(index1.base_index) |
| # because the default beam factor for RCQ is 4 |
| irq.rq.max_beam_size = 4 |
| |
| index2 = faiss.index_factory(ds.d, "IVF64(RCQ2x3),Flat") |
| index2.train(ds.get_train()) |
| quantizer = faiss.downcast_index(index2.quantizer) |
| quantizer.rq = irq.rq |
| index2.is_trained = True |
| |
| codes1 = index1.sa_encode(ds.get_database()) |
| codes2 = index2.sa_encode(ds.get_database()) |
| |
| np.testing.assert_array_equal(codes1, codes2) |
| |
| 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) |
