455 lines
13 KiB
Python
455 lines
13 KiB
Python
# 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.
|
|
|
|
"""
|
|
Tests for the implementation of Local Search Quantizer
|
|
"""
|
|
|
|
import numpy as np
|
|
|
|
import faiss
|
|
import unittest
|
|
|
|
from faiss.contrib import datasets
|
|
|
|
|
|
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
|
|
|
|
|
|
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)
|
|
|
|
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_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_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, 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, 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, binary terms 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)
|
|
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(
|
|
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 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 eval_codec(q, xb):
|
|
codes = q.compute_codes(xb)
|
|
decoded = q.decode(codes)
|
|
return ((xb - decoded) ** 2).sum()
|
|
|
|
|
|
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)
|
|
|
|
print(err_lsq, err_pq)
|
|
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
|
|
)
|
|
|
|
|
|
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
|
|
)
|
|
|
|
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))
|
|
# print("nprobe=", nprobe, "inter=", inter)
|
|
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")
|