facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 02:13:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

more formatting (#4568)

Browse Source 
Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4568

just formatting lints on Faiss

Reviewed By: junjieqi

Differential Revision: D81286229

fbshipit-source-id: 208efcf07b92a6dec9f78056a07f0c5bf7199aea
This commit is contained in:
Tia Lim
2025-09-05 00:54:43 -07:00
committed by Facebook GitHub Bot
parent 1ed2611c03
commit 3671c61af1
9 changed files with 220 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchs/bench_all_ivf/bench_all_ivf.py
+1
View File
@@ -400,6 +400,7 @@ def run_search(args, ds, index, res):
# Driver function
######################################################
def main():
parser = argparse.ArgumentParser()
benchs/bench_all_ivf/cmp_with_scann.py
+17 -12
View File
@@ -55,7 +55,8 @@ def main():
aa('--download', default=False, action="store_true")
aa('--lib', default='faiss', help='library to use (faiss or scann)')
aa('--thenscann', default=False, action="store_true")
aa('--base_dir', default='/checkpoint/matthijs/faiss_improvements/cmp_ivf_scan_2')
aa('--base_dir',
default='/checkpoint/matthijs/faiss_improvements/cmp_ivf_scan_2')
group = parser.add_argument_group('searching')
aa('--k', default=10, type=int, help='nb of nearest neighbors')
@@ -83,11 +84,11 @@ def main():
print(ds)
# store for SCANN
os.system(f"rm -rf {cache_dir}; mkdir -p {cache_dir}")
tosave = dict(
xb = ds.get_database(),
xq = ds.get_queries(),
gt = ds.get_groundtruth()
)
tosave = {
"xb": ds.get_database(),
"xq": ds.get_queries(),
"gt": ds.get_groundtruth(),
}
for name, v in tosave.items():
fname = cache_dir + "/" + name + ".npy"
print("save", fname)
@@ -137,13 +138,15 @@ def main():
if os.path.exists(scann_dir + "/scann_config.pb"):
searcher = scann_ops_pybind.load_searcher(scann_dir)
else:
searcher = scann_make_index(xb, name1_to_name2[distance_measure], scann_dir, 0)
searcher = scann_make_index(
xb, name1_to_name2[distance_measure], scann_dir, 0)
scann_dir = cache_dir + "/scann1.1.1_serialized_reorder"
if os.path.exists(scann_dir + "/scann_config.pb"):
searcher_reo = scann_ops_pybind.load_searcher(scann_dir)
else:
searcher_reo = scann_make_index(xb, name1_to_name2[distance_measure], scann_dir, 100)
searcher_reo = scann_make_index(
xb, name1_to_name2[distance_measure], scann_dir, 100)
scann_eval_search(
searcher, searcher_reo,
@@ -180,8 +183,10 @@ def scann_make_index(xb, distance_measure, scann_dir, reorder_k):
else:
thr = 0
k = 10
sb = scann.scann_ops_pybind.builder(xb, k, distance_measure)
sb = sb.tree(num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
sb = scann.scann_ops_pybind.builder(
xb, k, distance_measure)
sb = sb.tree(
num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
sb = sb.score_ah(2, anisotropic_quantization_threshold=thr)
if reorder_k > 0:
@@ -198,6 +203,7 @@ def scann_make_index(xb, distance_measure, scann_dir, reorder_k):
searcher.serialize(scann_dir)
return searcher
def scann_eval_search(
searcher, searcher_reo,
xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt,
@@ -235,8 +241,6 @@ def scann_eval_search(
eval_inters(header, I, gt, times)
###############################################################
# Faiss
###############################################################
@@ -261,6 +265,7 @@ def faiss_make_index(xb, metric_type, fname):
return index
def faiss_eval_search(
index, xq, xb, nprobe_tab, pre_reorder_k_tab,
k, gt, nrun, measure
benchs/bench_gpu_1bn.py
+37 -22
View File
@@ -90,22 +90,38 @@ args = sys.argv[1:]
while args:
a = args.pop(0)
if a == '-h': usage()
elif a == '-ngpu': ngpu = int(args.pop(0))
elif a == '-R': replicas = int(args.pop(0))
elif a == '-noptables': use_precomputed_tables = False
elif a == '-abs': add_batch_size = int(args.pop(0))
elif a == '-qbs': query_batch_size = int(args.pop(0))
elif a == '-nnn': nnn = int(args.pop(0))
elif a == '-tempmem': tempmem = int(args.pop(0))
elif a == '-nocache': use_cache = False
elif a == '-knngraph': knngraph = True
elif a == '-altadd': altadd = True
elif a == '-float16': use_float16 = True
elif a == '-nprobe': nprobes = [int(x) for x in args.pop(0).split(',')]
elif a == '-max_add': max_add = int(args.pop(0))
elif not dbname: dbname = a
elif not index_key: index_key = a
if a == '-h':
usage()
elif a == '-ngpu':
ngpu = int(args.pop(0))
elif a == '-R':
replicas = int(args.pop(0))
elif a == '-noptables':
use_precomputed_tables = False
elif a == '-abs':
add_batch_size = int(args.pop(0))
elif a == '-qbs':
query_batch_size = int(args.pop(0))
elif a == '-nnn':
nnn = int(args.pop(0))
elif a == '-tempmem':
tempmem = int(args.pop(0))
elif a == '-nocache':
use_cache = False
elif a == '-knngraph':
knngraph = True
elif a == '-altadd':
altadd = True
elif a == '-float16':
use_float16 = True
elif a == '-nprobe':
nprobes = [int(x) for x in args.pop(0).split(',')]
elif a == '-max_add':
max_add = int(args.pop(0))
elif not dbname:
dbname = a
elif not index_key:
index_key = a
else:
print("argument %s unknown" % a, file=sys.stderr)
sys.exit(1)
@@ -239,7 +255,6 @@ print("sizes: B %s Q %s T %s gt %s" % (
gt_I.shape if gt_I is not None else None))
#################################################################
# Parse index_key and set cache files
#
@@ -305,7 +320,7 @@ print("preparing resources for %d GPUs" % ngpu)
gpu_resources = []
for i in range(ngpu):
for _ in range(ngpu):
res = faiss.StandardGpuResources()
if tempmem >= 0:
res.setTempMemory(tempmem)
@@ -608,7 +623,6 @@ def compute_populated_index_2(preproc):
return None, indexall
def get_populated_index(preproc):
if not index_cachefile or not os.path.exists(index_cachefile):
@@ -644,7 +658,7 @@ def get_populated_index(preproc):
index = gpu_index
else:
del gpu_index # We override the GPU index
del gpu_index # We override the GPU index
print("Copy CPU index to %d sharded GPU indexes" % replicas)
@@ -719,8 +733,9 @@ def eval_dataset(index, preproc):
print(" probe=%-3d: %.3f s" % (nprobe, t1 - t0), end=' ')
gtc = gt_I[:, :1]
nq = xq.shape[0]
for rank in 1, 10, 100:
if rank > nnn: continue
for rank in (1, 10, 100):
if rank > nnn:
continue
nok = (I[:, :rank] == gtc).sum()
print("1-R@%d: %.4f" % (rank, nok / float(nq)), end=' ')
print()
benchs/distributed_ondisk/search_server.py
+7 -8
View File
@@ -10,12 +10,16 @@ from faiss.contrib import rpc
import combined_index
import argparse
from multiprocessing.pool import ThreadPool
import faiss
import numpy as np
############################################################
# Server implementation
############################################################
class MyServer(rpc.Server):
""" Assign version that can be exposed via RPC """
def __init__(self, s, index):
@@ -56,6 +60,7 @@ def main():
args.port, report_to_file=when_ready,
v6=not args.ipv4)
if __name__ == '__main__':
main()
@@ -64,11 +69,6 @@ if __name__ == '__main__':
# Client implementation
############################################################
from multiprocessing.pool import ThreadPool
import faiss
import numpy as np
class ResultHeap:
""" Combine query results from a sliced dataset (for k-nn search) """
@@ -111,7 +111,6 @@ def distribute_weights(weights, nbin):
return bins, assign
class SplitPerListIndex:
"""manages a local index, that does the coarse quantization and a set
of sub_indexes. The sub_indexes search a subset of the inverted
@@ -177,7 +176,7 @@ class SplitPerListIndex:
t0 = time.time()
Di, Ii = sub_index.ivf_search_preassigned(
xqo, list_nos_i, coarse_dis, k)
#print(list_nos_i, Ii)
# print(list_nos_i, Ii)
if self.verbose:
print('client %d: %.3f s' % (i, time.time() - t0))
return Di, Ii
@@ -185,7 +184,7 @@ class SplitPerListIndex:
rh = ResultHeap(x.shape[0], k)
for Di, Ii in self.pool.imap(do_query, range(self.ni)):
#print("ADD", Ii, rh.I)
# print("ADD", Ii, rh.I)
rh.add_batch_result(Di, Ii, 0)
rh.finalize()
return rh.D, rh.I
contrib/torch_utils.py
+98 -46
View File
@@ -32,6 +32,7 @@ import numpy as np
# Equivalent of swig_ptr for Torch tensors
##################################################################
def swig_ptr_from_UInt8Tensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
assert x.is_contiguous()
@@ -97,7 +98,8 @@ def using_stream(res, pytorch_stream=None):
pytorch_stream = torch.cuda.current_stream()
# This is the cudaStream_t that we wish to use
cuda_stream_s = faiss.cast_integer_to_cudastream_t(pytorch_stream.cuda_stream)
cuda_stream_s = faiss.cast_integer_to_cudastream_t(
pytorch_stream.cuda_stream)
# So we can revert GpuResources stream state upon exit
prior_dev = torch.cuda.current_device()
@@ -111,6 +113,7 @@ def using_stream(res, pytorch_stream=None):
finally:
res.setDefaultStream(prior_dev, prior_stream)
def torch_replace_method(the_class, name, replacement,
ignore_missing=False, ignore_no_base=False):
try:
@@ -132,8 +135,9 @@ def torch_replace_method(the_class, name, replacement,
# Setup wrappers
##################################################################
def handle_torch_Index(the_class):
def torch_replacement_add(self, x, numeric_type = faiss.Float32):
def torch_replacement_add(self, x, numeric_type=faiss.Float32):
if type(x) is np.ndarray:
# forward to faiss __init__.py base method
return self.add_numpy(x)
@@ -146,11 +150,13 @@ def handle_torch_Index(the_class):
elif numeric_type == faiss.Float16:
x_ptr = swig_ptr_from_HalfTensor(x)
else:
raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ")
raise ValueError(
"numeric type must be either faiss.Float32 or faiss.Float16"
)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -159,7 +165,8 @@ def handle_torch_Index(the_class):
# CPU torch
self.add_ex(n, x_ptr, numeric_type)
def torch_replacement_add_with_ids(self, x, ids, numeric_type = faiss.Float32):
def torch_replacement_add_with_ids(self, x, ids,
numeric_type=faiss.Float32):
if type(x) is np.ndarray:
# forward to faiss __init__.py base method
return self.add_with_ids_numpy(x, ids)
@@ -172,14 +179,17 @@ def handle_torch_Index(the_class):
elif numeric_type == faiss.Float16:
x_ptr = swig_ptr_from_HalfTensor(x)
else:
raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ")
raise ValueError(
"numeric type must be either faiss.Float32 or faiss.Float16"
)
assert type(ids) is torch.Tensor
assert ids.shape == (n, ), 'not same number of vectors as ids'
ids_ptr = swig_ptr_from_IndicesTensor(ids)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -206,7 +216,8 @@ def handle_torch_Index(the_class):
L_ptr = swig_ptr_from_IndicesTensor(labels)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -217,7 +228,7 @@ def handle_torch_Index(the_class):
return labels
def torch_replacement_train(self, x, numeric_type = faiss.Float32):
def torch_replacement_train(self, x, numeric_type=faiss.Float32):
if type(x) is np.ndarray:
# forward to faiss __init__.py base method
return self.train_numpy(x)
@@ -230,10 +241,13 @@ def handle_torch_Index(the_class):
elif numeric_type == faiss.Float16:
x_ptr = swig_ptr_from_HalfTensor(x)
else:
raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ")
raise ValueError(
"numeric type must be either faiss.Float32 or faiss.Float16"
)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -249,7 +263,9 @@ def handle_torch_Index(the_class):
elif numeric_type == faiss.Float16:
x_ptr = swig_ptr_from_HalfTensor(x)
else:
raise ValueError("numeric type must be either faiss.Float32 or faiss.Float16 ")
raise ValueError(
"numeric type must be either faiss.Float32 or faiss.Float16"
)
if D is None:
D = torch.empty(n, k, device=x.device, dtype=torch.float32)
@@ -267,7 +283,8 @@ def handle_torch_Index(the_class):
return x_ptr, D_ptr, I_ptr, D, I
def torch_replacement_search(self, x, k, D=None, I=None, numeric_type=faiss.Float32):
def torch_replacement_search(self, x, k, D=None, I=None,
numeric_type=faiss.Float32):
if type(x) is np.ndarray:
# forward to faiss __init__.py base method
return self.search_numpy(x, k, D=D, I=I)
@@ -279,7 +296,8 @@ def handle_torch_Index(the_class):
x_ptr, D_ptr, I_ptr, D, I = search_methods_common(x, k, D, I)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -290,7 +308,8 @@ def handle_torch_Index(the_class):
return D, I
def torch_replacement_search_and_reconstruct(self, x, k, D=None, I=None, R=None):
def torch_replacement_search_and_reconstruct(self, x, k, D=None, I=None,
R=None):
if type(x) is np.ndarray:
# Forward to faiss __init__.py base method
return self.search_and_reconstruct_numpy(x, k, D=D, I=I, R=R)
@@ -309,18 +328,21 @@ def handle_torch_Index(the_class):
R_ptr = swig_ptr_from_FloatTensor(R)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
self.search_and_reconstruct_c(n, x_ptr, k, D_ptr, I_ptr, R_ptr)
else:
# CPU torch
self.search_and_reconstruct_c(n, x_ptr, k, D_ptr, I_ptr, R_ptr)
self.search_and_reconstruct_c(n, x_ptr, k, D_ptr, I_ptr,
R_ptr)
return D, I, R
def torch_replacement_search_preassigned(self, x, k, Iq, Dq, *, D=None, I=None):
def torch_replacement_search_preassigned(self, x, k, Iq, Dq, *,
D=None, I=None):
if type(x) is np.ndarray:
# forward to faiss __init__.py base method
return self.search_preassigned_numpy(x, k, Iq, Dq, D=D, I=I)
@@ -343,20 +365,24 @@ def handle_torch_Index(the_class):
Dq_ptr = None
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
self.search_preassigned_c(n, x_ptr, k, Iq_ptr, Dq_ptr, D_ptr, I_ptr, False)
self.search_preassigned_c(n, x_ptr, k, Iq_ptr, Dq_ptr,
D_ptr, I_ptr, False)
else:
# CPU torch
self.search_preassigned_c(n, x_ptr, k, Iq_ptr, Dq_ptr, D_ptr, I_ptr, False)
self.search_preassigned_c(n, x_ptr, k, Iq_ptr, Dq_ptr,
D_ptr, I_ptr, False)
return D, I
def torch_replacement_remove_ids(self, x):
# Not yet implemented
assert type(x) is not torch.Tensor, 'remove_ids not yet implemented for torch'
assert type(x) is not torch.Tensor, \
'remove_ids not yet implemented for torch'
return self.remove_ids_numpy(x)
def torch_replacement_reconstruct(self, key, x=None):
@@ -382,7 +408,8 @@ def handle_torch_Index(the_class):
x_ptr = swig_ptr_from_FloatTensor(x)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -419,7 +446,8 @@ def handle_torch_Index(the_class):
x_ptr = swig_ptr_from_FloatTensor(x)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -444,7 +472,8 @@ def handle_torch_Index(the_class):
x_ptr = swig_ptr_from_FloatTensor(x)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -465,8 +494,10 @@ def handle_torch_Index(the_class):
assert d == self.d
x_ptr = swig_ptr_from_FloatTensor(x)
assert not x.is_cuda, 'Range search using GPU tensor not yet implemented'
assert not hasattr(self, 'getDevice'), 'Range search on GPU index not yet implemented'
assert not x.is_cuda, \
'Range search using GPU tensor not yet implemented'
assert not hasattr(self, 'getDevice'), \
'Range search on GPU index not yet implemented'
res = faiss.RangeSearchResult(n)
self.range_search_c(n, x_ptr, thresh, res)
@@ -475,7 +506,8 @@ def handle_torch_Index(the_class):
# FIXME: no rev_swig_ptr equivalent for torch.Tensor, just convert
# np to torch
# NOTE: torch does not support np.uint64, just np.int64
lims = torch.from_numpy(faiss.rev_swig_ptr(res.lims, n + 1).copy().astype('int64'))
lims = torch.from_numpy(
faiss.rev_swig_ptr(res.lims, n + 1).copy().astype('int64'))
nd = int(lims[-1])
D = torch.from_numpy(faiss.rev_swig_ptr(res.distances, nd).copy())
I = torch.from_numpy(faiss.rev_swig_ptr(res.labels, nd).copy())
@@ -499,7 +531,8 @@ def handle_torch_Index(the_class):
codes_ptr = swig_ptr_from_UInt8Tensor(codes)
if x.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -528,7 +561,8 @@ def handle_torch_Index(the_class):
x_ptr = swig_ptr_from_FloatTensor(x)
if codes.is_cuda:
assert hasattr(self, 'getDevice'), 'GPU tensor on CPU index not allowed'
assert hasattr(self, 'getDevice'), \
'GPU tensor on CPU index not allowed'
# On the GPU, use proper stream ordering
with using_stream(self.getResources()):
@@ -539,25 +573,33 @@ def handle_torch_Index(the_class):
return x
torch_replace_method(the_class, 'add', torch_replacement_add)
torch_replace_method(the_class, 'add_with_ids', torch_replacement_add_with_ids)
torch_replace_method(the_class, 'add_with_ids',
torch_replacement_add_with_ids)
torch_replace_method(the_class, 'assign', torch_replacement_assign)
torch_replace_method(the_class, 'train', torch_replacement_train)
torch_replace_method(the_class, 'search', torch_replacement_search)
torch_replace_method(the_class, 'remove_ids', torch_replacement_remove_ids)
torch_replace_method(the_class, 'reconstruct', torch_replacement_reconstruct)
torch_replace_method(the_class, 'reconstruct_n', torch_replacement_reconstruct_n)
torch_replace_method(the_class, 'range_search', torch_replacement_range_search)
torch_replace_method(the_class, 'update_vectors', torch_replacement_update_vectors,
torch_replace_method(the_class, 'remove_ids',
torch_replacement_remove_ids)
torch_replace_method(the_class, 'reconstruct',
torch_replacement_reconstruct)
torch_replace_method(the_class, 'reconstruct_n',
torch_replacement_reconstruct_n)
torch_replace_method(the_class, 'range_search',
torch_replacement_range_search)
torch_replace_method(the_class, 'update_vectors',
torch_replacement_update_vectors,
ignore_missing=True)
torch_replace_method(the_class, 'search_and_reconstruct',
torch_replacement_search_and_reconstruct, ignore_missing=True)
torch_replacement_search_and_reconstruct,
ignore_missing=True)
torch_replace_method(the_class, 'search_preassigned',
torch_replacement_search_preassigned, ignore_missing=True)
torch_replacement_search_preassigned,
ignore_missing=True)
torch_replace_method(the_class, 'sa_encode', torch_replacement_sa_encode)
torch_replace_method(the_class, 'sa_decode', torch_replacement_sa_decode)
faiss_module = sys.modules['faiss']
# Re-patch anything that inherits from faiss.Index to add the torch bindings
@@ -573,7 +615,8 @@ for symbol in dir(faiss_module):
def torch_replacement_knn(xq, xb, k, metric=faiss.METRIC_L2, metric_arg=0):
if type(xb) is np.ndarray:
# Forward to faiss __init__.py base method
return faiss.knn_numpy(xq, xb, k, metric=metric, metric_arg=metric_arg)
return faiss.knn_numpy(xq, xb, k, metric=metric,
metric_arg=metric_arg)
nb, d = xb.size()
assert xb.is_contiguous()
@@ -616,7 +659,9 @@ torch_replace_method(faiss_module, 'knn', torch_replacement_knn, True, True)
# allows torch tensor usage with bfKnn
def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRIC_L2, device=-1, use_cuvs=False):
def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
metric=faiss.METRIC_L2, device=-1,
use_cuvs=False):
if type(xb) is np.ndarray:
# Forward to faiss __init__.py base method
return faiss.knn_gpu_numpy(res, xq, xb, k, D, I, metric, device)
@@ -710,10 +755,15 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None, metric=faiss.METRI
return D, I
torch_replace_method(faiss_module, 'knn_gpu', torch_replacement_knn_gpu, True, True)
torch_replace_method(faiss_module, 'knn_gpu', torch_replacement_knn_gpu,
True, True)
# allows torch tensor usage with bfKnn for all pairwise distances
def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None, metric=faiss.METRIC_L2, device=-1):
def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
metric=faiss.METRIC_L2,
device=-1):
if type(xb) is np.ndarray:
# Forward to faiss __init__.py base method
return faiss.pairwise_distance_gpu_numpy(res, xq, xb, D, metric)
@@ -766,7 +816,7 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None, metric=faiss.ME
args = faiss.GpuDistanceParams()
args.metric = metric
args.k = -1 # selects all pairwise distance
args.k = -1 # selects all pairwise distance
args.dims = d
args.vectors = xb_ptr
args.vectorsRowMajor = xb_row_major
@@ -784,4 +834,6 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None, metric=faiss.ME
return D
torch_replace_method(faiss_module, 'pairwise_distance_gpu', torch_replacement_pairwise_distance_gpu, True, True)
torch_replace_method(faiss_module, 'pairwise_distance_gpu',
torch_replacement_pairwise_distance_gpu, True, True)
faiss/gpu/test/test_binary_cagra.py
+1
View File
@@ -26,6 +26,7 @@ import faiss
from faiss.contrib import evaluation
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
faiss/gpu/test/test_cagra.py
+29 -22
View File
@@ -14,8 +14,6 @@ import numpy as np
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
class TestComputeGT(unittest.TestCase):
def do_compute_GT(self, metric, numeric_type):
@@ -23,20 +21,24 @@ class TestComputeGT(unittest.TestCase):
k = 12
if numeric_type == faiss.Int8:
data_base_nt = np.random.randint(-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(-128, 128, size=(100, d), dtype=np.int8)
data_base_nt = np.random.randint(
-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(
-128, 128, size=(100, d), dtype=np.int8)
data_base = data_base_nt.astype(np.float32)
data_query = data_query_nt.astype(np.float32)
else:
ds = datasets.SyntheticDataset(d, 0, 10000, 100)
data_base = ds.get_database() #fp32
data_query = ds.get_queries() #fp32
data_base = ds.get_database() # fp32
data_query = ds.get_queries() # fp32
# Normalize for inner product to avoid duplicate neighbors
if metric == faiss.METRIC_INNER_PRODUCT:
# Normalize database vectors
data_base = data_base / np.linalg.norm(data_base, axis=1, keepdims=True)
data_base = data_base / np.linalg.norm(
data_base, axis=1, keepdims=True)
# Normalize query vectors
data_query = data_query / np.linalg.norm(data_query, axis=1, keepdims=True)
data_query = data_query / np.linalg.norm(
data_query, axis=1, keepdims=True)
if numeric_type == faiss.Float16:
data_base_nt = data_base.astype(np.float16)
data_query_nt = data_query.astype(np.float16)
@@ -81,25 +83,26 @@ class TestComputeGT(unittest.TestCase):
def test_compute_GT_IP_Int8(self):
self.do_compute_GT(faiss.METRIC_INNER_PRODUCT, faiss.Int8)
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
class TestInterop(unittest.TestCase):
def do_interop(self, metric, numeric_type):
d = 64
k = 12
if numeric_type == faiss.Int8:
data_base_nt = np.random.randint(-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(-128, 128, size=(100, d), dtype=np.int8)
data_base_nt = np.random.randint(
-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(
-128, 128, size=(100, d), dtype=np.int8)
data_base = data_base_nt.astype(np.float32)
data_query = data_query_nt.astype(np.float32)
else:
ds = datasets.SyntheticDataset(d, 0, 10000, 100)
data_base = ds.get_database() #fp32
data_query = ds.get_queries() #fp32
data_base = ds.get_database() # fp32
data_query = ds.get_queries() # fp32
if numeric_type == faiss.Float16:
data_base_nt = data_base.astype(np.float16)
data_query_nt = data_query.astype(np.float16)
@@ -122,7 +125,9 @@ class TestInterop(unittest.TestCase):
deserialized_index = faiss.deserialize_index(
faiss.serialize_index(cpu_index))
gpu_index = faiss.index_cpu_to_gpu(res, 0, deserialized_index)
gpu_index = faiss.index_cpu_to_gpu(
res, 0, deserialized_index
)
Dnew2, Inew2 = gpu_index.search(data_query_nt, k, numeric_type=numeric_type)
evaluation.check_ref_knn_with_draws(Dnew2, Inew2, Dnew, Inew, k)
@@ -149,22 +154,22 @@ class TestInterop(unittest.TestCase):
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
class TestIDMapCagra(unittest.TestCase):
def do_IDMapCagra(self, metric, numeric_type):
d = 64
k = 12
if numeric_type == faiss.Int8:
data_base_nt = np.random.randint(-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(-128, 128, size=(100, d), dtype=np.int8)
data_base_nt = np.random.randint(
-128, 128, size=(10000, d), dtype=np.int8)
data_query_nt = np.random.randint(
-128, 128, size=(100, d), dtype=np.int8)
data_base = data_base_nt.astype(np.float32)
data_query = data_query_nt.astype(np.float32)
else:
ds = datasets.SyntheticDataset(d, 0, 10000, 100)
data_base = ds.get_database() #fp32
data_query = ds.get_queries() #fp32
data_base = ds.get_database() # fp32
data_query = ds.get_queries() # fp32
if numeric_type == faiss.Float16:
data_base_nt = data_base.astype(np.float16)
data_query_nt = data_query.astype(np.float16)
@@ -180,7 +185,9 @@ class TestIDMapCagra(unittest.TestCase):
idMapIndex = faiss.IndexIDMap(index)
idMapIndex.train(data_base_nt, numeric_type=numeric_type)
ids = np.array([i for i in range(10000)])
idMapIndex.add_with_ids(data_base_nt, ids, numeric_type=numeric_type)
idMapIndex.add_with_ids(
data_base_nt, ids, numeric_type=numeric_type
)
Dnew, Inew = idMapIndex.search(data_query_nt, k, numeric_type=numeric_type)
evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k)
faiss/gpu/test/test_gpu_index_ivfsq.py
+23 -15
View File
@@ -8,6 +8,7 @@ import unittest
import numpy as np
import faiss
def make_t(num, d, clamp=False):
rs = np.random.RandomState(123)
x = rs.rand(num, d).astype('float32')
@@ -15,12 +16,13 @@ def make_t(num, d, clamp=False):
x = (x * 255).astype('uint8').astype('float32')
return x
def make_indices_copy_from_cpu(nlist, d, qtype, by_residual, metric, clamp):
to_train = make_t(10000, d, clamp)
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
qtype, metric, by_residual)
idx_cpu = faiss.IndexIVFScalarQuantizer(
quantizer_cp, d, nlist, qtype, metric, by_residual)
idx_cpu.train(to_train)
idx_cpu.add(to_train)
@@ -41,14 +43,14 @@ def make_indices_copy_from_gpu(nlist, d, qtype, by_residual, metric, clamp):
res.noTempMemory()
config = faiss.GpuIndexIVFScalarQuantizerConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist,
qtype, metric, by_residual, config)
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, metric, by_residual, config)
idx_gpu.train(to_train)
idx_gpu.add(to_train)
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
qtype, metric, by_residual)
idx_cpu = faiss.IndexIVFScalarQuantizer(
quantizer_cp, d, nlist, qtype, metric, by_residual)
idx_gpu.copyTo(idx_cpu)
return idx_cpu, idx_gpu
@@ -58,9 +60,9 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
to_train = make_t(10000, d, clamp)
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(quantizer_cp, d, nlist,
qtype, metric, by_residual)
assert(by_residual == idx_cpu.by_residual)
idx_cpu = faiss.IndexIVFScalarQuantizer(
quantizer_cp, d, nlist, qtype, metric, by_residual)
assert by_residual == idx_cpu.by_residual
idx_cpu.train(to_train)
idx_cpu.add(to_train)
@@ -69,9 +71,9 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
res.noTempMemory()
config = faiss.GpuIndexIVFScalarQuantizerConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist,
qtype, metric, by_residual, config)
assert(by_residual == idx_gpu.by_residual)
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, metric, by_residual, config)
assert by_residual == idx_gpu.by_residual
idx_gpu.train(to_train)
idx_gpu.add(to_train)
@@ -82,6 +84,7 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
# Testing functions
#
def summarize_results(dist, idx):
valid = []
invalid = []
@@ -100,6 +103,7 @@ def summarize_results(dist, idx):
return valid, invalid
def compare_results(d1, i1, d2, i2):
# Count number of index differences
idx_diffs = {}
@@ -115,10 +119,10 @@ def compare_results(d1, i1, d2, i2):
if (len(inv1) != len(inv2)):
print('mismatch ', len(inv1), len(inv2), inv2[0])
assert(len(inv1) == len(inv2))
assert len(inv1) == len(inv2)
idx_invalid += len(inv2)
for x1, x2 in zip(inv1, inv2):
assert(x1 == x2)
assert x1 == x2
for _, (query1, query2) in enumerate(zip(valid1, valid2)):
for idx1, order_d1 in query1.items():
@@ -135,6 +139,7 @@ def compare_results(d1, i1, d2, i2):
return idx_diffs, idx_diffs_inf, idx_invalid
def check_diffs(total_num, in_window_thresh, diffs, diff_inf, invalid):
# We require a certain fraction of results to be within +/- diff_window
# index differences
@@ -147,7 +152,8 @@ def check_diffs(total_num, in_window_thresh, diffs, diff_inf, invalid):
if (in_window < in_window_thresh):
print('error {} {}'.format(in_window, in_window_thresh))
assert(in_window >= in_window_thresh)
assert in_window >= in_window_thresh
def do_test_with_index(ci, gi, nprobe, k, clamp, in_window_thresh):
num_query = 11
@@ -161,6 +167,7 @@ def do_test_with_index(ci, gi, nprobe, k, clamp, in_window_thresh):
*compare_results(*ci.search(to_query, k),
*gi.search(to_query, k)))
def do_test(nlist, d, qtype, by_residual, metric, nprobe, k):
clamp = (qtype == faiss.ScalarQuantizer.QT_8bit_direct)
ci, gi = make_indices_copy_from_cpu(nlist, d, qtype,
@@ -181,6 +188,7 @@ def do_test(nlist, d, qtype, by_residual, metric, nprobe, k):
# and residuals
do_test_with_index(ci, gi, nprobe, k, clamp, 0.8)
def do_multi_test(qtype):
nlist = 100
nprobe = 10
tests/test_io.py
+7 -6
View File
@@ -21,6 +21,7 @@ nt = 2000
nb = 1000
nq = 200
class TestIOVariants(unittest.TestCase):
def test_io_error(self):
@@ -192,7 +193,6 @@ class TestCallbacks(unittest.TestCase):
if os.path.exists(fname):
os.unlink(fname)
def test_transfer_pipe(self):
""" transfer an index through a Unix pipe """
@@ -469,7 +469,8 @@ class TestIVFPQRead(unittest.TestCase):
faiss.write_index(index, fname)
index_a = faiss.read_index(fname)
index_b = faiss.read_index(fname, faiss.IO_FLAG_SKIP_PRECOMPUTE_TABLE)
index_b = faiss.read_index(
fname, faiss.IO_FLAG_SKIP_PRECOMPUTE_TABLE)
Da, Ia = index_a.search(xq, 10)
Db, Ib = index_b.search(xq, 10)
@@ -485,7 +486,6 @@ class TestIVFPQRead(unittest.TestCase):
os.unlink(fname)
class TestIOFlatMMap(unittest.TestCase):
@unittest.skipIf(
platform.system() not in ["Windows", "Linux"],
@@ -520,15 +520,16 @@ class TestIOFlatMMap(unittest.TestCase):
except:
pass
def test_zerocopy(self):
def test_zerocopy(self):
xt, xb, xq = get_dataset_2(32, 0, 100, 50)
index = faiss.index_factory(32, "SQfp16", faiss.METRIC_L2)
# does not need training
# does not need training
index.add(xb)
Dref, Iref = index.search(xq, 10)
serialized_index = faiss.serialize_index(index)
reader = faiss.ZeroCopyIOReader(faiss.swig_ptr(serialized_index), serialized_index.size)
reader = faiss.ZeroCopyIOReader(
faiss.swig_ptr(serialized_index), serialized_index.size)
index2 = faiss.read_index(reader)
Dnew, Inew = index2.search(xq, 10)
np.testing.assert_array_equal(Iref, Inew)