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, atol=1e-3)
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 bianry 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 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)
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, binary terms and unary terms
codes = rs.randint(0, K, (n, M)).astype(np.int32)
new_codes = codes.copy()
unaries = rs.rand(n, M, K).astype(np.float32)
binaries = rs.rand(M, M, K, K).astype(np.float32)
# do icm encoding given binary and unary terms
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
lsq.icm_encode_step(
faiss.swig_ptr(unaries),
faiss.swig_ptr(binaries),
faiss.swig_ptr(new_codes), n)
# 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 subtest_gpu_encoding(self, ngpus):
"""check that the error is in the same as cpu."""
ds = datasets.SyntheticDataset(32, 1000, 1000, 0)
xt = ds.get_train()
xb = ds.get_database()
M = 4
nbits = 8
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
lsq.train(xt)
err_cpu = self.eval_codec(lsq, xb)
lsq = faiss.LocalSearchQuantizer(ds.d, M, nbits)
lsq.train(xt)
lsq.icm_encoder_factory = faiss.GpuIcmEncoderFactory(ngpus)
err_gpu = self.eval_codec(lsq, xb)
# 13804.411 vs 13814.794, 1 gpu
print(err_gpu, err_cpu)
self.assertLess(err_gpu, err_cpu * 1.05)
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
print(err_double, err_float)
self.assertLess(err_double, err_float)
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)
print(err_plsq, err_lsq)
self.assertEqual(err_plsq, err_lsq)
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)
print(err_lsq, err_pq)
self.assertLess(err_lsq, err_pq)
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, atol=1e-4)
print(f"eval on {M}x{nbits} maxtrain={maxtrain}")
xq = ds.get_queries()
xb = ds.get_database()
gt = ds.get_groundtruth()
xt = ds.get_train(maxtrain=maxtrain)
nb, d = xb.shape
nq, d = xq.shape
nt, d = xt.shape
# fastest to slowest
if 'lsq-gpu' in todo:
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
ngpus = faiss.get_num_gpus()
lsq.icm_encoder_factory = faiss.GpuIcmEncoderFactory(ngpus)
lsq.verbose = True
eval_quantizer(lsq, xb, xt, 'lsq-gpu')
if 'pq' in todo:
pq = faiss.ProductQuantizer(d, M, nbits)
print("===== PQ")
eval_quantizer(pq, xq, xb, gt, xt)
if 'opq' in todo:
d2 = ((d + M - 1) // M) * M
print("OPQ d2=", d2)
opq = faiss.OPQMatrix(d, M, d2)
opq.train(xt)
