mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Gergely Szilvasy
f4338ade1b
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5134 Expand DD test coverage by applying for_all_simd_levels to existing test classes that exercise DD-dispatched code paths but were previously pinned to a single SIMD level. This is the "mega decorator" diff from the DD test coverage gaps plan -- pure decorator additions, no new test logic. Follow-up diffs add new test files and numerical cross-level assertions for gaps the decorator alone cannot close. Classes decorated (grouped by area): * Binary Hamming non-IVF: TestRange and TestKnn in test_binary_hashindex.py; TestBinarySearchParams in test_binary_search_params.py; TestIndexBinaryFromFloat in test_index_binary_from_float.py; TestSpectralHash in test_index_accuracy.py. * IVFPQ / search: EvalIVFPQAccuracy in test_index.py; TestSelector and TestSearchParams in test_search_params.py. * Flat / refine / Panorama: TestIndexFlat, TestIndexFlatL2, TestIndexFlatL2Panorama, TestScalarQuantizer in test_index.py; TestDistanceComputer, TestIndexRefineSearchParams, TestIndexRefineRangeSearch in test_refine.py; TestIndexRefinePanorama, TestIndexFlatPanorama, TestIndexHNSWFlatPanorama, TestIndexIVFFlatPanorama in their respective files. * Quantizer encode: TestResidualQuantizer, TestIndexResidualQuantizerSearch in test_residual_quantizer.py; TestComponents, TestLocalSearchQuantizer in test_local_search_quantizer.py. * Fast scan: TestFastScanFiltering, TestBlockSkipConsistency, TestFastScanRangeSearchFilter in test_fastscan_filter.py. * Broader index tests (search-exercising classes only): TestParameterSpace in test_autotune.py; TestSpectralHash in test_factory.py; TestMerge1, TestMerge2, TestRemoveFastScan in test_merge_index.py. BUCK changes move decorated tests from the legacy-listing lists to the simd_levels lists (which use supports_static_listing = False, required for dynamic class name generation), and add the new entries test_binary_search_params, test_fastscan_filter, test_refine_panorama, test_hnsw_panorama. ___ overriding_review_checks_triggers_an_audit_and_retroactive_review Oncall Short Name: fair_umami_cluster Differential Revision: D101822335 fbshipit-source-id: 9a432c6d6bee201c3731713226405a8c8ecebbe6
687 lines
21 KiB
Python
687 lines
21 KiB
Python
# 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.
|
|
|
|
"""
|
|
Tests for the implementation of Local Search Quantizer
|
|
"""
|
|
|
|
import numpy as np
|
|
import platform
|
|
|
|
import faiss
|
|
import unittest
|
|
|
|
from common_faiss_tests import for_all_simd_levels
|
|
from faiss.contrib import datasets
|
|
|
|
faiss.omp_set_num_threads(4)
|
|
|
|
sp = faiss.swig_ptr
|
|
|
|
|
|
def construct_sparse_matrix(codes, K):
|
|
n, M = codes.shape
|
|
B = np.zeros((n, M * K), dtype=np.float32)
|
|
for i in range(n):
|
|
for j in range(M):
|
|
code = codes[i, j]
|
|
B[i, j * K + code] = 1
|
|
return B
|
|
|
|
|
|
def update_codebooks_ref(x, codes, K, lambd):
|
|
n, d = x.shape
|
|
M = codes.shape[1]
|
|
|
|
B = construct_sparse_matrix(codes, K)
|
|
reg = np.identity(M * K) * float(lambd)
|
|
reg = reg.astype(np.float32)
|
|
|
|
# C = (B'B + lambd * I)^(-1)B'X
|
|
bb = np.linalg.inv(B.T @ B + reg)
|
|
C = bb @ B.T @ x
|
|
C = C.reshape(M, K, d)
|
|
|
|
return C
|
|
|
|
|
|
def compute_binary_terms_ref(codebooks):
|
|
M, K, d = codebooks.shape
|
|
|
|
codebooks_t = np.swapaxes(codebooks, 1, 2) # [M, d, K]
|
|
binaries = 2 * codebooks.dot(codebooks_t) # [M, K, M, K]
|
|
binaries = np.swapaxes(binaries, 1, 2) # [M, M, K, K]
|
|
|
|
return binaries
|
|
|
|
|
|
def compute_unary_terms_ref(codebooks, x):
|
|
codebooks_t = np.swapaxes(codebooks, 1, 2) # [M, d, K]
|
|
unaries = -2 * x.dot(codebooks_t) # [n, M, K]
|
|
code_norms = np.sum(codebooks * codebooks, axis=2) # [M, K]
|
|
unaries += code_norms
|
|
unaries = np.swapaxes(unaries, 0, 1) # [M, n, K]
|
|
|
|
return unaries
|
|
|
|
|
|
def icm_encode_step_ref(unaries, binaries, codes):
|
|
M, n, K = unaries.shape
|
|
|
|
for m in range(M):
|
|
objs = unaries[m].copy() # [n, K]
|
|
|
|
for m2 in range(M): # pair, m2 != m
|
|
if m2 == m:
|
|
continue
|
|
|
|
for i in range(n):
|
|
for code in range(K):
|
|
code2 = codes[i, m2]
|
|
objs[i, code] += binaries[m, m2, code, code2]
|
|
|
|
codes[:, m] = np.argmin(objs, axis=1)
|
|
|
|
return codes
|
|
|
|
|
|
def decode_ref(x, codebooks, codes):
|
|
n, d = x.shape
|
|
_, M = codes.shape
|
|
decoded_x = np.zeros((n, d), dtype=np.float32)
|
|
for i in range(n):
|
|
for m in range(M):
|
|
decoded_x[i] += codebooks[m, codes[i, m]]
|
|
|
|
return decoded_x
|
|
|
|
|
|
def icm_encode_ref(x, codebooks, codes):
|
|
n, d = x.shape
|
|
M, K, d = codebooks.shape
|
|
|
|
codes = codes.copy()
|
|
for m in range(M):
|
|
objs = np.zeros((n, K), dtype=np.float32) # [n, K]
|
|
for code in range(K):
|
|
new_codes = codes.copy()
|
|
new_codes[:, m] = code
|
|
|
|
# decode x
|
|
decoded_x = decode_ref(x, codebooks, new_codes)
|
|
objs[:, code] = np.sum((x - decoded_x) ** 2, axis=1)
|
|
|
|
codes[:, m] = np.argmin(objs, axis=1)
|
|
|
|
return codes
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestComponents(unittest.TestCase):
|
|
|
|
def test_decode(self):
|
|
"""Test LSQ decode"""
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 6
|
|
K = (1 << nbits)
|
|
|
|
rs = np.random.RandomState(123)
|
|
x = rs.rand(n, d).astype(np.float32)
|
|
codes = rs.randint(0, K, (n, M)).astype(np.int32)
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.train(x)
|
|
|
|
# decode x
|
|
pack_codes = np.zeros((n, lsq.code_size)).astype(np.uint8)
|
|
decoded_x = np.zeros((n, d)).astype(np.float32)
|
|
lsq.pack_codes(n, sp(codes), sp(pack_codes))
|
|
lsq.decode_c(sp(pack_codes), sp(decoded_x), n)
|
|
|
|
# decode in Python
|
|
codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
codebooks = codebooks.reshape(M, K, d).copy()
|
|
decoded_x_ref = decode_ref(x, codebooks, codes)
|
|
|
|
np.testing.assert_allclose(decoded_x, decoded_x_ref, rtol=1e-6)
|
|
|
|
@unittest.skipIf(platform.system() == 'Windows',
|
|
'Does not work on Windows after numpy 2 upgrade.')
|
|
def test_update_codebooks(self):
|
|
"""Test codebooks updatation."""
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 6
|
|
K = (1 << nbits)
|
|
|
|
# set a larger value to make the updating process more stable
|
|
lambd = 1e-2
|
|
|
|
rs = np.random.RandomState(123)
|
|
x = rs.rand(n, d).astype(np.float32)
|
|
codes = rs.randint(0, K, (n, M)).astype(np.int32)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.lambd = lambd
|
|
lsq.train(x) # just for allocating memory for codebooks
|
|
|
|
codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
codebooks = codebooks.reshape(M, K, d).copy()
|
|
|
|
lsq.update_codebooks(sp(x), sp(codes), n)
|
|
new_codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
new_codebooks = new_codebooks.reshape(M, K, d).copy()
|
|
|
|
ref_codebooks = update_codebooks_ref(x, codes, K, lambd)
|
|
|
|
np.testing.assert_allclose(new_codebooks, ref_codebooks, rtol=1e-3, atol=1e-3)
|
|
|
|
def test_update_codebooks_with_double(self):
|
|
"""If the data is not zero-centering, it would be more accurate to
|
|
use double-precision floating-point numbers."""
|
|
ds = datasets.SyntheticDataset(16, 1000, 1000, 0)
|
|
|
|
xt = ds.get_train() + 1000
|
|
xb = ds.get_database() + 1000
|
|
|
|
M = 4
|
|
nbits = 4
|
|
|
|
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
|
|
lsq.train(xt)
|
|
err_double = eval_codec(lsq, xb)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
|
|
lsq.update_codebooks_with_double = False
|
|
lsq.train(xt)
|
|
err_float = eval_codec(lsq, xb)
|
|
|
|
# 6533.377 vs 25457.99
|
|
self.assertLess(err_double, err_float)
|
|
|
|
def test_compute_binary_terms(self):
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 6
|
|
K = (1 << nbits)
|
|
|
|
rs = np.random.RandomState(123)
|
|
x = rs.rand(n, d).astype(np.float32)
|
|
binaries = np.zeros((M, M, K, K)).astype(np.float32)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.train(x) # just for allocating memory for codebooks
|
|
|
|
lsq.compute_binary_terms(sp(binaries))
|
|
|
|
codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
codebooks = codebooks.reshape(M, K, d).copy()
|
|
ref_binaries = compute_binary_terms_ref(codebooks)
|
|
|
|
np.testing.assert_allclose(
|
|
binaries, ref_binaries, rtol=1e-4, atol=1e-4
|
|
)
|
|
|
|
def test_compute_unary_terms(self):
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 6
|
|
K = (1 << nbits)
|
|
|
|
rs = np.random.RandomState(123)
|
|
x = rs.rand(n, d).astype(np.float32)
|
|
unaries = np.zeros((M, n, K)).astype(np.float32)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.train(x) # just for allocating memory for codebooks
|
|
|
|
lsq.compute_unary_terms(sp(x), sp(unaries), n)
|
|
|
|
codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
codebooks = codebooks.reshape(M, K, d).copy()
|
|
ref_unaries = compute_unary_terms_ref(codebooks, x)
|
|
|
|
np.testing.assert_allclose(unaries, ref_unaries, rtol=1e-4, atol=1e-4)
|
|
|
|
def test_icm_encode_step(self):
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 6
|
|
K = (1 << nbits)
|
|
|
|
rs = np.random.RandomState(123)
|
|
|
|
# randomly generate codes and unary terms
|
|
codes = rs.randint(0, K, (n, M)).astype(np.int32)
|
|
new_codes = codes.copy()
|
|
unaries = rs.rand(M, n, K).astype(np.float32)
|
|
|
|
# binary terms should be symmetric, because binary terms
|
|
# represent cached dot products between the code C1 in codebook M1
|
|
# and the code C2 in codebook M2.
|
|
# so, binaries[M1, M2, C1, C2] == binaries[M2, M1, C2, C1]
|
|
#
|
|
# generate binary terms in a standard way that provides
|
|
# the needed symmetry
|
|
codebooks = rs.rand(M, K, d).astype(np.float32)
|
|
binaries = compute_binary_terms_ref(codebooks)
|
|
binaries = np.ascontiguousarray(binaries)
|
|
|
|
# do icm encoding given binary and unary terms
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.icm_encode_step(
|
|
sp(new_codes),
|
|
sp(unaries),
|
|
sp(binaries),
|
|
n,
|
|
1)
|
|
|
|
# do icm encoding given binary and unary terms in Python
|
|
ref_codes = icm_encode_step_ref(unaries, binaries, codes)
|
|
np.testing.assert_array_equal(new_codes, ref_codes)
|
|
|
|
def test_icm_encode(self):
|
|
d = 16
|
|
n = 500
|
|
M = 4
|
|
nbits = 4
|
|
K = (1 << nbits)
|
|
|
|
rs = np.random.RandomState(123)
|
|
x = rs.rand(n, d).astype(np.float32)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
|
|
lsq.train(x) # just for allocating memory for codebooks
|
|
|
|
# compute binary terms
|
|
binaries = np.zeros((M, M, K, K)).astype(np.float32)
|
|
lsq.compute_binary_terms(sp(binaries))
|
|
|
|
# compute unary terms
|
|
unaries = np.zeros((M, n, K)).astype(np.float32)
|
|
lsq.compute_unary_terms(sp(x), sp(unaries), n)
|
|
|
|
# randomly generate codes
|
|
codes = rs.randint(0, K, (n, M)).astype(np.int32)
|
|
new_codes = codes.copy()
|
|
|
|
# do icm encoding given binary and unary terms
|
|
lsq.icm_encode_step(
|
|
sp(new_codes),
|
|
sp(unaries),
|
|
sp(binaries),
|
|
n,
|
|
1)
|
|
|
|
# do icm encoding without pre-computed unary and binary terms in Python
|
|
codebooks = faiss.vector_float_to_array(lsq.codebooks)
|
|
codebooks = codebooks.reshape(M, K, d).copy()
|
|
ref_codes = icm_encode_ref(x, codebooks, codes)
|
|
|
|
np.testing.assert_array_equal(new_codes, ref_codes)
|
|
|
|
|
|
def eval_codec(q, xb):
|
|
codes = q.compute_codes(xb)
|
|
decoded = q.decode(codes)
|
|
return ((xb - decoded) ** 2).sum()
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def test_training(self):
|
|
"""check that the error is in the same ballpark as PQ."""
|
|
ds = datasets.SyntheticDataset(32, 3000, 3000, 0)
|
|
|
|
xt = ds.get_train()
|
|
xb = ds.get_database()
|
|
|
|
M = 4
|
|
nbits = 4
|
|
|
|
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
|
|
lsq.train(xt)
|
|
err_lsq = eval_codec(lsq, xb)
|
|
|
|
pq = faiss.ProductQuantizer(ds.d, M, nbits)
|
|
pq.train(xt)
|
|
err_pq = eval_codec(pq, xb)
|
|
|
|
self.assertLess(err_lsq, err_pq)
|
|
|
|
|
|
class TestIndexLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def test_IndexLocalSearchQuantizer(self):
|
|
ds = datasets.SyntheticDataset(32, 1000, 200, 100)
|
|
gt = ds.get_groundtruth(10)
|
|
|
|
ir = faiss.IndexLocalSearchQuantizer(ds.d, 4, 5)
|
|
ir.train(ds.get_train())
|
|
ir.add(ds.get_database())
|
|
Dref, Iref = ir.search(ds.get_queries(), 10)
|
|
inter_ref = faiss.eval_intersection(Iref, gt)
|
|
|
|
# 467
|
|
self.assertGreater(inter_ref, 460)
|
|
|
|
AQ = faiss.AdditiveQuantizer
|
|
ir2 = faiss.IndexLocalSearchQuantizer(
|
|
ds.d, 4, 5, faiss.METRIC_L2, AQ.ST_norm_float)
|
|
|
|
ir2.train(ds.get_train()) # just to set flags properly
|
|
ir2.lsq.codebooks = ir.lsq.codebooks
|
|
|
|
ir2.add(ds.get_database())
|
|
D2, I2 = ir2.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_almost_equal(Dref, D2, decimal=5)
|
|
self.assertLess((Iref != I2).sum(), Iref.size * 0.01)
|
|
|
|
# test I/O
|
|
ir3 = faiss.deserialize_index(faiss.serialize_index(ir))
|
|
D3, I3 = ir3.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(Iref, I3)
|
|
np.testing.assert_array_equal(Dref, D3)
|
|
|
|
def test_coarse_quantizer(self):
|
|
ds = datasets.SyntheticDataset(32, 5000, 1000, 100)
|
|
gt = ds.get_groundtruth(10)
|
|
|
|
quantizer = faiss.LocalSearchCoarseQuantizer(ds.d, 2, 4)
|
|
quantizer.lsq.nperts
|
|
quantizer.lsq.nperts = 2
|
|
|
|
index = faiss.IndexIVFFlat(quantizer, ds.d, 256)
|
|
index.quantizer_trains_alone = True
|
|
|
|
index.train(ds.get_train())
|
|
|
|
index.add(ds.get_database())
|
|
index.nprobe = 4
|
|
|
|
Dref, Iref = index.search(ds.get_queries(), 10)
|
|
|
|
inter_ref = faiss.eval_intersection(Iref, gt)
|
|
|
|
# 249
|
|
self.assertGreater(inter_ref, 235)
|
|
|
|
def test_factory(self):
|
|
index = faiss.index_factory(20, "LSQ5x6_Nqint8")
|
|
self.assertEqual(index.lsq.M, 5)
|
|
self.assertEqual(index.lsq.K, 1 << 6)
|
|
self.assertEqual(
|
|
index.lsq.search_type,
|
|
faiss.AdditiveQuantizer.ST_norm_qint8
|
|
)
|
|
|
|
index = faiss.index_factory(20, "LSQ5x6_Ncqint8")
|
|
self.assertEqual(
|
|
index.lsq.search_type,
|
|
faiss.AdditiveQuantizer.ST_norm_cqint8
|
|
)
|
|
|
|
index = faiss.index_factory(20, "LSQ5x6_Ncqint4")
|
|
self.assertEqual(
|
|
index.lsq.search_type,
|
|
faiss.AdditiveQuantizer.ST_norm_cqint4
|
|
)
|
|
|
|
|
|
|
|
class TestIndexIVFLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def test_factory(self):
|
|
index = faiss.index_factory(20, "IVF1024,LSQ5x6_Nqint8")
|
|
self.assertEqual(index.nlist, 1024)
|
|
self.assertEqual(index.lsq.M, 5)
|
|
self.assertEqual(index.lsq.K, 1 << 6)
|
|
self.assertEqual(
|
|
index.lsq.search_type,
|
|
faiss.AdditiveQuantizer.ST_norm_qint8
|
|
)
|
|
|
|
index = faiss.index_factory(20, "IVF1024,LSQ5x6_Ncqint8")
|
|
self.assertEqual(
|
|
index.lsq.search_type,
|
|
faiss.AdditiveQuantizer.ST_norm_cqint8
|
|
)
|
|
|
|
def eval_index_accuracy(self, factory_key):
|
|
# just do a single test, most search functions are already stress
|
|
# tested in test_residual_quantizer.py
|
|
ds = datasets.SyntheticDataset(32, 3000, 1000, 100)
|
|
index = faiss.index_factory(ds.d, factory_key)
|
|
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
|
|
inters = []
|
|
for nprobe in 1, 2, 5, 10, 20, 50:
|
|
index.nprobe = nprobe
|
|
D, I = index.search(ds.get_queries(), 10)
|
|
inter = faiss.eval_intersection(I, ds.get_groundtruth(10))
|
|
inters.append(inter)
|
|
|
|
inters = np.array(inters)
|
|
# in fact the results should be the same for the decoding and the
|
|
# reconstructing versions
|
|
self.assertTrue(np.all(inters[1:] >= inters[:-1]))
|
|
|
|
# do a little I/O test
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
D2, I2 = index2.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(I2, I)
|
|
np.testing.assert_array_equal(D2, D)
|
|
|
|
def test_index_accuracy_reconstruct(self):
|
|
self.eval_index_accuracy("IVF100,LSQ4x5")
|
|
|
|
def test_index_accuracy_reconstruct_LUT(self):
|
|
self.eval_index_accuracy("IVF100,LSQ4x5_Nfloat")
|
|
|
|
def test_index_accuracy_cqint(self):
|
|
self.eval_index_accuracy("IVF100,LSQ4x5_Ncqint8")
|
|
|
|
def test_deterministic(self):
|
|
ds = datasets.SyntheticDataset(d=16, nt=1000, nb=10005, nq=1000)
|
|
index = faiss.index_factory(ds.d, "IVF100,LSQ4x4_Nqint8")
|
|
|
|
k = 1
|
|
xq = ds.get_queries()
|
|
xb = ds.get_database()
|
|
xt = ds.get_train()
|
|
|
|
index.train(xt)
|
|
index.add(xb)
|
|
D, I = index.search(xq, k=k)
|
|
|
|
index.reset()
|
|
index.train(xt)
|
|
index.add(xb)
|
|
D2, I2 = index.search(xq, k=k)
|
|
|
|
np.testing.assert_array_equal(I, I2)
|
|
|
|
|
|
class TestProductLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def test_codec(self):
|
|
"""check that the error is in the same ballpark as PQ."""
|
|
ds = datasets.SyntheticDataset(64, 3000, 3000, 0)
|
|
|
|
xt = ds.get_train()
|
|
xb = ds.get_database()
|
|
|
|
nsplits = 2
|
|
Msub = 2
|
|
nbits = 4
|
|
|
|
plsq = faiss.ProductLocalSearchQuantizer(ds.d, nsplits, Msub, nbits)
|
|
plsq.train(xt)
|
|
err_plsq = eval_codec(plsq, xb)
|
|
|
|
pq = faiss.ProductQuantizer(ds.d, nsplits * Msub, nbits)
|
|
pq.train(xt)
|
|
err_pq = eval_codec(pq, xb)
|
|
|
|
self.assertLess(err_plsq, err_pq)
|
|
|
|
def test_with_lsq(self):
|
|
"""compare with LSQ when nsplits = 1"""
|
|
ds = datasets.SyntheticDataset(32, 3000, 3000, 0)
|
|
|
|
xt = ds.get_train()
|
|
xb = ds.get_database()
|
|
|
|
M = 4
|
|
nbits = 4
|
|
|
|
plsq = faiss.ProductLocalSearchQuantizer(ds.d, 1, M, nbits)
|
|
plsq.train(xt)
|
|
err_plsq = eval_codec(plsq, xb)
|
|
|
|
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
|
|
lsq.train(xt)
|
|
err_lsq = eval_codec(lsq, xb)
|
|
|
|
self.assertEqual(err_plsq, err_lsq)
|
|
|
|
def test_lut(self):
|
|
"""test compute_LUT function"""
|
|
ds = datasets.SyntheticDataset(16, 1000, 0, 100)
|
|
|
|
xt = ds.get_train()
|
|
xq = ds.get_queries()
|
|
|
|
nsplits = 2
|
|
Msub = 2
|
|
nbits = 4
|
|
nq, d = xq.shape
|
|
dsub = d // nsplits
|
|
|
|
plsq = faiss.ProductLocalSearchQuantizer(ds.d, nsplits, Msub, nbits)
|
|
plsq.train(xt)
|
|
|
|
subcodebook_size = Msub * (1 << nbits)
|
|
codebook_size = nsplits * subcodebook_size
|
|
lut = np.zeros((nq, codebook_size), dtype=np.float32)
|
|
plsq.compute_LUT(nq, sp(xq), sp(lut))
|
|
|
|
codebooks = faiss.vector_to_array(plsq.codebooks)
|
|
codebooks = codebooks.reshape(nsplits, subcodebook_size, dsub)
|
|
xq = xq.reshape(nq, nsplits, dsub)
|
|
lut_ref = np.zeros((nq, nsplits, subcodebook_size), dtype=np.float32)
|
|
for i in range(nsplits):
|
|
lut_ref[:, i] = xq[:, i] @ codebooks[i].T
|
|
lut_ref = lut_ref.reshape(nq, codebook_size)
|
|
|
|
np.testing.assert_allclose(lut, lut_ref, rtol=5e-04)
|
|
|
|
|
|
class TestIndexProductLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def test_accuracy1(self):
|
|
"""check that the error is in the same ballpark as LSQ."""
|
|
recall1 = self.eval_index_accuracy("PLSQ4x3x5_Nqint8")
|
|
recall2 = self.eval_index_accuracy("LSQ12x5_Nqint8")
|
|
self.assertGreaterEqual(recall1, recall2) # 622 vs 551
|
|
|
|
def test_accuracy2(self):
|
|
"""when nsplits = 1, PLSQ should be almost the same as LSQ"""
|
|
recall1 = self.eval_index_accuracy("PLSQ1x3x5_Nqint8")
|
|
recall2 = self.eval_index_accuracy("LSQ3x5_Nqint8")
|
|
diff = abs(recall1 - recall2) # 273 vs 275 in OSX
|
|
self.assertGreaterEqual(5, diff)
|
|
|
|
def eval_index_accuracy(self, index_key):
|
|
ds = datasets.SyntheticDataset(32, 1000, 1000, 100)
|
|
index = faiss.index_factory(ds.d, index_key)
|
|
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
D, I = index.search(ds.get_queries(), 10)
|
|
inter = faiss.eval_intersection(I, ds.get_groundtruth(10))
|
|
|
|
# do a little I/O test
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
D2, I2 = index2.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(I2, I)
|
|
np.testing.assert_array_equal(D2, D)
|
|
|
|
return inter
|
|
|
|
def test_factory(self):
|
|
AQ = faiss.AdditiveQuantizer
|
|
ns, Msub, nbits = 2, 4, 8
|
|
index = faiss.index_factory(64, f"PLSQ{ns}x{Msub}x{nbits}_Nqint8")
|
|
assert isinstance(index, faiss.IndexProductLocalSearchQuantizer)
|
|
self.assertEqual(index.plsq.nsplits, ns)
|
|
self.assertEqual(index.plsq.subquantizer(0).M, Msub)
|
|
self.assertEqual(index.plsq.subquantizer(0).nbits.at(0), nbits)
|
|
self.assertEqual(index.plsq.search_type, AQ.ST_norm_qint8)
|
|
|
|
code_size = (ns * Msub * nbits + 7) // 8 + 1
|
|
self.assertEqual(index.plsq.code_size, code_size)
|
|
|
|
|
|
class TestIndexIVFProductLocalSearchQuantizer(unittest.TestCase):
|
|
|
|
def eval_index_accuracy(self, factory_key):
|
|
ds = datasets.SyntheticDataset(32, 1000, 1000, 100)
|
|
index = faiss.index_factory(ds.d, factory_key)
|
|
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
|
|
inters = []
|
|
for nprobe in 1, 2, 4, 8, 16:
|
|
index.nprobe = nprobe
|
|
D, I = index.search(ds.get_queries(), 10)
|
|
inter = faiss.eval_intersection(I, ds.get_groundtruth(10))
|
|
inters.append(inter)
|
|
|
|
inters = np.array(inters)
|
|
self.assertTrue(np.all(inters[1:] >= inters[:-1]))
|
|
|
|
# do a little I/O test
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
D2, I2 = index2.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(I2, I)
|
|
np.testing.assert_array_equal(D2, D)
|
|
|
|
return inter
|
|
|
|
def test_index_accuracy(self):
|
|
self.eval_index_accuracy("IVF32,PLSQ2x2x5_Nqint8")
|
|
|
|
@unittest.skipIf(platform.system() == 'Windows',
|
|
'Does not work on Windows-2022+.')
|
|
def test_index_accuracy2(self):
|
|
"""check that the error is in the same ballpark as LSQ."""
|
|
inter1 = self.eval_index_accuracy("IVF32,PLSQ2x2x5_Nqint8")
|
|
inter2 = self.eval_index_accuracy("IVF32,LSQ4x5_Nqint8")
|
|
self.assertGreaterEqual(inter1 * 1.1, inter2)
|
|
|
|
def test_factory(self):
|
|
AQ = faiss.AdditiveQuantizer
|
|
ns, Msub, nbits = 2, 4, 8
|
|
index = faiss.index_factory(64, f"IVF32,PLSQ{ns}x{Msub}x{nbits}_Nqint8")
|
|
assert isinstance(index, faiss.IndexIVFProductLocalSearchQuantizer)
|
|
self.assertEqual(index.nlist, 32)
|
|
self.assertEqual(index.plsq.nsplits, ns)
|
|
self.assertEqual(index.plsq.subquantizer(0).M, Msub)
|
|
self.assertEqual(index.plsq.subquantizer(0).nbits.at(0), nbits)
|
|
self.assertEqual(index.plsq.search_type, AQ.ST_norm_qint8)
|
|
|
|
code_size = (ns * Msub * nbits + 7) // 8 + 1
|
|
self.assertEqual(index.plsq.code_size, code_size)
|