mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Michael Norris
a74867cf56
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5611 D112025877 added the `faiss::hnsw_deterministic_build` flag. The flag chose between the lock-based graph build and the deterministic graph build. This diff deletes the flag. The deterministic build is now the only build. Reviewed By: trang-nm-nguyen Differential Revision: D118342471 fbshipit-source-id: d3c5e5238aca2bcda463eb27959ec22827222c15
847 lines
28 KiB
Python
847 lines
28 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.
|
|
|
|
"""a few tests for graph-based indices (HNSW, nndescent and NSG)"""
|
|
|
|
import numpy as np
|
|
import unittest
|
|
import faiss
|
|
|
|
from common_faiss_tests import get_dataset_2, for_all_simd_levels
|
|
|
|
|
|
@for_all_simd_levels
|
|
class TestHNSW(unittest.TestCase):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
unittest.TestCase.__init__(self, *args, **kwargs)
|
|
d = 32
|
|
nt = 0
|
|
nb = 1500
|
|
nq = 500
|
|
|
|
(_, self.xb, self.xq) = get_dataset_2(d, nt, nb, nq)
|
|
index = faiss.IndexFlatL2(d)
|
|
index.add(self.xb)
|
|
Dref, Iref = index.search(self.xq, 1)
|
|
self.Iref = Iref
|
|
|
|
def test_hnsw(self):
|
|
d = self.xq.shape[1]
|
|
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.add(self.xb)
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 460)
|
|
|
|
self.io_and_retest(index, Dhnsw, Ihnsw)
|
|
|
|
def test_deterministic_build(self):
|
|
# The HNSW build is deterministic and lock-free: repeated builds, and
|
|
# builds at different thread counts, produce byte-identical graphs.
|
|
# nb must be large enough that the batch-size cap (2% of n) exceeds
|
|
# the >100 threshold gating the parallel build region; otherwise the
|
|
# build runs single-threaded and determinism is trivially untested.
|
|
d = 32
|
|
nb = 20000
|
|
xb = np.random.RandomState(123).random((nb, d)).astype("float32")
|
|
|
|
# build() mutates the process-global OpenMP thread count; restore it so
|
|
# it does not leak into subsequent tests.
|
|
self.addCleanup(faiss.omp_set_num_threads, faiss.omp_get_max_threads())
|
|
|
|
def build(nthreads):
|
|
faiss.omp_set_num_threads(nthreads)
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.add(xb)
|
|
return index
|
|
|
|
# `index.hnsw` returns a fresh (copied) proxy and `.neighbors` /
|
|
# `.offsets` are views into it, so chaining `index.hnsw.neighbors`
|
|
# inline reads a freed temporary. Bind the hnsw proxy to keep it alive
|
|
# across the (copying) vector_to_array calls.
|
|
def graph_of(index):
|
|
hnsw = index.hnsw
|
|
return (
|
|
hnsw.max_level,
|
|
faiss.vector_to_array(hnsw.neighbors),
|
|
faiss.vector_to_array(hnsw.offsets),
|
|
)
|
|
|
|
idx8 = build(8)
|
|
max_level, g8a, offs = graph_of(idx8)
|
|
|
|
# graph is non-trivial: multi-level, many real edges
|
|
self.assertGreaterEqual(max_level, 1)
|
|
self.assertGreater(int((g8a >= 0).sum()), nb)
|
|
|
|
# reproducible across repeated builds and across thread counts
|
|
g8b = graph_of(build(8))[1]
|
|
g1 = graph_of(build(1))[1]
|
|
np.testing.assert_array_equal(g8a, g8b)
|
|
np.testing.assert_array_equal(g8a, g1)
|
|
|
|
# no self-loops
|
|
for i in range(nb):
|
|
seg = g8a[int(offs[i]) : int(offs[i + 1])]
|
|
self.assertNotIn(i, seg[seg >= 0].tolist())
|
|
|
|
def test_range_search(self):
|
|
index_flat = faiss.IndexFlat(self.xb.shape[1])
|
|
index_flat.add(self.xb)
|
|
D, _ = index_flat.search(self.xq, 10)
|
|
radius = np.median(D[:, -1])
|
|
lims_ref, Dref, Iref = index_flat.range_search(self.xq, radius)
|
|
|
|
index = faiss.IndexHNSWFlat(self.xb.shape[1], 16)
|
|
index.add(self.xb)
|
|
lims, D, I = index.range_search(self.xq, radius)
|
|
|
|
nmiss = 0
|
|
# check if returned results are a subset of the reference results
|
|
for i in range(len(self.xq)):
|
|
ref = Iref[lims_ref[i] : lims_ref[i + 1]]
|
|
new = I[lims[i] : lims[i + 1]]
|
|
self.assertLessEqual(set(new), set(ref))
|
|
nmiss += len(ref) - len(new)
|
|
# currently we miss 405 / 6019 neighbors
|
|
self.assertLessEqual(nmiss, lims_ref[-1] * 0.1)
|
|
|
|
def test_hnsw_unbounded_queue(self):
|
|
d = self.xq.shape[1]
|
|
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.add(self.xb)
|
|
index.hnsw.search_bounded_queue = False
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 460)
|
|
|
|
self.io_and_retest(index, Dhnsw, Ihnsw)
|
|
|
|
def test_hnsw_no_init_level0(self):
|
|
d = self.xq.shape[1]
|
|
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.init_level0 = False
|
|
index.add(self.xb)
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
# This is expected to be smaller because we are not initializing
|
|
# vectors into level 0.
|
|
self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 25)
|
|
|
|
self.io_and_retest(index, Dhnsw, Ihnsw)
|
|
|
|
def io_and_retest(self, index, Dhnsw, Ihnsw):
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
Dhnsw2, Ihnsw2 = index2.search(self.xq, 1)
|
|
|
|
self.assertTrue(np.all(Dhnsw2 == Dhnsw))
|
|
self.assertTrue(np.all(Ihnsw2 == Ihnsw))
|
|
|
|
# Verify deserialized index is serializable again
|
|
index2b = faiss.deserialize_index(faiss.serialize_index(index2))
|
|
Dhnsw2b, Ihnsw2b = index2b.search(self.xq, 1)
|
|
self.assertTrue(np.all(Dhnsw2b == Dhnsw))
|
|
self.assertTrue(np.all(Ihnsw2b == Ihnsw))
|
|
|
|
# also test clone
|
|
index3 = faiss.clone_index(index)
|
|
Dhnsw3, Ihnsw3 = index3.search(self.xq, 1)
|
|
|
|
self.assertTrue(np.all(Dhnsw3 == Dhnsw))
|
|
self.assertTrue(np.all(Ihnsw3 == Ihnsw))
|
|
|
|
def test_hnsw_2level(self):
|
|
d = self.xq.shape[1]
|
|
|
|
quant = faiss.IndexFlatL2(d)
|
|
|
|
index = faiss.IndexHNSW2Level(quant, 256, 8, 8)
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 307)
|
|
|
|
self.io_and_retest(index, Dhnsw, Ihnsw)
|
|
|
|
def test_hnsw_2level_mixed_search(self):
|
|
d = self.xq.shape[1]
|
|
|
|
quant = faiss.IndexFlatL2(d)
|
|
|
|
storage = faiss.IndexIVFPQ(quant, d, 32, 8, 8)
|
|
storage.make_direct_map()
|
|
index = faiss.IndexHNSW2Level(quant, 32, 8, 8)
|
|
index.storage = storage
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
# It is expected that the mixed search will perform worse.
|
|
self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 200)
|
|
|
|
self.io_and_retest(index, Dhnsw, Ihnsw)
|
|
|
|
def test_add_0_vecs(self):
|
|
index = faiss.IndexHNSWFlat(10, 16)
|
|
zero_vecs = np.zeros((0, 10), dtype="float32")
|
|
# infinite loop
|
|
index.add(zero_vecs)
|
|
self.assertEqual(index.ntotal, 0)
|
|
|
|
def test_hnsw_IP(self):
|
|
d = self.xq.shape[1]
|
|
|
|
index_IP = faiss.IndexFlatIP(d)
|
|
index_IP.add(self.xb)
|
|
Dref, Iref = index_IP.search(self.xq, 1)
|
|
|
|
index = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT)
|
|
index.add(self.xb)
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
|
|
self.assertGreaterEqual((Iref == Ihnsw).sum(), 470)
|
|
|
|
mask = Iref[:, 0] == Ihnsw[:, 0]
|
|
self.assertTrue(np.allclose(Dref[mask, 0], Dhnsw[mask, 0]))
|
|
|
|
def test_ndis_stats(self):
|
|
d = self.xq.shape[1]
|
|
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.add(self.xb)
|
|
stats = faiss.cvar.hnsw_stats
|
|
stats.reset()
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
self.assertGreater(stats.ndis, len(self.xq) * index.hnsw.efSearch)
|
|
|
|
def test_io_no_storage(self):
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexHNSWFlat(d, 16)
|
|
index.add(self.xb)
|
|
|
|
Dref, Iref = index.search(self.xq, 5)
|
|
|
|
# test writing without storage
|
|
index2 = faiss.deserialize_index(
|
|
faiss.serialize_index(index, faiss.IO_FLAG_SKIP_STORAGE)
|
|
)
|
|
self.assertEqual(index2.storage, None)
|
|
self.assertRaises(RuntimeError, index2.search, self.xb, 1)
|
|
|
|
# make sure we can store an index with empty storage
|
|
faiss.deserialize_index(faiss.serialize_index(index2))
|
|
|
|
# add storage afterwards
|
|
index.storage = faiss.clone_index(index.storage)
|
|
index.own_fields = True
|
|
|
|
Dnew, Inew = index.search(self.xq, 5)
|
|
np.testing.assert_array_equal(Dnew, Dref)
|
|
np.testing.assert_array_equal(Inew, Iref)
|
|
|
|
def test_hnsw_reset(self):
|
|
d = self.xb.shape[1]
|
|
index_flat = faiss.IndexFlat(d)
|
|
index_flat.add(self.xb)
|
|
self.assertEqual(index_flat.ntotal, self.xb.shape[0])
|
|
index_hnsw = faiss.IndexHNSW(index_flat)
|
|
index_hnsw.add(self.xb)
|
|
# * 2 because we add to storage twice. This is just for testing
|
|
# that storage gets cleared correctly.
|
|
self.assertEqual(index_hnsw.ntotal, self.xb.shape[0] * 2)
|
|
|
|
index_hnsw.reset()
|
|
|
|
self.assertEqual(index_flat.ntotal, 0)
|
|
self.assertEqual(index_hnsw.ntotal, 0)
|
|
|
|
|
|
class TestHNSWSimilarity(unittest.TestCase):
|
|
"""Coverage for HNSW variants built with a similarity metric
|
|
(METRIC_INNER_PRODUCT). These tests exercise the
|
|
HNSW::is_similarity dispatch in IndexHNSW{Flat,PQ,SQ,Cagra}: build,
|
|
search, and (for IndexHNSWFlat) save/load round trips that have to
|
|
re-derive `hnsw.is_similarity` from the persisted metric type.
|
|
"""
|
|
|
|
def setUp(self):
|
|
d = 32
|
|
nt = 1500
|
|
nb = 1500
|
|
nq = 200
|
|
(self.xt, self.xb, self.xq) = get_dataset_2(d, nt, nb, nq)
|
|
# Reference IP search results.
|
|
ref = faiss.IndexFlatIP(d)
|
|
ref.add(self.xb)
|
|
self.Dref_ip, self.Iref_ip = ref.search(self.xq, 1)
|
|
|
|
def test_hnsw_flat_IP_is_similarity_set(self):
|
|
d = self.xb.shape[1]
|
|
index = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT)
|
|
self.assertTrue(index.hnsw.is_similarity)
|
|
|
|
index_l2 = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_L2)
|
|
self.assertFalse(index_l2.hnsw.is_similarity)
|
|
|
|
def _check_quantized_IP(self, index, min_agreement):
|
|
"""Build (train if needed), search, then save / load and assert
|
|
that the round trip preserves ordering and `hnsw.is_similarity`.
|
|
The agreement floors are calibrated against the observed values
|
|
on the test dataset with a small (~5%) headroom for RNG noise;
|
|
they catch ordering regressions (which collapse agreement to
|
|
~0) and quantizer-quality regressions (which would shift
|
|
agreement by tens of points)."""
|
|
self.assertTrue(index.hnsw.is_similarity)
|
|
if not index.is_trained:
|
|
index.train(self.xt)
|
|
index.add(self.xb)
|
|
|
|
Dhnsw, Ihnsw = index.search(self.xq, 1)
|
|
agreement = (self.Iref_ip[:, 0] == Ihnsw[:, 0]).sum()
|
|
self.assertGreaterEqual(
|
|
agreement,
|
|
min_agreement,
|
|
f"{type(index).__name__} agreement={agreement} < {min_agreement}",
|
|
)
|
|
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
# `hnsw.is_similarity` is not serialized; index_read must
|
|
# re-derive it from the persisted metric_type.
|
|
self.assertTrue(index2.hnsw.is_similarity)
|
|
self.assertEqual(index2.metric_type, faiss.METRIC_INNER_PRODUCT)
|
|
D2, I2 = index2.search(self.xq, 1)
|
|
np.testing.assert_array_equal(I2, Ihnsw)
|
|
np.testing.assert_array_equal(D2, Dhnsw)
|
|
|
|
# Agreement floors below are calibrated against measured values
|
|
# (Flat: 199, PQ: 26-30, SQ: 195, Cagra: 200 over nq=200) with a
|
|
# small headroom for RNG noise.
|
|
|
|
def test_hnsw_flat_IP(self):
|
|
d = self.xb.shape[1]
|
|
self._check_quantized_IP(
|
|
faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT),
|
|
# observed ~199/200; headroom for OpenMP-parallel-build
|
|
# nondeterminism
|
|
min_agreement=180,
|
|
)
|
|
|
|
def test_hnsw_pq_IP(self):
|
|
d = self.xb.shape[1]
|
|
self._check_quantized_IP(
|
|
faiss.IndexHNSWPQ(d, 8, 16, 8, faiss.METRIC_INNER_PRODUCT),
|
|
min_agreement=20,
|
|
)
|
|
|
|
def test_hnsw_sq_IP(self):
|
|
d = self.xb.shape[1]
|
|
self._check_quantized_IP(
|
|
faiss.IndexHNSWSQ(
|
|
d,
|
|
faiss.ScalarQuantizer.QT_8bit,
|
|
16,
|
|
faiss.METRIC_INNER_PRODUCT,
|
|
),
|
|
min_agreement=190,
|
|
)
|
|
|
|
def test_hnsw_cagra_IP(self):
|
|
d = self.xb.shape[1]
|
|
self._check_quantized_IP(
|
|
faiss.IndexHNSWCagra(d, 16, faiss.METRIC_INNER_PRODUCT),
|
|
min_agreement=195,
|
|
)
|
|
|
|
def test_hnsw_cagra_IP_base_level_only(self):
|
|
"""IndexHNSWCagra.base_level_only takes a separate code path
|
|
(random entry-point selection + search_level_0) that also
|
|
needs to dispatch on is_similarity."""
|
|
d = self.xb.shape[1]
|
|
index = faiss.IndexHNSWCagra(d, 16, faiss.METRIC_INNER_PRODUCT)
|
|
index.add(self.xb)
|
|
index.base_level_only = True
|
|
index.num_base_level_search_entrypoints = 64
|
|
|
|
_, Ihnsw = index.search(self.xq, 1)
|
|
agreement = (self.Iref_ip[:, 0] == Ihnsw[:, 0]).sum()
|
|
# Measured value on this dataset is 200/200; leave headroom for
|
|
# the RNG-driven entry-point selection (this path is unseeded).
|
|
self.assertGreaterEqual(
|
|
agreement,
|
|
190,
|
|
f"Cagra base_only agreement={agreement}",
|
|
)
|
|
|
|
|
|
class TestHNSWNaN(unittest.TestCase):
|
|
"""Adding a vector with NaN to an IVF+HNSW index used to crash because
|
|
NaN distances corrupt the MinimaxHeap ordering in HNSW search. The fix
|
|
converts NaN to +inf in MinimaxHeap::push so the heap stays well-ordered.
|
|
"""
|
|
|
|
def test_add_nan_vector_to_ivf_hnsw(self):
|
|
d = 64
|
|
nt = 2000
|
|
nb = 1000
|
|
xt = np.random.default_rng(42).random((nt, d), dtype="float32")
|
|
xb = np.random.default_rng(43).random((nb, d), dtype="float32")
|
|
|
|
index = faiss.index_factory(d, "IVF256_HNSW32,SQ8")
|
|
index.train(xt)
|
|
index.add(xb)
|
|
|
|
# Create a vector with NaN in the first component
|
|
vec = np.zeros((1, d), dtype="float32")
|
|
vec[0, 0] = np.nan
|
|
|
|
# This should not crash
|
|
index.add(vec)
|
|
self.assertEqual(index.ntotal, nb + 1)
|
|
|
|
|
|
class Issue3684(unittest.TestCase):
|
|
|
|
def test_issue3684(self):
|
|
np.random.seed(1234) # For reproducibility
|
|
d = 256 # Example dimension
|
|
nb = 10 # Number of database vectors
|
|
nq = 2 # Number of query vectors
|
|
xb = np.random.random((nb, d)).astype("float32")
|
|
xq = np.random.random((nq, d)).astype("float32")
|
|
|
|
faiss.normalize_L2(xb) # Normalize both query and database vectors
|
|
faiss.normalize_L2(xq)
|
|
|
|
hnsw_index_ip = faiss.IndexHNSWFlat(256, 16, faiss.METRIC_INNER_PRODUCT)
|
|
hnsw_index_ip.hnsw.efConstruction = 512
|
|
hnsw_index_ip.hnsw.efSearch = 512
|
|
hnsw_index_ip.add(xb)
|
|
|
|
# test knn
|
|
D, I = hnsw_index_ip.search(xq, 10)
|
|
self.assertTrue(np.all(D[:, :-1] >= D[:, 1:]))
|
|
|
|
# test range search
|
|
radius = 0.74 # Cosine similarity threshold
|
|
lims, D, I = hnsw_index_ip.range_search(xq, radius)
|
|
self.assertTrue(np.all(D >= radius))
|
|
|
|
|
|
class TestNSG(unittest.TestCase):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
unittest.TestCase.__init__(self, *args, **kwargs)
|
|
d = 32
|
|
nt = 0
|
|
nb = 1500
|
|
nq = 500
|
|
self.GK = 32
|
|
|
|
_, self.xb, self.xq = get_dataset_2(d, nt, nb, nq)
|
|
|
|
def make_knn_graph(self, metric):
|
|
n = self.xb.shape[0]
|
|
d = self.xb.shape[1]
|
|
index = faiss.IndexFlat(d, metric)
|
|
index.add(self.xb)
|
|
_, I = index.search(self.xb, self.GK + 1)
|
|
knn_graph = np.zeros((n, self.GK), dtype=np.int64)
|
|
|
|
# For the inner product distance, the distance between a vector and
|
|
# itself may not be the smallest, so it is not guaranteed that I[:, 0]
|
|
# is the query itself.
|
|
for i in range(n):
|
|
cnt = 0
|
|
for j in range(self.GK + 1):
|
|
if I[i, j] != i:
|
|
knn_graph[i, cnt] = I[i, j]
|
|
cnt += 1
|
|
if cnt == self.GK:
|
|
break
|
|
return knn_graph
|
|
|
|
def subtest_io_and_clone(self, index, Dnsg, Insg):
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
|
|
Dnsg2, Insg2 = index2.search(self.xq, 1)
|
|
np.testing.assert_array_equal(Dnsg2, Dnsg)
|
|
np.testing.assert_array_equal(Insg2, Insg)
|
|
|
|
# Verify deserialized index is serializable again
|
|
index2b = faiss.deserialize_index(faiss.serialize_index(index2))
|
|
Dnsg2b, Insg2b = index2b.search(self.xq, 1)
|
|
np.testing.assert_array_equal(Dnsg2b, Dnsg)
|
|
np.testing.assert_array_equal(Insg2b, Insg)
|
|
|
|
# also test clone
|
|
index3 = faiss.clone_index(index)
|
|
Dnsg3, Insg3 = index3.search(self.xq, 1)
|
|
np.testing.assert_array_equal(Dnsg3, Dnsg)
|
|
np.testing.assert_array_equal(Insg3, Insg)
|
|
|
|
def subtest_connectivity(self, index, nb):
|
|
vt = faiss.VisitedTableVector(nb)
|
|
count = index.nsg.dfs(vt, index.nsg.enterpoint, 0)
|
|
self.assertEqual(count, nb)
|
|
|
|
def subtest_add(self, build_type, thresh, metric=faiss.METRIC_L2):
|
|
d = self.xq.shape[1]
|
|
|
|
flat_index = faiss.IndexFlat(d, metric)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, 1)
|
|
|
|
index = faiss.IndexNSGFlat(d, 16, metric)
|
|
index.verbose = True
|
|
index.build_type = build_type
|
|
index.GK = self.GK
|
|
index.add(self.xb)
|
|
Dnsg, Insg = index.search(self.xq, 1)
|
|
|
|
recalls = (Iref == Insg).sum()
|
|
self.assertGreaterEqual(recalls, thresh)
|
|
self.subtest_connectivity(index, self.xb.shape[0])
|
|
self.subtest_io_and_clone(index, Dnsg, Insg)
|
|
|
|
def subtest_build(self, knn_graph, thresh, metric=faiss.METRIC_L2):
|
|
d = self.xq.shape[1]
|
|
|
|
flat_index = faiss.IndexFlat(d, metric)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, 1)
|
|
|
|
index = faiss.IndexNSGFlat(d, 16, metric)
|
|
index.verbose = True
|
|
|
|
index.build(self.xb, knn_graph)
|
|
Dnsg, Insg = index.search(self.xq, 1)
|
|
|
|
recalls = (Iref == Insg).sum()
|
|
self.assertGreaterEqual(recalls, thresh)
|
|
self.subtest_connectivity(index, self.xb.shape[0])
|
|
|
|
def test_add_bruteforce_L2(self):
|
|
self.subtest_add(0, 475, faiss.METRIC_L2)
|
|
|
|
def test_add_nndescent_L2(self):
|
|
self.subtest_add(1, 475, faiss.METRIC_L2)
|
|
|
|
def test_add_bruteforce_IP(self):
|
|
self.subtest_add(0, 480, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_add_nndescent_IP(self):
|
|
self.subtest_add(1, 480, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_build_L2(self):
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_L2)
|
|
self.subtest_build(knn_graph, 475, faiss.METRIC_L2)
|
|
|
|
def test_build_IP(self):
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_INNER_PRODUCT)
|
|
self.subtest_build(knn_graph, 480, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_build_invalid_knng(self):
|
|
"""Make some invalid entries in the input knn graph.
|
|
|
|
It would cause a warning but IndexNSG should be able
|
|
to handle this.
|
|
"""
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_L2)
|
|
knn_graph[:100, 5] = -111
|
|
self.subtest_build(knn_graph, 475, faiss.METRIC_L2)
|
|
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_INNER_PRODUCT)
|
|
knn_graph[:100, 5] = -111
|
|
self.subtest_build(knn_graph, 480, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def test_reset(self):
|
|
"""test IndexNSG.reset()"""
|
|
d = self.xq.shape[1]
|
|
|
|
metric = faiss.METRIC_L2
|
|
flat_index = faiss.IndexFlat(d, metric)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, 1)
|
|
|
|
index = faiss.IndexNSGFlat(d, 16)
|
|
index.verbose = True
|
|
index.GK = 32
|
|
|
|
index.add(self.xb)
|
|
Dnsg, Insg = index.search(self.xq, 1)
|
|
recalls = (Iref == Insg).sum()
|
|
self.assertGreaterEqual(recalls, 475)
|
|
self.subtest_connectivity(index, self.xb.shape[0])
|
|
|
|
index.reset()
|
|
index.add(self.xb)
|
|
Dnsg, Insg = index.search(self.xq, 1)
|
|
recalls = (Iref == Insg).sum()
|
|
self.assertGreaterEqual(recalls, 475)
|
|
self.subtest_connectivity(index, self.xb.shape[0])
|
|
|
|
def test_order(self):
|
|
"""make sure that output results are sorted"""
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNSGFlat(d, 32)
|
|
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
|
|
k = 10
|
|
nq = self.xq.shape[0]
|
|
D, _ = index.search(self.xq, k)
|
|
|
|
indices = np.argsort(D, axis=1)
|
|
gt = np.arange(0, k)[np.newaxis, :] # [1, k]
|
|
gt = np.repeat(gt, nq, axis=0) # [nq, k]
|
|
np.testing.assert_array_equal(indices, gt)
|
|
|
|
def test_nsg_supports_pre_built_knn_graph(self):
|
|
"""Test IndexNSGBuild"""
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_L2)
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNSGFlat(d, 16)
|
|
index.build(self.xb, knn_graph)
|
|
index.search(self.xq, k=1)
|
|
self.assertTrue(index.is_built)
|
|
|
|
def test_nsg_with_pre_built_knn_graph_throws_when_rebuilding_via_add(self):
|
|
"""Test IndexNSGBuild"""
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_L2)
|
|
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNSGFlat(d, 16)
|
|
index.build(self.xb, knn_graph)
|
|
index.search(self.xq, k=1)
|
|
self.assertTrue(index.is_built)
|
|
|
|
index.GK = 32
|
|
index.train(self.xb)
|
|
with self.assertRaises(RuntimeError) as context:
|
|
index.add(self.xb)
|
|
|
|
self.assertIn(
|
|
"NSG does not support incremental addition", str(context.exception)
|
|
)
|
|
|
|
def test_nsg_rebuild_throws_with_pre_built_knn_graph(self):
|
|
"""Test IndexNSGBuild"""
|
|
knn_graph = self.make_knn_graph(faiss.METRIC_L2)
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNSGFlat(d, 16)
|
|
index.build(self.xb, knn_graph)
|
|
index.search(self.xq, k=1)
|
|
|
|
with self.assertRaises(RuntimeError) as context:
|
|
index.build(self.xb, knn_graph)
|
|
|
|
self.assertIn("The IndexNSG is already built", str(context.exception))
|
|
|
|
def test_nsg_pq(self):
|
|
"""Test IndexNSGPQ"""
|
|
d = self.xq.shape[1]
|
|
R, pq_M = 32, 4
|
|
index = faiss.index_factory(d, f"NSG{R}_PQ{pq_M}np")
|
|
self.assertIsInstance(index, faiss.IndexNSGPQ)
|
|
idxpq = faiss.downcast_index(index.storage)
|
|
self.assertEqual(index.nsg.R, R)
|
|
self.assertEqual(idxpq.pq.M, pq_M)
|
|
|
|
flat_index = faiss.IndexFlat(d)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, k=1)
|
|
|
|
index.GK = 32
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
D, I = index.search(self.xq, k=1)
|
|
|
|
# test accuracy
|
|
recalls = (Iref == I).sum()
|
|
# observed ~193/500; PQ4 is coarse -- headroom for OpenMP
|
|
# FP-reduction / tie-break nondeterminism in NSG build + PQ assign
|
|
self.assertGreaterEqual(recalls, 180)
|
|
|
|
# test I/O
|
|
self.subtest_io_and_clone(index, D, I)
|
|
|
|
def test_nsg_sq(self):
|
|
"""Test IndexNSGSQ"""
|
|
d = self.xq.shape[1]
|
|
R = 32
|
|
index = faiss.index_factory(d, f"NSG{R}_SQ8")
|
|
self.assertIsInstance(index, faiss.IndexNSGSQ)
|
|
idxsq = faiss.downcast_index(index.storage)
|
|
self.assertEqual(index.nsg.R, R)
|
|
self.assertEqual(idxsq.sq.qtype, faiss.ScalarQuantizer.QT_8bit)
|
|
|
|
flat_index = faiss.IndexFlat(d)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, k=1)
|
|
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
D, I = index.search(self.xq, k=1)
|
|
|
|
# test accuracy
|
|
recalls = (Iref == I).sum()
|
|
# nominal ~411/500; loosened for OpenMP graph-build nondeterminism
|
|
self.assertGreaterEqual(recalls, 395)
|
|
|
|
# test I/O
|
|
self.subtest_io_and_clone(index, D, I)
|
|
|
|
|
|
class TestNNDescent(unittest.TestCase):
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
unittest.TestCase.__init__(self, *args, **kwargs)
|
|
d = 32
|
|
nt = 0
|
|
nb = 1500
|
|
nq = 500
|
|
self.GK = 32
|
|
|
|
_, self.xb, self.xq = get_dataset_2(d, nt, nb, nq)
|
|
|
|
def test_nndescentflat(self):
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNNDescentFlat(d, 32)
|
|
index.nndescent.search_L = 8
|
|
|
|
flat_index = faiss.IndexFlat(d)
|
|
flat_index.add(self.xb)
|
|
Dref, Iref = flat_index.search(self.xq, k=1)
|
|
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
D, I = index.search(self.xq, k=1)
|
|
|
|
# test accuracy
|
|
recalls = (Iref == I).sum()
|
|
self.assertGreaterEqual(recalls, 450) # 462
|
|
|
|
# do some IO tests
|
|
index2 = faiss.deserialize_index(faiss.serialize_index(index))
|
|
|
|
D2, I2 = index2.search(self.xq, 1)
|
|
np.testing.assert_array_equal(D2, D)
|
|
np.testing.assert_array_equal(I2, I)
|
|
|
|
# Verify deserialized index is serializable again
|
|
index2b = faiss.deserialize_index(faiss.serialize_index(index2))
|
|
D2b, I2b = index2b.search(self.xq, 1)
|
|
np.testing.assert_array_equal(D2b, D)
|
|
np.testing.assert_array_equal(I2b, I)
|
|
|
|
# also test clone
|
|
index3 = faiss.clone_index(index)
|
|
D3, I3 = index3.search(self.xq, 1)
|
|
np.testing.assert_array_equal(D3, D)
|
|
np.testing.assert_array_equal(I3, I)
|
|
|
|
def test_order(self):
|
|
"""make sure that output results are sorted"""
|
|
d = self.xq.shape[1]
|
|
index = faiss.IndexNNDescentFlat(d, 32)
|
|
|
|
index.train(self.xb)
|
|
index.add(self.xb)
|
|
|
|
k = 10
|
|
nq = self.xq.shape[0]
|
|
D, _ = index.search(self.xq, k)
|
|
|
|
indices = np.argsort(D, axis=1)
|
|
gt = np.arange(0, k)[np.newaxis, :] # [1, k]
|
|
gt = np.repeat(gt, nq, axis=0) # [nq, k]
|
|
np.testing.assert_array_equal(indices, gt)
|
|
|
|
|
|
class TestNNDescentGenRandom(unittest.TestCase):
|
|
"""Regression tests for gen_random edge cases in NNDescent."""
|
|
|
|
def test_search_L_equals_ntotal(self):
|
|
"""gen_random(size, N) crashed with division by zero when size == N.
|
|
|
|
In search(), L_2 = max(search_L, topk). When search_L >= ntotal,
|
|
gen_random is called with size == N, causing rng() % 0.
|
|
"""
|
|
d = 32
|
|
nb = 200 # just above NUM_EVAL_POINTS=100
|
|
xb = np.random.default_rng(42).random((nb, d)).astype("float32")
|
|
xq = np.random.default_rng(43).random((10, d)).astype("float32")
|
|
|
|
index = faiss.IndexNNDescentFlat(d, 32)
|
|
index.nndescent.search_L = nb # triggers gen_random(size=nb, N=nb)
|
|
index.train(xb)
|
|
index.add(xb)
|
|
|
|
# This crashed with division by zero before the fix
|
|
D, I = index.search(xq, k=1)
|
|
self.assertEqual(I.shape, (xq.shape[0], 1))
|
|
self.assertEqual(D.shape, (xq.shape[0], 1))
|
|
|
|
|
|
class TestNNDescentKNNG(unittest.TestCase):
|
|
|
|
def test_knng_L2(self):
|
|
self.subtest(32, 10, faiss.METRIC_L2)
|
|
|
|
def test_knng_IP(self):
|
|
self.subtest(32, 10, faiss.METRIC_INNER_PRODUCT)
|
|
|
|
def subtest(self, d, K, metric):
|
|
nb = 1000
|
|
_, xb, _ = get_dataset_2(d, 0, nb, 0)
|
|
|
|
_, knn = faiss.knn(xb, xb, K + 1, metric)
|
|
knn = knn[:, 1:]
|
|
|
|
index = faiss.IndexNNDescentFlat(d, K, metric)
|
|
index.nndescent.S = 10
|
|
index.nndescent.R = 32
|
|
index.nndescent.L = K + 20
|
|
index.nndescent.iter = 5
|
|
index.verbose = True
|
|
|
|
index.add(xb)
|
|
graph = index.nndescent.final_graph
|
|
graph = faiss.vector_to_array(graph)
|
|
graph = graph.reshape(nb, K)
|
|
|
|
recalls = 0
|
|
for i in range(nb):
|
|
for j in range(K):
|
|
for k in range(K):
|
|
if graph[i, j] == knn[i, k]:
|
|
recalls += 1
|
|
break
|
|
recall = 1.0 * recalls / (nb * K)
|
|
self.assertGreater(recall, 0.99)
|
|
|
|
def test_small_nndescent(self):
|
|
"""building a too small graph used to crash, make sure it raises
|
|
an exception instead.
|
|
TODO: build the exact knn graph for small cases
|
|
"""
|
|
d = 32
|
|
K = 10
|
|
index = faiss.IndexNNDescentFlat(d, K, faiss.METRIC_L2)
|
|
index.nndescent.S = 10
|
|
index.nndescent.R = 32
|
|
index.nndescent.L = K + 20
|
|
index.nndescent.iter = 5
|
|
index.verbose = True
|
|
|
|
xb = np.zeros((78, d), dtype="float32")
|
|
self.assertRaises(RuntimeError, index.add, xb)
|