mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-11 22:50:00 +00:00
Michael Norris
ec94295adc
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5658 A fast scan kernel sums `M` uint8 table entries into a uint16 that can wrap around to 0. The scale that quantizes the table must therefore bound the whole sum that gets accumulated, not only one entry. Seven call sites each held their own copy of that rule. Only four applied it. This adds `quantize_lut::fastscan_lut_scale` and routes all seven through it. The bound also reserves a margin for `round_tab`, which rounds half up, so `n` terms can gain `n/2` over the unrounded sum. `IndexIVFRaBitQFastScan` also folds its per-probe bias into the bound, because that bias is rounded into the same accumulator. `IndexPQFastScan` only reaches the bound when a sub-vector holds one dimension. A nearest centroid keeps the selected entries near the bottom of each column. `IndexSQFastScan` sets `M = d`, so it fails above about `d = 512`. At `d = 1024` it goes from R@1 0.00 to 1.00. Search speed does not change. The scale is computed once per query, and the inner loop is untouched. An index at `M <= 256` quantizes to the same bytes as before, because the sum term does not bind there. Reviewed By: junjieqi Differential Revision: D120638521 fbshipit-source-id: 27d5dd15f5a60a9c80cec1f63f53ac9c6c51b8ff
809 lines
25 KiB
Python
809 lines
25 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.
|
|
|
|
|
|
import unittest
|
|
import time
|
|
|
|
import numpy as np
|
|
import faiss
|
|
|
|
from faiss.contrib import datasets
|
|
|
|
from common_faiss_tests import (
|
|
compare_binary_result_lists,
|
|
for_all_simd_levels,
|
|
NoneSIMDLevel,
|
|
)
|
|
|
|
# the tests tend to timeout in stress modes + dev otherwise
|
|
faiss.omp_set_num_threads(4)
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestSearch(unittest.TestCase):
|
|
|
|
def test_PQ4_accuracy(self):
|
|
ds = datasets.SyntheticDataset(32, 2000, 5000, 1000)
|
|
|
|
index_gt = faiss.IndexFlatL2(32)
|
|
index_gt.add(ds.get_database())
|
|
Dref, Iref = index_gt.search(ds.get_queries(), 10)
|
|
|
|
index = faiss.index_factory(32, "PQ16x4fs")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Da, Ia = index.search(ds.get_queries(), 10)
|
|
|
|
nq = Iref.shape[0]
|
|
recall_at_1 = (Iref[:, 0] == Ia[:, 0]).sum() / nq
|
|
assert recall_at_1 > 0.6
|
|
|
|
def test_PQ4_search_matches_none(self):
|
|
"""Fast-scan PQ search dispatches via the 256/512-bit QBS path.
|
|
|
|
The integer QBS kernel is deterministic across SIMD levels *only if*
|
|
the upstream float PQ distance table is bit-identical AND the
|
|
handler's end() method (int-to-float conversion) compiles to
|
|
identical FP code. Two cases break this:
|
|
|
|
1. aarch64 NEON/SVE: float distance-table computation reduces
|
|
in a different order, producing non-bit-identical LUT.
|
|
2. DD mode: per-SIMD TUs compile with different flags (-mfma
|
|
vs no -mfma), so handler end() may contract mul+add into
|
|
fma, producing 1-ULP float distance differences.
|
|
|
|
In both cases, fall back to a recall@1 check.
|
|
"""
|
|
ds = datasets.SyntheticDataset(32, 2000, 5000, 200)
|
|
index = faiss.index_factory(32, "PQ16x4fs")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
|
|
xq = ds.get_queries()
|
|
nq = xq.shape[0]
|
|
M, ksub = index.pq.M, index.pq.ksub
|
|
|
|
tab = np.zeros((nq, M, ksub), dtype="float32")
|
|
index.pq.compute_distance_tables(
|
|
nq, faiss.swig_ptr(xq), faiss.swig_ptr(tab)
|
|
)
|
|
D, I = index.search(xq, 10)
|
|
|
|
with NoneSIMDLevel():
|
|
tab_none = np.zeros((nq, M, ksub), dtype="float32")
|
|
index.pq.compute_distance_tables(
|
|
nq, faiss.swig_ptr(xq), faiss.swig_ptr(tab_none)
|
|
)
|
|
D_none, I_none = index.search(xq, 10)
|
|
|
|
lut_diffs = int((tab != tab_none).sum())
|
|
dd_mode = "DD" in faiss.get_compile_options()
|
|
if lut_diffs == 0 and not dd_mode:
|
|
compare_binary_result_lists(D, I, D_none, I_none)
|
|
else:
|
|
recall_at_1 = (I[:, 0] == I_none[:, 0]).mean()
|
|
self.assertGreater(
|
|
recall_at_1,
|
|
0.95,
|
|
f"LUT diverges ({lut_diffs} entries differ, DD={dd_mode}); "
|
|
f"recall@1 vs NONE = {recall_at_1:.3f}",
|
|
)
|
|
|
|
|
|
class TestFastScanDDSpeed(unittest.TestCase):
|
|
"""Verify DD-mode FastScan kernels use real SIMD."""
|
|
|
|
def _build_index(self):
|
|
d, n, nq, k = 128, 100_000, 1000, 10
|
|
np.random.seed(123)
|
|
xb = np.random.randn(n, d).astype("float32")
|
|
xq = np.random.randn(nq, d).astype("float32")
|
|
index = faiss.index_factory(d, "PQ16x4fs")
|
|
index.train(xb)
|
|
index.add(xb)
|
|
index.search(xq, k)
|
|
return index, xq, k
|
|
|
|
def _bench(self, index, xq, k, level, iters=50):
|
|
faiss.SIMDConfig.set_level(level)
|
|
t0 = time.time()
|
|
for _ in range(iters):
|
|
index.search(xq, k)
|
|
return time.time() - t0
|
|
|
|
def _check_speedup(self, none_time, simd_time, label):
|
|
speedup = none_time / simd_time
|
|
print(
|
|
f"FastScan DD: NONE={none_time:.3f}s, "
|
|
f"{label}={simd_time:.3f}s, {speedup:.1f}x"
|
|
)
|
|
# 2x is conservative: scalar-to-AVX2 theoretical max is ~32x,
|
|
# observed 5-27x in isolation. 2x leaves margin for CI noise.
|
|
self.assertGreater(
|
|
speedup,
|
|
2.0,
|
|
f"{label} should be >2x faster than NONE, " f"got {speedup:.1f}x",
|
|
)
|
|
|
|
@unittest.skipUnless(
|
|
"DD" in faiss.get_compile_options()
|
|
and "OPTIMIZE" in faiss.get_compile_options()
|
|
and faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_AVX2),
|
|
"DD-only, needs opt mode and AVX2",
|
|
)
|
|
def test_dd_fastscan_avx2_faster_than_none(self):
|
|
index, xq, k = self._build_index()
|
|
none_t = self._bench(index, xq, k, faiss.SIMDLevel_NONE)
|
|
avx2_t = self._bench(index, xq, k, faiss.SIMDLevel_AVX2)
|
|
faiss.SIMDConfig.set_level(faiss.SIMDConfig.get_dispatched_level())
|
|
self._check_speedup(none_t, avx2_t, "AVX2")
|
|
|
|
@unittest.skipUnless(
|
|
"DD" in faiss.get_compile_options()
|
|
and "OPTIMIZE" in faiss.get_compile_options()
|
|
and faiss.SIMDConfig.is_simd_level_available(faiss.SIMDLevel_AVX512),
|
|
"DD-only, needs opt mode and AVX512",
|
|
)
|
|
def test_dd_fastscan_avx512_faster_than_none(self):
|
|
index, xq, k = self._build_index()
|
|
none_t = self._bench(index, xq, k, faiss.SIMDLevel_NONE)
|
|
avx512_t = self._bench(index, xq, k, faiss.SIMDLevel_AVX512)
|
|
faiss.SIMDConfig.set_level(faiss.SIMDConfig.get_dispatched_level())
|
|
self._check_speedup(none_t, avx512_t, "AVX512")
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestRounding(unittest.TestCase):
|
|
|
|
def do_test_rounding(self, implem=4, metric=faiss.METRIC_L2):
|
|
ds = datasets.SyntheticDataset(32, 2000, 5000, 200)
|
|
|
|
index = faiss.index_factory(32, "PQ16x4", metric)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Dref, Iref = index.search(ds.get_queries(), 10)
|
|
nq = Iref.shape[0]
|
|
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
|
|
# simply repro normal search
|
|
index2.implem = 2
|
|
D2, I2 = index2.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(I2, Iref)
|
|
np.testing.assert_array_equal(D2, Dref)
|
|
|
|
# rounded LUT with correction
|
|
index2.implem = implem
|
|
D4, I4 = index2.search(ds.get_queries(), 10)
|
|
# check accuracy of indexes
|
|
recalls = {}
|
|
for rank in 1, 10:
|
|
recalls[rank] = (Iref[:, :1] == I4[:, :rank]).sum() / nq
|
|
|
|
min_r1 = 0.98 if metric == faiss.METRIC_INNER_PRODUCT else 0.99
|
|
self.assertGreaterEqual(recalls[1], min_r1)
|
|
self.assertGreater(recalls[10], 0.995)
|
|
# check accuracy of distances
|
|
# err3 = ((D3 - D2) ** 2).sum()
|
|
err4 = ((D4 - D2) ** 2).sum()
|
|
nf = (D2**2).sum()
|
|
self.assertLess(err4, nf * 1e-4)
|
|
|
|
def test_implem_4(self):
|
|
self.do_test_rounding(4)
|
|
|
|
def test_implem_4_ip(self):
|
|
self.do_test_rounding(4, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_implem_12(self):
|
|
self.do_test_rounding(12)
|
|
|
|
def test_implem_12_ip(self):
|
|
self.do_test_rounding(12, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_implem_14(self):
|
|
self.do_test_rounding(14)
|
|
|
|
def test_implem_14_ip(self):
|
|
self.do_test_rounding(12, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
|
|
class TestManyPQSubQuantizers(unittest.TestCase):
|
|
def do_test_recall_matches_pq(self, d):
|
|
ds = datasets.SyntheticDataset(d, 2000, 5000, 50)
|
|
# M = d gives one dimension per sub-quantizer.
|
|
pq = faiss.IndexPQ(d, d, 4)
|
|
pq.train(ds.get_train())
|
|
pq.add(ds.get_database())
|
|
|
|
fs = faiss.IndexPQFastScan(d, d, 4)
|
|
fs.train(ds.get_train())
|
|
fs.add(ds.get_database())
|
|
|
|
Dref, Iref = pq.search(ds.get_queries(), 10)
|
|
D, I = fs.search(ds.get_queries(), 10)
|
|
agree = (Iref[:, :1] == I[:, :1]).sum() / Iref.shape[0]
|
|
self.assertGreater(agree, 0.9)
|
|
|
|
def test_m_1024(self):
|
|
self.do_test_recall_matches_pq(1024)
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestReconstruct(unittest.TestCase):
|
|
|
|
def test_pqfastscan(self):
|
|
ds = datasets.SyntheticDataset(20, 1000, 1000, 0)
|
|
|
|
index = faiss.index_factory(20, "PQ5x4")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
recons = index.reconstruct_n(0, index.ntotal)
|
|
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
recons2 = index2.reconstruct_n(0, index.ntotal)
|
|
|
|
np.testing.assert_array_equal(recons, recons2)
|
|
|
|
def test_aqfastscan(self):
|
|
ds = datasets.SyntheticDataset(20, 1000, 1000, 0)
|
|
|
|
index = faiss.index_factory(20, "RQ5x4_Nrq2x4")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
recons = index.reconstruct_n(0, index.ntotal)
|
|
|
|
index2 = faiss.IndexAdditiveQuantizerFastScan(index)
|
|
recons2 = index2.reconstruct_n(0, index.ntotal)
|
|
|
|
np.testing.assert_array_equal(recons, recons2)
|
|
|
|
|
|
#########################################################
|
|
# Kernel unit test
|
|
#########################################################
|
|
|
|
|
|
def reference_accu(codes, LUT):
|
|
nq, nsp, is_16 = LUT.shape
|
|
nb, nsp_2 = codes.shape
|
|
assert is_16 == 16
|
|
assert nsp_2 == nsp // 2
|
|
accu = np.zeros((nq, nb), "uint16")
|
|
for i in range(nq):
|
|
for j in range(nb):
|
|
a = np.uint16(0)
|
|
for sp in range(0, nsp, 2):
|
|
c = codes[j, sp // 2]
|
|
a += LUT[i, sp, c & 15].astype("uint16")
|
|
a += LUT[i, sp + 1, c >> 4].astype("uint16")
|
|
accu[i, j] = a
|
|
return accu
|
|
|
|
|
|
# disabled because the function to write to mem is not implemented currently
|
|
|
|
|
|
class ThisIsNotATestLoop5: # (unittest.TestCase):
|
|
|
|
def do_loop5_kernel(self, nq, bb):
|
|
"""unit test for the accumulation kernel"""
|
|
nb = bb * 32 # databse size
|
|
nsp = 24 # number of sub-quantizers
|
|
|
|
rs = np.random.RandomState(123)
|
|
codes = rs.randint(256, size=(nb, nsp // 2)).astype("uint8")
|
|
LUT = rs.randint(256, size=(nq, nsp, 16)).astype("uint8")
|
|
accu_ref = reference_accu(codes, LUT)
|
|
|
|
def to_A(x):
|
|
return faiss.array_to_AlignedTable(x.ravel())
|
|
|
|
sp = faiss.swig_ptr
|
|
|
|
LUT_a = faiss.AlignedTableUint8(LUT.size)
|
|
faiss.pq4_pack_LUT(nq, nsp, sp(LUT), LUT_a.get())
|
|
|
|
codes_a = faiss.AlignedTableUint8(codes.size)
|
|
faiss.pq4_pack_codes(sp(codes), nb, nsp, nb, nb, nsp, codes_a.get())
|
|
|
|
accu_a = faiss.AlignedTableUint16(nq * nb)
|
|
accu_a.clear()
|
|
faiss.loop5_kernel_accumulate_1_block_to_mem(
|
|
nq, nb, nsp, codes_a.get(), LUT_a.get(), accu_a.get()
|
|
)
|
|
accu = faiss.AlignedTable_to_array(accu_a).reshape(nq, nb)
|
|
np.testing.assert_array_equal(accu_ref, accu)
|
|
|
|
def test_11(self):
|
|
self.do_loop5_kernel(1, 1)
|
|
|
|
def test_21(self):
|
|
self.do_loop5_kernel(2, 1)
|
|
|
|
def test_12(self):
|
|
self.do_loop5_kernel(1, 2)
|
|
|
|
def test_22(self):
|
|
self.do_loop5_kernel(2, 2)
|
|
|
|
|
|
##########################################################
|
|
# Tests for various IndexPQFastScan implementations
|
|
##########################################################
|
|
|
|
|
|
def verify_with_draws(testcase, Dref, Iref, Dnew, Inew):
|
|
"""Verify a list of results where there are draws in the distances
|
|
(because they are integer)."""
|
|
np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5)
|
|
# here we have to be careful because of draws
|
|
for i in range(len(Iref)):
|
|
if np.all(Iref[i] == Inew[i]):
|
|
continue
|
|
# we can deduce nothing about the latest line
|
|
skip_dis = Dref[i, -1]
|
|
for dis in np.unique(Dref):
|
|
if dis == skip_dis:
|
|
continue
|
|
mask = Dref[i, :] == dis
|
|
testcase.assertEqual(set(Iref[i, mask]), set(Inew[i, mask]))
|
|
|
|
|
|
class TestImplems(unittest.TestCase):
|
|
|
|
def __init__(self, *args):
|
|
unittest.TestCase.__init__(self, *args)
|
|
self.cache = {}
|
|
self.k = 10
|
|
|
|
def get_index(self, d, metric):
|
|
if (d, metric) not in self.cache:
|
|
ds = datasets.SyntheticDataset(d, 1000, 2000, 200)
|
|
target_size = d // 2
|
|
index = faiss.index_factory(d, "PQ%dx4" % target_size, metric)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
# uint8 LUT but no SIMD
|
|
index2.implem = 4
|
|
Dref, Iref = index2.search(ds.get_queries(), 10)
|
|
|
|
# check CodePacker
|
|
codes_ref = faiss.vector_to_array(index.codes)
|
|
codes_ref = codes_ref.reshape(-1, index.code_size)
|
|
index2codes = faiss.vector_to_array(index2.codes)
|
|
code_packer = index2.get_CodePacker()
|
|
index2codes = index2codes.reshape(-1, code_packer.block_size)
|
|
|
|
for i in range(0, len(codes_ref), 13):
|
|
code_new = code_packer.unpack_1(index2codes, i)
|
|
np.testing.assert_array_equal(codes_ref[i], code_new)
|
|
|
|
self.cache[(d, metric)] = (ds, index, Dref, Iref)
|
|
|
|
return self.cache[(d, metric)]
|
|
|
|
def do_with_params(self, d, params, metric=faiss.METRIC_L2):
|
|
ds, index, Dref, Iref = self.get_index(d, metric)
|
|
|
|
index2 = self.build_fast_scan_index(index, params)
|
|
|
|
Dnew, Inew = index2.search(ds.get_queries(), self.k)
|
|
|
|
Dref = Dref[:, : self.k]
|
|
Iref = Iref[:, : self.k]
|
|
|
|
verify_with_draws(self, Dref, Iref, Dnew, Inew)
|
|
|
|
def build_fast_scan_index(self, index, params):
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
index2.implem = 5
|
|
return index2
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestImplem12(TestImplems):
|
|
|
|
def build_fast_scan_index(self, index, qbs):
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
index2.qbs = qbs
|
|
index2.implem = 12
|
|
return index2
|
|
|
|
def test_qbs7(self):
|
|
self.do_with_params(32, 0x223)
|
|
|
|
def test_qbs7b(self):
|
|
self.do_with_params(32, 0x133)
|
|
|
|
def test_qbs6(self):
|
|
self.do_with_params(32, 0x33)
|
|
|
|
def test_qbs6_ip(self):
|
|
self.do_with_params(32, 0x33, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_qbs6b(self):
|
|
# test codepath where qbs is not known at compile time
|
|
self.do_with_params(32, 0x1113)
|
|
|
|
def test_qbs6_odd_dim(self):
|
|
self.do_with_params(30, 0x33)
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestImplem13(TestImplems):
|
|
|
|
def build_fast_scan_index(self, index, qbs):
|
|
index2 = faiss.IndexPQFastScan(index)
|
|
index2.qbs = qbs
|
|
index2.implem = 13
|
|
return index2
|
|
|
|
def test_qbs7(self):
|
|
self.do_with_params(32, 0x223)
|
|
|
|
def test_qbs7_k1(self):
|
|
self.k = 1
|
|
self.do_with_params(32, 0x223)
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestImplem14(TestImplems):
|
|
|
|
def build_fast_scan_index(self, index, params):
|
|
qbs, bbs = params
|
|
index2 = faiss.IndexPQFastScan(index, bbs)
|
|
index2.qbs = qbs
|
|
index2.implem = 14
|
|
return index2
|
|
|
|
def test_1_32(self):
|
|
self.do_with_params(32, (1, 32))
|
|
|
|
def test_1_64(self):
|
|
self.do_with_params(32, (1, 64))
|
|
|
|
def test_2_32(self):
|
|
self.do_with_params(32, (2, 32))
|
|
|
|
def test_2_64(self):
|
|
self.do_with_params(32, (2, 64))
|
|
|
|
def test_qbs_1_32_k1(self):
|
|
self.k = 1
|
|
self.do_with_params(32, (1, 32))
|
|
|
|
def test_qbs_1_64_k1(self):
|
|
self.k = 1
|
|
self.do_with_params(32, (1, 64))
|
|
|
|
def test_1_32_odd_dim(self):
|
|
self.do_with_params(30, (1, 32))
|
|
|
|
def test_1_64_odd_dim(self):
|
|
self.do_with_params(30, (1, 64))
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestImplem15(TestImplems):
|
|
|
|
def build_fast_scan_index(self, index, params):
|
|
qbs, bbs = params
|
|
index2 = faiss.IndexPQFastScan(index, bbs)
|
|
index2.qbs = qbs
|
|
index2.implem = 15
|
|
return index2
|
|
|
|
def test_1_32(self):
|
|
self.do_with_params(32, (1, 32))
|
|
|
|
def test_2_64(self):
|
|
self.do_with_params(32, (2, 64))
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestAdd(unittest.TestCase):
|
|
|
|
def do_test_add(self, d, bbs):
|
|
|
|
ds = datasets.SyntheticDataset(d, 2000, 5000, 200)
|
|
|
|
index = faiss.index_factory(d, f"PQ{d // 2}x4np")
|
|
index.train(ds.get_train())
|
|
|
|
xb = ds.get_database()
|
|
index.add(xb[:1235])
|
|
|
|
index2 = faiss.IndexPQFastScan(index, bbs)
|
|
index2.add(xb[1235:])
|
|
new_codes = faiss.AlignedTable_to_array(index2.codes)
|
|
|
|
index.add(xb[1235:])
|
|
index3 = faiss.IndexPQFastScan(index, bbs)
|
|
ref_codes = faiss.AlignedTable_to_array(index3.codes)
|
|
self.assertEqual(index3.ntotal, index2.ntotal)
|
|
|
|
np.testing.assert_array_equal(ref_codes, new_codes)
|
|
|
|
def test_add(self):
|
|
self.do_test_add(32, 32)
|
|
|
|
def test_add_bbs64(self):
|
|
self.do_test_add(32, 64)
|
|
|
|
def test_add_odd_d(self):
|
|
self.do_test_add(30, 64)
|
|
|
|
def test_constructor(self):
|
|
d = 32
|
|
ds = datasets.SyntheticDataset(d, 2000, 5000, 200)
|
|
|
|
index = faiss.index_factory(d, f"PQ{d // 2}x4np")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Dref, Iref = index.search(ds.get_queries(), 10)
|
|
nq = Iref.shape[0]
|
|
|
|
index2 = faiss.IndexPQFastScan(d, d // 2, 4)
|
|
index2.train(ds.get_train())
|
|
index2.add(ds.get_database())
|
|
Dnew, Inew = index2.search(ds.get_queries(), 10)
|
|
|
|
recall_at_1 = (Iref[:, 0] == Inew[:, 0]).sum() / nq
|
|
|
|
self.assertGreaterEqual(recall_at_1, 0.99)
|
|
|
|
data = faiss.serialize_index(index2)
|
|
index3 = faiss.deserialize_index(data)
|
|
|
|
self.assertEqual(index2.implem, index3.implem)
|
|
|
|
D3, I3 = index3.search(ds.get_queries(), 10)
|
|
np.testing.assert_array_equal(D3, Dnew)
|
|
np.testing.assert_array_equal(I3, Inew)
|
|
|
|
|
|
class TestAQFastScan(unittest.TestCase):
|
|
|
|
def subtest_accuracy(self, aq, st, implem, metric_type="L2"):
|
|
"""
|
|
Compare IndexAdditiveQuantizerFastScan with IndexAQ (qint8)
|
|
"""
|
|
d = 16
|
|
ds = datasets.SyntheticDataset(d, 1000, 1000, 500, metric_type)
|
|
gt = ds.get_groundtruth(k=1)
|
|
|
|
if metric_type == "L2":
|
|
metric = faiss.METRIC_L2
|
|
postfix1 = "_Nqint8"
|
|
postfix2 = f"_N{st}2x4"
|
|
else:
|
|
metric = faiss.METRIC_INNER_PRODUCT
|
|
postfix1 = postfix2 = ""
|
|
|
|
index = faiss.index_factory(d, f"{aq}3x4{postfix1}", metric)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Dref, Iref = index.search(ds.get_queries(), 1)
|
|
|
|
indexfs = faiss.index_factory(d, f"{aq}3x4fs_32{postfix2}", metric)
|
|
indexfs.train(ds.get_train())
|
|
indexfs.add(ds.get_database())
|
|
indexfs.implem = implem
|
|
Da, Ia = indexfs.search(ds.get_queries(), 1)
|
|
|
|
nq = Iref.shape[0]
|
|
recall_ref = (Iref == gt).sum() / nq
|
|
recall = (Ia == gt).sum() / nq
|
|
|
|
assert abs(recall_ref - recall) < 0.05
|
|
|
|
def xx_test_accuracy(self):
|
|
for metric in "L2", "IP":
|
|
for implem in 0, 12, 13, 14, 15:
|
|
self.subtest_accuracy("RQ", "rq", implem, metric)
|
|
self.subtest_accuracy("LSQ", "lsq", implem, metric)
|
|
|
|
def subtest_from_idxaq(self, implem, metric):
|
|
if metric == "L2":
|
|
metric_type = faiss.METRIC_L2
|
|
st = "_Nrq2x4"
|
|
else:
|
|
metric_type = faiss.METRIC_INNER_PRODUCT
|
|
st = ""
|
|
|
|
d = 16
|
|
ds = datasets.SyntheticDataset(d, 1000, 2000, 1000, metric=metric)
|
|
gt = ds.get_groundtruth(k=1)
|
|
index = faiss.index_factory(d, "RQ8x4" + st, metric_type)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Dref, Iref = index.search(ds.get_queries(), 1)
|
|
|
|
indexfs = faiss.IndexAdditiveQuantizerFastScan(index)
|
|
indexfs.implem = implem
|
|
D1, I1 = indexfs.search(ds.get_queries(), 1)
|
|
|
|
nq = Iref.shape[0]
|
|
recall_ref = (Iref == gt).sum() / nq
|
|
recall1 = (I1 == gt).sum() / nq
|
|
assert abs(recall_ref - recall1) < 0.05
|
|
|
|
def xx_test_from_idxaq(self):
|
|
for implem in 2, 3, 4:
|
|
self.subtest_from_idxaq(implem, "L2")
|
|
self.subtest_from_idxaq(implem, "IP")
|
|
|
|
def subtest_factory(self, aq, M, bbs, st):
|
|
"""
|
|
Format: {AQ}{M}x4fs_{bbs}_N{st}
|
|
|
|
AQ (str): `LSQ` or `RQ`
|
|
M (int): number of subquantizers
|
|
bbs (int): build block size
|
|
st (str): search type, `lsq2x4` or `rq2x4`
|
|
"""
|
|
AQ = faiss.AdditiveQuantizer
|
|
d = 16
|
|
|
|
if bbs > 0:
|
|
index = faiss.index_factory(d, f"{aq}{M}x4fs_{bbs}_N{st}2x4")
|
|
else:
|
|
index = faiss.index_factory(d, f"{aq}{M}x4fs_N{st}2x4")
|
|
bbs = 32
|
|
|
|
assert index.bbs == bbs
|
|
aq = faiss.downcast_AdditiveQuantizer(index.aq)
|
|
assert aq.M == M
|
|
|
|
if aq == "LSQ":
|
|
assert isinstance(aq, faiss.LocalSearchQuantizer)
|
|
if aq == "RQ":
|
|
assert isinstance(aq, faiss.ResidualQuantizer)
|
|
|
|
if st == "lsq":
|
|
assert aq.search_type == AQ.ST_norm_lsq2x4
|
|
if st == "rq":
|
|
assert aq.search_type == AQ.ST_norm_rq2x4
|
|
|
|
def test_factory(self):
|
|
self.subtest_factory("LSQ", 16, 64, "lsq")
|
|
self.subtest_factory("LSQ", 16, 64, "rq")
|
|
self.subtest_factory("RQ", 16, 64, "rq")
|
|
self.subtest_factory("RQ", 16, 64, "lsq")
|
|
self.subtest_factory("LSQ", 64, 0, "lsq")
|
|
|
|
def subtest_io(self, factory_str):
|
|
d = 8
|
|
ds = datasets.SyntheticDataset(d, 1000, 500, 100)
|
|
|
|
index = faiss.index_factory(d, factory_str)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
D1, I1 = index.search(ds.get_queries(), 1)
|
|
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
D2, I2 = index2.search(ds.get_queries(), 1)
|
|
np.testing.assert_array_equal(I1, I2)
|
|
|
|
def test_io(self):
|
|
self.subtest_io("LSQ4x4fs_Nlsq2x4")
|
|
self.subtest_io("LSQ4x4fs_Nrq2x4")
|
|
self.subtest_io("RQ4x4fs_Nrq2x4")
|
|
self.subtest_io("RQ4x4fs_Nlsq2x4")
|
|
|
|
|
|
# programatically generate tests to get finer test granularity.
|
|
|
|
|
|
def add_TestAQFastScan_subset_accuracy(aq, st, implem, metric):
|
|
setattr(
|
|
TestAQFastScan,
|
|
f"test_accuracy_{metric}_{aq}_implem{implem}",
|
|
lambda self: self.subtest_accuracy(aq, st, implem, metric),
|
|
)
|
|
|
|
|
|
for metric in "L2", "IP":
|
|
for implem in 0, 12, 13, 14, 15:
|
|
add_TestAQFastScan_subset_accuracy("LSQ", "lsq", implem, metric)
|
|
add_TestAQFastScan_subset_accuracy("RQ", "rq", implem, metric)
|
|
|
|
|
|
def add_TestAQFastScan_subtest_from_idxaq(implem, metric):
|
|
setattr(
|
|
TestAQFastScan,
|
|
f"test_from_idxaq_{metric}_implem{implem}",
|
|
lambda self: self.subtest_from_idxaq(implem, metric),
|
|
)
|
|
|
|
|
|
for implem in 2, 3, 4:
|
|
add_TestAQFastScan_subtest_from_idxaq(implem, "L2")
|
|
add_TestAQFastScan_subtest_from_idxaq(implem, "IP")
|
|
|
|
# Apply decorator after dynamic method generation.
|
|
TestAQFastScan = for_all_simd_levels(TestAQFastScan)
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestPAQFastScan(unittest.TestCase):
|
|
|
|
def subtest_accuracy(self, paq):
|
|
"""
|
|
Compare IndexPAQFastScan with IndexPAQ (qint8)
|
|
"""
|
|
d = 16
|
|
ds = datasets.SyntheticDataset(d, 1000, 1000, 500)
|
|
gt = ds.get_groundtruth(k=1)
|
|
|
|
index = faiss.index_factory(d, f"{paq}2x3x4_Nqint8")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
Dref, Iref = index.search(ds.get_queries(), 1)
|
|
|
|
indexfs = faiss.index_factory(d, f"{paq}2x3x4fs_Nlsq2x4")
|
|
indexfs.train(ds.get_train())
|
|
indexfs.add(ds.get_database())
|
|
Da, Ia = indexfs.search(ds.get_queries(), 1)
|
|
|
|
nq = Iref.shape[0]
|
|
recall_ref = (Iref == gt).sum() / nq
|
|
recall = (Ia == gt).sum() / nq
|
|
|
|
assert abs(recall_ref - recall) < 0.05
|
|
|
|
def test_accuracy_PLSQ(self):
|
|
self.subtest_accuracy("PLSQ")
|
|
|
|
def test_accuracy_PRQ(self):
|
|
self.subtest_accuracy("PRQ")
|
|
|
|
def subtest_factory(self, paq):
|
|
index = faiss.index_factory(16, f"{paq}2x3x4fs_Nlsq2x4")
|
|
q = faiss.downcast_Quantizer(index.aq)
|
|
self.assertEqual(q.nsplits, 2)
|
|
self.assertEqual(q.subquantizer(0).M, 3)
|
|
|
|
def test_factory(self):
|
|
self.subtest_factory("PRQ")
|
|
self.subtest_factory("PLSQ")
|
|
|
|
def subtest_io(self, factory_str):
|
|
d = 8
|
|
ds = datasets.SyntheticDataset(d, 1000, 500, 100)
|
|
|
|
index = faiss.index_factory(d, factory_str)
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
D1, I1 = index.search(ds.get_queries(), 1)
|
|
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
D2, I2 = index2.search(ds.get_queries(), 1)
|
|
np.testing.assert_array_equal(I1, I2)
|
|
|
|
def test_io(self):
|
|
self.subtest_io("PLSQ2x3x4fs_Nlsq2x4")
|
|
self.subtest_io("PRQ2x3x4fs_Nrq2x4")
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestBlockDecode(unittest.TestCase):
|
|
|
|
def test_issue_2739(self):
|
|
ds = datasets.SyntheticDataset(960, 200, 1, 0)
|
|
M = 32
|
|
index = faiss.index_factory(ds.d, f"PQ{M}x4fs")
|
|
index.train(ds.get_train())
|
|
index.add(ds.get_database())
|
|
|
|
decoded = index.pq.decode(index.pq.compute_codes(ds.get_database()))[
|
|
0, ::100
|
|
]
|
|
np.testing.assert_array_equal(decoded, index.reconstruct(0)[::100])
|