363 lines
12 KiB
Python
363 lines
12 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 faiss
|
|
import time
|
|
import numpy as np
|
|
|
|
import logging
|
|
|
|
LOG = logging.getLogger(__name__)
|
|
|
|
def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2):
|
|
"""Computes the exact KNN search results for a dataset that possibly
|
|
does not fit in RAM but for which we have an iterator that
|
|
returns it block by block.
|
|
"""
|
|
LOG.info("knn_ground_truth queries size %s k=%d" % (xq.shape, k))
|
|
t0 = time.time()
|
|
nq, d = xq.shape
|
|
keep_max = faiss.is_similarity_metric(metric_type)
|
|
rh = faiss.ResultHeap(nq, k, keep_max=keep_max)
|
|
|
|
index = faiss.IndexFlat(d, metric_type)
|
|
if faiss.get_num_gpus():
|
|
LOG.info('running on %d GPUs' % faiss.get_num_gpus())
|
|
index = faiss.index_cpu_to_all_gpus(index)
|
|
|
|
# compute ground-truth by blocks, and add to heaps
|
|
i0 = 0
|
|
for xbi in db_iterator:
|
|
ni = xbi.shape[0]
|
|
index.add(xbi)
|
|
D, I = index.search(xq, k)
|
|
I += i0
|
|
rh.add_result(D, I)
|
|
index.reset()
|
|
i0 += ni
|
|
LOG.info("%d db elements, %.3f s" % (i0, time.time() - t0))
|
|
|
|
rh.finalize()
|
|
LOG.info("GT time: %.3f s (%d vectors)" % (time.time() - t0, i0))
|
|
|
|
return rh.D, rh.I
|
|
|
|
# knn function used to be here
|
|
knn = faiss.knn
|
|
|
|
|
|
|
|
|
|
def range_search_gpu(xq, r2, index_gpu, index_cpu, gpu_k=1024):
|
|
"""GPU does not support range search, so we emulate it with
|
|
knn search + fallback to CPU index.
|
|
|
|
The index_cpu can either be:
|
|
- a CPU index that supports range search
|
|
- a numpy table, that will be used to construct a Flat index if needed.
|
|
- None. In that case, at most gpu_k results will be returned
|
|
"""
|
|
nq, d = xq.shape
|
|
is_binary_index = isinstance(index_gpu, faiss.IndexBinary)
|
|
keep_max = faiss.is_similarity_metric(index_gpu.metric_type)
|
|
r2 = int(r2) if is_binary_index else float(r2)
|
|
k = min(index_gpu.ntotal, gpu_k)
|
|
LOG.debug(
|
|
f"GPU search {nq} queries with {k=:} {is_binary_index=:} {keep_max=:}")
|
|
t0 = time.time()
|
|
D, I = index_gpu.search(xq, k)
|
|
t1 = time.time() - t0
|
|
if is_binary_index:
|
|
assert d * 8 < 32768 # let's compact the distance matrix
|
|
D = D.astype('int16')
|
|
t2 = 0
|
|
lim_remain = None
|
|
if index_cpu is not None:
|
|
if not keep_max:
|
|
mask = D[:, k - 1] < r2
|
|
else:
|
|
mask = D[:, k - 1] > r2
|
|
if mask.sum() > 0:
|
|
LOG.debug("CPU search remain %d" % mask.sum())
|
|
t0 = time.time()
|
|
if isinstance(index_cpu, np.ndarray):
|
|
# then it in fact an array that we have to make flat
|
|
xb = index_cpu
|
|
if is_binary_index:
|
|
index_cpu = faiss.IndexBinaryFlat(d * 8)
|
|
else:
|
|
index_cpu = faiss.IndexFlat(d, index_gpu.metric_type)
|
|
index_cpu.add(xb)
|
|
lim_remain, D_remain, I_remain = index_cpu.range_search(xq[mask], r2)
|
|
if is_binary_index:
|
|
D_remain = D_remain.astype('int16')
|
|
t2 = time.time() - t0
|
|
LOG.debug("combine")
|
|
t0 = time.time()
|
|
|
|
CombinerRangeKNN = (
|
|
faiss.CombinerRangeKNNint16 if is_binary_index else
|
|
faiss.CombinerRangeKNNfloat
|
|
)
|
|
|
|
combiner = CombinerRangeKNN(nq, k, r2, keep_max)
|
|
if True:
|
|
sp = faiss.swig_ptr
|
|
combiner.I = sp(I)
|
|
combiner.D = sp(D)
|
|
# combiner.set_knn_result(sp(I), sp(D))
|
|
if lim_remain is not None:
|
|
combiner.mask = sp(mask)
|
|
combiner.D_remain = sp(D_remain)
|
|
combiner.lim_remain = sp(lim_remain.view("int64"))
|
|
combiner.I_remain = sp(I_remain)
|
|
# combiner.set_range_result(sp(mask), sp(lim_remain.view("int64")), sp(D_remain), sp(I_remain))
|
|
L_res = np.empty(nq + 1, dtype='int64')
|
|
combiner.compute_sizes(sp(L_res))
|
|
nres = L_res[-1]
|
|
D_res = np.empty(nres, dtype=D.dtype)
|
|
I_res = np.empty(nres, dtype='int64')
|
|
combiner.write_result(sp(D_res), sp(I_res))
|
|
else:
|
|
D_res, I_res = [], []
|
|
nr = 0
|
|
for i in range(nq):
|
|
if not mask[i]:
|
|
if index_gpu.metric_type == faiss.METRIC_L2:
|
|
nv = (D[i, :] < r2).sum()
|
|
else:
|
|
nv = (D[i, :] > r2).sum()
|
|
D_res.append(D[i, :nv])
|
|
I_res.append(I[i, :nv])
|
|
else:
|
|
l0, l1 = lim_remain[nr], lim_remain[nr + 1]
|
|
D_res.append(D_remain[l0:l1])
|
|
I_res.append(I_remain[l0:l1])
|
|
nr += 1
|
|
L_res = np.cumsum([0] + [len(di) for di in D_res])
|
|
D_res = np.hstack(D_res)
|
|
I_res = np.hstack(I_res)
|
|
t3 = time.time() - t0
|
|
LOG.debug(f"times {t1:.3f}s {t2:.3f}s {t3:.3f}s")
|
|
return L_res, D_res, I_res
|
|
|
|
|
|
def range_ground_truth(xq, db_iterator, threshold, metric_type=faiss.METRIC_L2,
|
|
shard=False, ngpu=-1):
|
|
"""Computes the range-search search results for a dataset that possibly
|
|
does not fit in RAM but for which we have an iterator that
|
|
returns it block by block.
|
|
"""
|
|
nq, d = xq.shape
|
|
t0 = time.time()
|
|
xq = np.ascontiguousarray(xq, dtype='float32')
|
|
|
|
index = faiss.IndexFlat(d, metric_type)
|
|
if ngpu == -1:
|
|
ngpu = faiss.get_num_gpus()
|
|
if ngpu:
|
|
LOG.info('running on %d GPUs' % ngpu)
|
|
co = faiss.GpuMultipleClonerOptions()
|
|
co.shard = shard
|
|
index_gpu = faiss.index_cpu_to_all_gpus(index, co=co, ngpu=ngpu)
|
|
|
|
# compute ground-truth by blocks
|
|
i0 = 0
|
|
D = [[] for _i in range(nq)]
|
|
I = [[] for _i in range(nq)]
|
|
for xbi in db_iterator:
|
|
ni = xbi.shape[0]
|
|
if ngpu > 0:
|
|
index_gpu.add(xbi)
|
|
lims_i, Di, Ii = range_search_gpu(xq, threshold, index_gpu, xbi)
|
|
index_gpu.reset()
|
|
else:
|
|
index.add(xbi)
|
|
lims_i, Di, Ii = index.range_search(xq, threshold)
|
|
index.reset()
|
|
Ii += i0
|
|
for j in range(nq):
|
|
l0, l1 = lims_i[j], lims_i[j + 1]
|
|
if l1 > l0:
|
|
D[j].append(Di[l0:l1])
|
|
I[j].append(Ii[l0:l1])
|
|
i0 += ni
|
|
LOG.info("%d db elements, %.3f s" % (i0, time.time() - t0))
|
|
|
|
empty_I = np.zeros(0, dtype='int64')
|
|
empty_D = np.zeros(0, dtype='float32')
|
|
# import pdb; pdb.set_trace()
|
|
D = [(np.hstack(i) if i != [] else empty_D) for i in D]
|
|
I = [(np.hstack(i) if i != [] else empty_I) for i in I]
|
|
sizes = [len(i) for i in I]
|
|
assert len(sizes) == nq
|
|
lims = np.zeros(nq + 1, dtype="uint64")
|
|
lims[1:] = np.cumsum(sizes)
|
|
return lims, np.hstack(D), np.hstack(I)
|
|
|
|
|
|
def threshold_radius_nres(nres, dis, ids, thresh, keep_max=False):
|
|
""" select a set of results """
|
|
if keep_max:
|
|
mask = dis > thresh
|
|
else:
|
|
mask = dis < thresh
|
|
new_nres = np.zeros_like(nres)
|
|
o = 0
|
|
for i, nr in enumerate(nres):
|
|
nr = int(nr) # avoid issues with int64 + uint64
|
|
new_nres[i] = mask[o:o + nr].sum()
|
|
o += nr
|
|
return new_nres, dis[mask], ids[mask]
|
|
|
|
|
|
def threshold_radius(lims, dis, ids, thresh, keep_max=False):
|
|
""" restrict range-search results to those below a given radius """
|
|
if keep_max:
|
|
mask = dis > thresh
|
|
else:
|
|
mask = dis < thresh
|
|
new_lims = np.zeros_like(lims)
|
|
n = len(lims) - 1
|
|
for i in range(n):
|
|
l0, l1 = lims[i], lims[i + 1]
|
|
new_lims[i + 1] = new_lims[i] + mask[l0:l1].sum()
|
|
return new_lims, dis[mask], ids[mask]
|
|
|
|
|
|
def apply_maxres(res_batches, target_nres, keep_max=False):
|
|
"""find radius that reduces number of results to target_nres, and
|
|
applies it in-place to the result batches used in
|
|
range_search_max_results"""
|
|
alldis = np.hstack([dis for _, dis, _ in res_batches])
|
|
assert len(alldis) > target_nres
|
|
if keep_max:
|
|
alldis.partition(len(alldis) - target_nres - 1)
|
|
radius = alldis[-1 - target_nres]
|
|
else:
|
|
alldis.partition(target_nres)
|
|
radius = alldis[target_nres]
|
|
|
|
if alldis.dtype == 'float32':
|
|
radius = float(radius)
|
|
else:
|
|
radius = int(radius)
|
|
LOG.debug(' setting radius to %s' % radius)
|
|
totres = 0
|
|
for i, (nres, dis, ids) in enumerate(res_batches):
|
|
nres, dis, ids = threshold_radius_nres(
|
|
nres, dis, ids, radius, keep_max=keep_max)
|
|
totres += len(dis)
|
|
res_batches[i] = nres, dis, ids
|
|
LOG.debug(' updated previous results, new nb results %d' % totres)
|
|
return radius, totres
|
|
|
|
|
|
def range_search_max_results(index, query_iterator, radius,
|
|
max_results=None, min_results=None,
|
|
shard=False, ngpu=0, clip_to_min=False):
|
|
"""Performs a range search with many queries (given by an iterator)
|
|
and adjusts the threshold on-the-fly so that the total results
|
|
table does not grow larger than max_results.
|
|
|
|
If ngpu != 0, the function moves the index to this many GPUs to
|
|
speed up search.
|
|
"""
|
|
# TODO: all result manipulations are in python, should move to C++ if perf
|
|
# critical
|
|
is_binary_index = isinstance(index, faiss.IndexBinary)
|
|
|
|
if min_results is None:
|
|
assert max_results is not None
|
|
min_results = int(0.8 * max_results)
|
|
|
|
if max_results is None:
|
|
assert min_results is not None
|
|
max_results = int(min_results * 1.5)
|
|
|
|
if ngpu == -1:
|
|
ngpu = faiss.get_num_gpus()
|
|
|
|
if ngpu:
|
|
LOG.info('running on %d GPUs' % ngpu)
|
|
co = faiss.GpuMultipleClonerOptions()
|
|
co.shard = shard
|
|
index_gpu = faiss.index_cpu_to_all_gpus(index, co=co, ngpu=ngpu)
|
|
else:
|
|
index_gpu = None
|
|
|
|
t_start = time.time()
|
|
t_search = t_post_process = 0
|
|
qtot = totres = raw_totres = 0
|
|
res_batches = []
|
|
|
|
for xqi in query_iterator:
|
|
t0 = time.time()
|
|
LOG.debug(f"searching {len(xqi)} vectors")
|
|
if index_gpu:
|
|
lims_i, Di, Ii = range_search_gpu(xqi, radius, index_gpu, index)
|
|
else:
|
|
lims_i, Di, Ii = index.range_search(xqi, radius)
|
|
|
|
nres_i = lims_i[1:] - lims_i[:-1]
|
|
raw_totres += len(Di)
|
|
qtot += len(xqi)
|
|
|
|
t1 = time.time()
|
|
if is_binary_index:
|
|
# weird Faiss quirk that returns floats for Hamming distances
|
|
Di = Di.astype('int16')
|
|
|
|
totres += len(Di)
|
|
res_batches.append((nres_i, Di, Ii))
|
|
|
|
if max_results is not None and totres > max_results:
|
|
LOG.info('too many results %d > %d, scaling back radius' %
|
|
(totres, max_results))
|
|
radius, totres = apply_maxres(
|
|
res_batches, min_results,
|
|
keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT
|
|
)
|
|
t2 = time.time()
|
|
t_search += t1 - t0
|
|
t_post_process += t2 - t1
|
|
LOG.debug(' [%.3f s] %d queries done, %d results' % (
|
|
time.time() - t_start, qtot, totres))
|
|
|
|
LOG.info(
|
|
'search done in %.3f s + %.3f s, total %d results, end threshold %g' % (
|
|
t_search, t_post_process, totres, radius)
|
|
)
|
|
|
|
if clip_to_min and totres > min_results:
|
|
radius, totres = apply_maxres(
|
|
res_batches, min_results,
|
|
keep_max=index.metric_type == faiss.METRIC_INNER_PRODUCT
|
|
)
|
|
|
|
nres = np.hstack([nres_i for nres_i, dis_i, ids_i in res_batches])
|
|
dis = np.hstack([dis_i for nres_i, dis_i, ids_i in res_batches])
|
|
ids = np.hstack([ids_i for nres_i, dis_i, ids_i in res_batches])
|
|
|
|
lims = np.zeros(len(nres) + 1, dtype='uint64')
|
|
lims[1:] = np.cumsum(nres)
|
|
|
|
return radius, lims, dis, ids
|
|
|
|
|
|
def exponential_query_iterator(xq, start_bs=32, max_bs=20000):
|
|
""" produces batches of progressively increasing sizes. This is useful to
|
|
adjust the search radius progressively without overflowing with
|
|
intermediate results """
|
|
nq = len(xq)
|
|
bs = start_bs
|
|
i = 0
|
|
while i < nq:
|
|
xqi = xq[i:i + bs]
|
|
yield xqi
|
|
if bs < max_bs:
|
|
bs *= 2
|
|
i += len(xqi)
|