facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-11 22:50:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Format faiss Python files with black (line length 80) (#5349)

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

# this is a no-op change

Ran `black -l 80` over all tracked Python files under `fbcode/faiss/`, using line length 80 to match faiss's documented coding style (`CONTRIBUTING.md`: "80 character line length (both for C++ and Python)"). 172 files were reformatted; 28 were already compliant.

This is a pure-formatting change. `black` runs its AST-equivalence safety check on every file, so no identifiers, string/numeric values, operators, or control flow were altered — only whitespace, indentation, line wrapping, blank lines, trailing commas, and string-quote normalization (single to double quotes).

`fbcode/faiss/` is intentionally excluded from Meta's standard Python autoformatter (`arc f` / `arc lint -a` report no formatter for these paths) because it mirrors to public GitHub, so the Python had drifted out of any consistent style. The C++ is already clang-format-clean (0 of 807 files change), so this effort is Python-only. This diff brings the Python to canonical `black` formatting.

Excluded `demos/index_pq_flat_separate_codes_from_codebook.py`: it contains a Jupyter `!rm` shell-magic line and is not valid Python, so `black` cannot parse it.

This is the base of a 3-diff stack: (1) this black reformat; (2) `.flake8` alignment (max-line-length 80 + black-compatible ignores E203/W503); (3) genuine lint fixes (F401/F403/F841/B011).

Reviewed By: mnorris11

Differential Revision: D109736099

fbshipit-source-id: 9a054635766ed11aa3760a788395babd3d198abf
This commit is contained in:
Tia Lim
2026-06-29 17:29:05 -07:00
committed by meta-codesync[bot]
parent 0fa5b0afcb
commit 0f0f728166
159 changed files with 6369 additions and 4696 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchs/bench_all_ivf/bench_all_ivf.py
+240 -127
View File
@@ -20,7 +20,6 @@ except ModuleNotFoundError:
sanitize = datasets.sanitize
def unwind_index_ivf(index):
if isinstance(index, faiss.IndexPreTransform):
assert index.chain.size() == 1
@@ -63,9 +62,8 @@ def apply_AQ_options(index, args):
index.rq.use_beam_LUT = args.RQ_use_beam_LUT
def eval_setting(index, xq, gt, k, inter, min_time):
""" evaluate searching in terms of precision vs. speed """
"""evaluate searching in terms of precision vs. speed"""
nq = xq.shape[0]
ivf_stats = faiss.cvar.indexIVF_stats
ivf_stats.reset()
@@ -77,29 +75,29 @@ def eval_setting(index, xq, gt, k, inter, min_time):
t1 = time.time()
if t1 - t0 > min_time:
break
ms_per_query = ((t1 - t0) * 1000.0 / nq / nrun)
res = {
"ms_per_query": ms_per_query,
"nrun": nrun
}
ms_per_query = (t1 - t0) * 1000.0 / nq / nrun
res = {"ms_per_query": ms_per_query, "nrun": nrun}
res["n"] = ms_per_query
if inter:
rank = k
inter_measure = faiss.eval_intersection(gt[:, :rank], I[:, :rank]) / (nq * rank)
print("%.4f" % inter_measure, end=' ')
inter_measure = faiss.eval_intersection(gt[:, :rank], I[:, :rank]) / (
nq * rank
)
print("%.4f" % inter_measure, end=" ")
res["inter_measure"] = inter_measure
else:
res["recalls"] = {}
for rank in 1, 10, 100:
recall = (I[:, :rank] == gt[:, :1]).sum() / float(nq)
print("%.4f" % recall, end=' ')
print("%.4f" % recall, end=" ")
res["recalls"][rank] = recall
print(" %9.5f " % ms_per_query, end=' ')
print("%12d " % (ivf_stats.ndis / nrun), end=' ')
print(" %9.5f " % ms_per_query, end=" ")
print("%12d " % (ivf_stats.ndis / nrun), end=" ")
print(nrun)
res["ndis"] = ivf_stats.ndis / nrun
return res
######################################################
# Training
######################################################
@@ -109,8 +107,9 @@ def run_train(args, ds, res):
print("build index, key=", args.indexkey)
index = faiss.index_factory(
ds.d, args.indexkey, faiss.METRIC_L2 if ds.metric == "L2" else
faiss.METRIC_INNER_PRODUCT
ds.d,
args.indexkey,
faiss.METRIC_L2 if ds.metric == "L2" else faiss.METRIC_INNER_PRODUCT,
)
index_ivf, vec_transform = unwind_index_ivf(index)
@@ -118,7 +117,7 @@ def run_train(args, ds, res):
if args.by_residual != -1:
by_residual = args.by_residual == 1
print("setting by_residual = ", by_residual)
index_ivf.by_residual # check if field exists
index_ivf.by_residual # check if field exists
index_ivf.by_residual = by_residual
if index_ivf:
@@ -127,16 +126,21 @@ def run_train(args, ds, res):
# because otherwise the assignment is inaccurate
quantizer = faiss.downcast_index(index_ivf.quantizer)
if isinstance(quantizer, faiss.IndexRefine):
print(" update quantizer k_factor=", quantizer.k_factor, end=" -> ")
print(
" update quantizer k_factor=", quantizer.k_factor, end=" -> "
)
quantizer.k_factor = 32 if index_ivf.nlist < 1e6 else 64
print(quantizer.k_factor)
base_index = faiss.downcast_index(quantizer.base_index)
if isinstance(base_index, faiss.IndexIVF):
print(" update quantizer nprobe=", base_index.nprobe, end=" -> ")
print(
" update quantizer nprobe=", base_index.nprobe, end=" -> "
)
base_index.nprobe = (
16 if base_index.nlist < 1e5 else
32 if base_index.nlist < 4e6 else
64)
16
if base_index.nlist < 1e5
else 32 if base_index.nlist < 4e6 else 64
)
print(base_index.nprobe)
elif isinstance(quantizer, faiss.IndexHNSW):
hnsw = quantizer.hnsw
@@ -159,7 +163,7 @@ def run_train(args, ds, res):
maxtrain = args.maxtrain
if maxtrain == 0:
if 'IMI' in args.indexkey:
if "IMI" in args.indexkey:
maxtrain = int(256 * 2 ** (np.log2(index_ivf.nlist) / 2))
elif index_ivf:
maxtrain = 50 * index_ivf.nlist
@@ -178,28 +182,31 @@ def run_train(args, ds, res):
print("train, size", xt2.shape)
assert np.all(np.isfinite(xt2))
if (isinstance(vec_transform, faiss.OPQMatrix) and
isinstance(index_ivf, faiss.IndexIVFPQFastScan)):
if isinstance(vec_transform, faiss.OPQMatrix) and isinstance(
index_ivf, faiss.IndexIVFPQFastScan
):
print(" Forcing OPQ training PQ to PQ4")
ref_pq = index_ivf.pq
training_pq = faiss.ProductQuantizer(
ref_pq.d, ref_pq.M, ref_pq.nbits
)
training_pq = faiss.ProductQuantizer(ref_pq.d, ref_pq.M, ref_pq.nbits)
vec_transform.pq
vec_transform.pq = training_pq
if args.get_centroids_from == '':
if args.get_centroids_from == "":
if args.clustering_niter >= 0:
print(("setting nb of clustering iterations to %d" %
args.clustering_niter))
print(
(
"setting nb of clustering iterations to %d"
% args.clustering_niter
)
)
index_ivf.cp.niter = args.clustering_niter
if args.train_on_gpu:
print("add a training index on GPU")
train_index = faiss.index_cpu_to_all_gpus(
faiss.IndexFlatL2(index_ivf.d))
faiss.IndexFlatL2(index_ivf.d)
)
index_ivf.clustering_index = train_index
else:
@@ -229,6 +236,7 @@ def run_train(args, ds, res):
print(" train in %.3f s" % res.train_time)
return index
######################################################
# Populating index
######################################################
@@ -242,14 +250,15 @@ def run_add(args, ds, index, res):
assert args.split == [1, 0], "split not supported with full batch add"
index.add(sanitize(ds.get_database()))
else:
totn = ds.nb // args.split[0] # approximate
totn = ds.nb // args.split[0] # approximate
i0 = 0
print(f"Adding in block sizes {args.add_bs} with split {args.split}")
for xblock in ds.database_iterator(bs=args.add_bs, split=args.split):
i1 = i0 + len(xblock)
print(" adding %d:%d / %d [%.3f s, RSS %d kiB] " % (
i0, i1, totn, time.time() - t0,
faiss.get_mem_usage_kb()))
print(
" adding %d:%d / %d [%.3f s, RSS %d kiB] "
% (i0, i1, totn, time.time() - t0, faiss.get_mem_usage_kb())
)
index.add(xblock)
i0 = i1
@@ -261,6 +270,7 @@ def run_add(args, ds, index, res):
# Search
######################################################
def run_search(args, ds, index, res):
index_ivf, vec_transform = unwind_index_ivf(index)
@@ -283,7 +293,7 @@ def run_search(args, ds, index, res):
print("current RSS:", faiss.get_mem_usage_kb() * 1024)
precomputed_table_size = 0
if hasattr(index_ivf, 'precomputed_table'):
if hasattr(index_ivf, "precomputed_table"):
precomputed_table_size = index_ivf.precomputed_table.size() * 4
print("precomputed tables size:", precomputed_table_size)
@@ -294,7 +304,7 @@ def run_search(args, ds, index, res):
nq, d = xq.shape
gt = ds.get_groundtruth(k=args.k)
if not args.accept_short_gt: # Deep1B has only a single NN per query
if not args.accept_short_gt: # Deep1B has only a single NN per query
assert gt.shape[1] == args.k
if args.searchthreads != -1:
@@ -309,26 +319,25 @@ def run_search(args, ds, index, res):
parametersets = args.searchparams
if args.inter:
header = (
'%-40s inter@%3d time(ms/q) nb distances #runs' %
("parameters", args.k)
header = "%-40s inter@%3d time(ms/q) nb distances #runs" % (
"parameters",
args.k,
)
else:
header = (
'%-40s R@1 R@10 R@100 time(ms/q) nb distances #runs' %
"parameters"
"%-40s R@1 R@10 R@100 time(ms/q) nb distances #runs"
% "parameters"
)
res.search_results = {}
if parametersets == ['autotune']:
if parametersets == ["autotune"]:
ps.n_experiments = args.n_autotune
ps.min_test_duration = args.min_test_duration
for kv in args.autotune_max:
k, vmax = kv.split(':')
k, vmax = kv.split(":")
vmax = float(vmax)
print("limiting %s to %g" % (k, vmax))
pr = ps.add_range(k)
@@ -337,8 +346,8 @@ def run_search(args, ds, index, res):
faiss.copy_array_to_vector(values, pr.values)
for kv in args.autotune_range:
k, vals = kv.split(':')
vals = np.fromstring(vals, sep=',')
k, vals = kv.split(":")
vals = np.fromstring(vals, sep=",")
print("setting %s to %s" % (k, vals))
pr = ps.add_range(k)
faiss.copy_array_to_vector(vals, pr.values)
@@ -353,10 +362,13 @@ def run_search(args, ds, index, res):
# by default, the criterion will request only 1 NN
crit.nnn = args.k
crit.set_groundtruth(None, gt.astype('int64'))
crit.set_groundtruth(None, gt.astype("int64"))
# then we let Faiss find the optimal parameters by itself
print("exploring operating points, %d threads" % faiss.omp_get_max_threads());
print(
"exploring operating points, %d threads"
% faiss.omp_get_max_threads()
)
ps.display()
t0 = time.time()
@@ -375,24 +387,27 @@ def run_search(args, ds, index, res):
ps.set_index_parameters(index, opt.key)
print(opt.key.ljust(maxw), end=' ')
print(opt.key.ljust(maxw), end=" ")
sys.stdout.flush()
res_i = eval_setting(index, xq, gt, args.k, args.inter, args.min_test_duration)
res_i = eval_setting(
index, xq, gt, args.k, args.inter, args.min_test_duration
)
res.search_results[opt.key] = res_i
else:
print(header)
for param in parametersets:
print("%-40s " % param, end=' ')
print("%-40s " % param, end=" ")
sys.stdout.flush()
ps.set_index_parameters(index, param)
res_i = eval_setting(index, xq, gt, args.k, args.inter, args.min_test_duration)
res_i = eval_setting(
index, xq, gt, args.k, args.inter, args.min_test_duration
)
res.search_results[param] = res_i
######################################################
# Driver function
######################################################
@@ -405,76 +420,173 @@ def main():
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('general options')
aa('--nthreads', default=-1, type=int,
help='nb of threads to use at train and add time')
aa('--json', default=False, action="store_true",
help="output stats in JSON format at the end")
aa('--todo', default=["check_files"],
choices=["train", "add", "search", "check_files"],
nargs="+", help='what to do (check_files means decide depending on which index files exist)')
group = parser.add_argument_group("general options")
aa(
"--nthreads",
default=-1,
type=int,
help="nb of threads to use at train and add time",
)
aa(
"--json",
default=False,
action="store_true",
help="output stats in JSON format at the end",
)
aa(
"--todo",
default=["check_files"],
choices=["train", "add", "search", "check_files"],
nargs="+",
help="what to do (check_files means decide depending on which index files exist)",
)
group = parser.add_argument_group('dataset options')
aa('--db', default='deep1M', help='dataset')
aa('--compute_gt', default=False, action='store_true',
help='compute and store the groundtruth')
aa('--force_IP', default=False, action="store_true",
help='force IP search instead of L2')
aa('--accept_short_gt', default=False, action='store_true',
help='work around a problem with Deep1B GT')
group = parser.add_argument_group("dataset options")
aa("--db", default="deep1M", help="dataset")
aa(
"--compute_gt",
default=False,
action="store_true",
help="compute and store the groundtruth",
)
aa(
"--force_IP",
default=False,
action="store_true",
help="force IP search instead of L2",
)
aa(
"--accept_short_gt",
default=False,
action="store_true",
help="work around a problem with Deep1B GT",
)
group = parser.add_argument_group('index construction')
aa('--indexkey', default='HNSW32', help='index_factory type')
aa('--trained_indexfile', default='',
help='file to read or write a trained index from')
aa('--maxtrain', default=256 * 256, type=int,
help='maximum number of training points (0 to set automatically)')
aa('--indexfile', default='', help='file to read or write index from')
aa('--split', default=[1, 0], type=int, nargs=2, help="database split")
aa('--add_bs', default=-1, type=int,
help='add elements index by batches of this size')
group = parser.add_argument_group("index construction")
aa("--indexkey", default="HNSW32", help="index_factory type")
aa(
"--trained_indexfile",
default="",
help="file to read or write a trained index from",
)
aa(
"--maxtrain",
default=256 * 256,
type=int,
help="maximum number of training points (0 to set automatically)",
)
aa("--indexfile", default="", help="file to read or write index from")
aa("--split", default=[1, 0], type=int, nargs=2, help="database split")
aa(
"--add_bs",
default=-1,
type=int,
help="add elements index by batches of this size",
)
group = parser.add_argument_group('IVF options')
aa('--by_residual', default=-1, type=int,
help="set if index should use residuals (default=unchanged)")
aa('--no_precomputed_tables', action='store_true', default=False,
help='disable precomputed tables (uses less memory)')
aa('--get_centroids_from', default='',
help='get the centroids from this index (to speed up training)')
aa('--clustering_niter', default=-1, type=int,
help='number of clustering iterations (-1 = leave default)')
aa('--train_on_gpu', default=False, action='store_true',
help='do training on GPU')
group = parser.add_argument_group("IVF options")
aa(
"--by_residual",
default=-1,
type=int,
help="set if index should use residuals (default=unchanged)",
)
aa(
"--no_precomputed_tables",
action="store_true",
default=False,
help="disable precomputed tables (uses less memory)",
)
aa(
"--get_centroids_from",
default="",
help="get the centroids from this index (to speed up training)",
)
aa(
"--clustering_niter",
default=-1,
type=int,
help="number of clustering iterations (-1 = leave default)",
)
aa(
"--train_on_gpu",
default=False,
action="store_true",
help="do training on GPU",
)
group = parser.add_argument_group('index-specific options')
aa('--M0', default=-1, type=int, help='size of base level for HNSW')
aa('--RQ_train_default', default=False, action="store_true",
help='disable progressive dim training for RQ')
aa('--RQ_beam_size', default=-1, type=int,
help='set beam size at add time')
aa('--LSQ_encode_ils_iters', default=-1, type=int,
help='ILS iterations for LSQ')
aa('--RQ_use_beam_LUT', default=-1, type=int,
help='use beam LUT at add time')
group = parser.add_argument_group("index-specific options")
aa("--M0", default=-1, type=int, help="size of base level for HNSW")
aa(
"--RQ_train_default",
default=False,
action="store_true",
help="disable progressive dim training for RQ",
)
aa("--RQ_beam_size", default=-1, type=int, help="set beam size at add time")
aa(
"--LSQ_encode_ils_iters",
default=-1,
type=int,
help="ILS iterations for LSQ",
)
aa(
"--RQ_use_beam_LUT",
default=-1,
type=int,
help="use beam LUT at add time",
)
group = parser.add_argument_group('searching')
aa('--k', default=100, type=int, help='nb of nearest neighbors')
aa('--inter', default=False, action='store_true',
help='use intersection measure instead of 1-recall as metric')
aa('--searchthreads', default=-1, type=int,
help='nb of threads to use at search time')
aa('--searchparams', nargs='+', default=['autotune'],
help="search parameters to use (can be autotune or a list of params)")
aa('--n_autotune', default=500, type=int,
help="max nb of autotune experiments")
aa('--autotune_max', default=[], nargs='*',
help='set max value for autotune variables format "var:val" (exclusive)')
aa('--autotune_range', default=[], nargs='*',
help='set complete autotune range, format "var:val1,val2,..."')
aa('--min_test_duration', default=3.0, type=float,
help='run test at least for so long to avoid jitter')
aa('--indexes_to_merge', default=[], nargs="*",
help="load these indexes to search and merge them before searching")
group = parser.add_argument_group("searching")
aa("--k", default=100, type=int, help="nb of nearest neighbors")
aa(
"--inter",
default=False,
action="store_true",
help="use intersection measure instead of 1-recall as metric",
)
aa(
"--searchthreads",
default=-1,
type=int,
help="nb of threads to use at search time",
)
aa(
"--searchparams",
nargs="+",
default=["autotune"],
help="search parameters to use (can be autotune or a list of params)",
)
aa(
"--n_autotune",
default=500,
type=int,
help="max nb of autotune experiments",
)
aa(
"--autotune_max",
default=[],
nargs="*",
help='set max value for autotune variables format "var:val" (exclusive)',
)
aa(
"--autotune_range",
default=[],
nargs="*",
help='set complete autotune range, format "var:val1,val2,..."',
)
aa(
"--min_test_duration",
default=3.0,
type=float,
help="run test at least for so long to avoid jitter",
)
aa(
"--indexes_to_merge",
default=[],
nargs="*",
help="load these indexes to search and merge them before searching",
)
args = parser.parse_args()
@@ -489,22 +601,23 @@ def main():
print("args:", args)
os.system('echo -n "nb processors "; '
'cat /proc/cpuinfo | grep ^processor | wc -l; '
'cat /proc/cpuinfo | grep ^"model name" | tail -1')
os.system(
'echo -n "nb processors "; '
"cat /proc/cpuinfo | grep ^processor | wc -l; "
'cat /proc/cpuinfo | grep ^"model name" | tail -1'
)
# object to collect results
res = argparse.Namespace()
res.args = args.__dict__
res.cpu_model = [
l for l in open("/proc/cpuinfo", "r")
if "model name" in l][0]
l for l in open("/proc/cpuinfo", "r") if "model name" in l
][0]
print("Load dataset")
ds = datasets.load_dataset(
dataset=args.db, compute_gt=args.compute_gt)
ds = datasets.load_dataset(dataset=args.db, compute_gt=args.compute_gt)
if args.force_IP:
ds.metric = "IP"
benchs/bench_all_ivf/bench_kmeans.py
+20 -18
View File
@@ -21,28 +21,30 @@ def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('dataset options')
group = parser.add_argument_group("dataset options")
aa('--db', default='deep1M', help='dataset')
aa('--nt', default=65536, type=int)
aa('--nb', default=100000, type=int)
aa('--nt_sample', default=0, type=int)
aa("--db", default="deep1M", help="dataset")
aa("--nt", default=65536, type=int)
aa("--nb", default=100000, type=int)
aa("--nt_sample", default=0, type=int)
group = parser.add_argument_group('kmeans options')
aa('--k', default=256, type=int)
aa('--seed', default=12345, type=int)
aa('--pcadim', default=-1, type=int, help='PCA to this dimension')
aa('--niter', default=25, type=int)
aa('--eval_freq', default=100, type=int)
group = parser.add_argument_group("kmeans options")
aa("--k", default=256, type=int)
aa("--seed", default=12345, type=int)
aa("--pcadim", default=-1, type=int, help="PCA to this dimension")
aa("--niter", default=25, type=int)
aa("--eval_freq", default=100, type=int)
args = parser.parse_args()
print("args:", args)
os.system('echo -n "nb processors "; '
'cat /proc/cpuinfo | grep ^processor | wc -l; '
'cat /proc/cpuinfo | grep ^"model name" | tail -1')
os.system(
'echo -n "nb processors "; '
"cat /proc/cpuinfo | grep ^processor | wc -l; "
'cat /proc/cpuinfo | grep ^"model name" | tail -1'
)
ngpu = faiss.get_num_gpus()
print("nb GPUs:", ngpu)
@@ -55,15 +57,15 @@ xt, xb, xq, gt = datasets.load_data(dataset=args.db)
if args.nt_sample == 0:
xt_pca = xt[args.nt:args.nt + 10000]
xt = xt[:args.nt]
xt_pca = xt[args.nt : args.nt + 10000]
xt = xt[: args.nt]
else:
xt_pca = xt[args.nt_sample:args.nt_sample + 10000]
xt_pca = xt[args.nt_sample : args.nt_sample + 10000]
rs = np.random.RandomState(args.seed)
idx = rs.choice(args.nt_sample, size=args.nt, replace=False)
xt = xt[idx]
xb = xb[:args.nb]
xb = xb[: args.nb]
d = xb.shape[1]
benchs/bench_all_ivf/cmp_with_scann.py
+87 -51
View File
@@ -47,31 +47,38 @@ def main():
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('dataset options')
group = parser.add_argument_group("dataset options")
aa('--db', default='deep1M', help='dataset')
aa('--measure', default="1-recall",
help="perf measure to use: 1-recall or inter")
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("--db", default="deep1M", help="dataset")
aa(
"--measure",
default="1-recall",
help="perf measure to use: 1-recall or inter",
)
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",
)
group = parser.add_argument_group('searching')
aa('--k', default=10, type=int, help='nb of nearest neighbors')
aa('--pre_reorder_k', default="0,10,100,1000", help='values for reorder_k')
aa('--nprobe', default="1,2,5,10,20,50,100,200", help='values for nprobe')
aa('--nrun', default=5, type=int, help='nb of runs to perform')
group = parser.add_argument_group("searching")
aa("--k", default=10, type=int, help="nb of nearest neighbors")
aa("--pre_reorder_k", default="0,10,100,1000", help="values for reorder_k")
aa("--nprobe", default="1,2,5,10,20,50,100,200", help="values for nprobe")
aa("--nrun", default=5, type=int, help="nb of runs to perform")
args = parser.parse_args()
print("args:", args)
pre_reorder_k_tab = [int(x) for x in args.pre_reorder_k.split(',')]
nprobe_tab = [int(x) for x in args.nprobe.split(',')]
pre_reorder_k_tab = [int(x) for x in args.pre_reorder_k.split(",")]
nprobe_tab = [int(x) for x in args.nprobe.split(",")]
os.system('echo -n "nb processors "; '
'cat /proc/cpuinfo | grep ^processor | wc -l; '
'cat /proc/cpuinfo | grep ^"model name" | tail -1')
os.system(
'echo -n "nb processors "; '
"cat /proc/cpuinfo | grep ^processor | wc -l; "
'cat /proc/cpuinfo | grep ^"model name" | tail -1'
)
cache_dir = args.base_dir + "/" + args.db + "/"
k = args.k
@@ -80,6 +87,7 @@ def main():
if not os.path.exists(cache_dir + "xb.npy"):
# prepare cache
from datasets import load_dataset
ds = load_dataset(args.db, download=args.download)
print(ds)
# store for SCANN
@@ -95,7 +103,7 @@ def main():
np.save(fname, v)
open(cache_dir + "metric", "w").write(ds.metric)
dataset = {}
for kn in "xb xq gt".split():
fname = cache_dir + "/" + kn + ".npy"
@@ -103,55 +111,69 @@ def main():
dataset[kn] = np.load(fname)
xb = dataset["xb"]
xq = dataset["xq"]
gt = dataset["gt"]
gt = dataset["gt"]
distance_measure = open(cache_dir + "metric").read()
if args.lib == "faiss":
import faiss
name1_to_metric = {
"IP": faiss.METRIC_INNER_PRODUCT,
"L2": faiss.METRIC_L2
"L2": faiss.METRIC_L2,
}
index_fname = cache_dir + "index.faiss"
if not os.path.exists(index_fname):
index = faiss_make_index(
xb, name1_to_metric[distance_measure], index_fname)
xb, name1_to_metric[distance_measure], index_fname
)
else:
index = faiss.read_index(index_fname)
faiss_eval_search(
index, xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt,
nrun, args.measure
index,
xq,
xb,
nprobe_tab,
pre_reorder_k_tab,
k,
gt,
nrun,
args.measure,
)
if args.lib == "scann":
from scann.scann_ops.py import scann_ops_pybind
name1_to_name2 = {
"IP": "dot_product",
"L2": "squared_l2"
}
name1_to_name2 = {"IP": "dot_product", "L2": "squared_l2"}
scann_dir = cache_dir + "/scann1.1.1_serialized"
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)
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)
xb, name1_to_name2[distance_measure], scann_dir, 100
)
scann_eval_search(
searcher, searcher_reo,
xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt,
nrun, args.measure
searcher,
searcher_reo,
xq,
xb,
nprobe_tab,
pre_reorder_k_tab,
k,
gt,
nrun,
args.measure,
)
if args.lib != "scann" and args.thenscann:
@@ -159,9 +181,10 @@ def main():
# options
cmdline = " ".join(sys.argv) + " --lib scann"
cmdline = (
". ~/anaconda3/etc/profile.d/conda.sh ; " +
"conda activate scann_1.1.1; "
"python -u " + cmdline)
". ~/anaconda3/etc/profile.d/conda.sh ; "
+ "conda activate scann_1.1.1; "
"python -u " + cmdline
)
print("running", cmdline)
@@ -183,10 +206,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 = scann.scann_ops_pybind.builder(xb, k, distance_measure)
sb = sb.tree(
num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
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:
@@ -205,9 +228,17 @@ def scann_make_index(xb, distance_measure, scann_dir, reorder_k):
def scann_eval_search(
searcher, searcher_reo,
xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt,
nrun, measure):
searcher,
searcher_reo,
xq,
xb,
nprobe_tab,
pre_reorder_k_tab,
k,
gt,
nrun,
measure,
):
# warmup
for _run in range(5):
@@ -228,8 +259,10 @@ def scann_eval_search(
else:
t0 = time.time()
I, D = searcher_reo.search_batched(
xq, leaves_to_search=nprobe, final_num_neighbors=k,
pre_reorder_num_neighbors=pre_reorder_k
xq,
leaves_to_search=nprobe,
final_num_neighbors=k,
pre_reorder_num_neighbors=pre_reorder_k,
)
t1 = time.time()
@@ -267,13 +300,16 @@ def faiss_make_index(xb, metric_type, fname):
def faiss_eval_search(
index, xq, xb, nprobe_tab, pre_reorder_k_tab,
k, gt, nrun, measure
):
index, xq, xb, nprobe_tab, pre_reorder_k_tab, k, gt, nrun, measure
):
import faiss
print("use precomputed table=", index.use_precomputed_table,
"by residual=", index.by_residual)
print(
"use precomputed table=",
index.use_precomputed_table,
"by residual=",
index.by_residual,
)
print("adding a refine index")
index_refine = faiss.IndexRefineFlat(index, faiss.swig_ptr(xb))
benchs/bench_all_ivf/datasets_oss.py
+23 -22
View File
@@ -15,11 +15,11 @@ from faiss.contrib import datasets as faiss_datasets
print("path:", faiss_datasets.__file__)
faiss_datasets.dataset_basedir = '/checkpoint/matthijs/simsearch/'
faiss_datasets.dataset_basedir = "/checkpoint/matthijs/simsearch/"
def sanitize(x):
return np.ascontiguousarray(x, dtype='float32')
return np.ascontiguousarray(x, dtype="float32")
#################################################################
@@ -53,44 +53,46 @@ class DatasetCentroids(faiss_datasets.Dataset):
def get_groundtruth(self, k=100):
return faiss.knn(
self.xq, self.xb, k,
faiss.METRIC_L2 if self.metric == 'L2' else faiss.METRIC_INNER_PRODUCT
self.xq,
self.xb,
k,
(
faiss.METRIC_L2
if self.metric == "L2"
else faiss.METRIC_INNER_PRODUCT
),
)[1]
def load_dataset(dataset='deep1M', compute_gt=False, download=False):
def load_dataset(dataset="deep1M", compute_gt=False, download=False):
print("load data", dataset)
if dataset == 'sift1M':
if dataset == "sift1M":
return faiss_datasets.DatasetSIFT1M()
elif dataset.startswith('bigann'):
elif dataset.startswith("bigann"):
dbsize = 1000 if dataset == "bigann1B" else int(dataset[6:-1])
return faiss_datasets.DatasetBigANN(nb_M=dbsize)
elif dataset.startswith("deep_centroids_"):
ncent = int(dataset[len("deep_centroids_"):])
ncent = int(dataset[len("deep_centroids_") :])
centdir = "/checkpoint/matthijs/bench_all_ivf/precomputed_clusters"
return DatasetCentroids(
faiss_datasets.DatasetDeep1B(nb=1000000),
f"{centdir}/clustering.dbdeep1M.IVF{ncent}.faissindex"
f"{centdir}/clustering.dbdeep1M.IVF{ncent}.faissindex",
)
elif dataset.startswith("deep"):
szsuf = dataset[4:]
if szsuf[-1] == 'M':
dbsize = 10 ** 6 * int(szsuf[:-1])
elif szsuf == '1B':
dbsize = 10 ** 9
elif szsuf[-1] == 'k':
if szsuf[-1] == "M":
dbsize = 10**6 * int(szsuf[:-1])
elif szsuf == "1B":
dbsize = 10**9
elif szsuf[-1] == "k":
dbsize = 1000 * int(szsuf[:-1])
else:
assert False, "did not recognize suffix " + szsuf
@@ -117,7 +119,7 @@ def evaluate_DI(D, I, gt):
rank = 1
while rank <= k:
recall = (I[:, :rank] == gt[:, :1]).sum() / float(nq)
print("R@%d: %.4f" % (rank, recall), end=' ')
print("R@%d: %.4f" % (rank, recall), end=" ")
rank *= 10
@@ -126,12 +128,11 @@ def evaluate(xq, gt, index, k=100, endl=True):
D, I = index.search(xq, k)
t1 = time.time()
nq = xq.shape[0]
print("\t %8.4f ms per query, " % (
(t1 - t0) * 1000.0 / nq), end=' ')
print("\t %8.4f ms per query, " % ((t1 - t0) * 1000.0 / nq), end=" ")
rank = 1
while rank <= k:
recall = (I[:, :rank] == gt[:, :1]).sum() / float(nq)
print("R@%d: %.4f" % (rank, recall), end=' ')
print("R@%d: %.4f" % (rank, recall), end=" ")
rank *= 10
if endl:
print()
benchs/bench_all_ivf/make_groundtruth.py
+2 -6
View File
@@ -7,7 +7,7 @@ import logging
# https://stackoverflow.com/questions/7016056/python-logging-not-outputting-anything
logging.basicConfig()
logger = logging.getLogger('faiss.contrib.exhaustive_search')
logger = logging.getLogger("faiss.contrib.exhaustive_search")
logger.setLevel(logging.INFO)
from faiss.contrib import datasets
@@ -18,10 +18,6 @@ ds = datasets.DatasetDeep1B(nb=int(1e9))
print("computing GT matches for", ds)
D, I = knn_ground_truth(
ds.get_queries(),
ds.database_iterator(bs=65536),
k=100
)
D, I = knn_ground_truth(ds.get_queries(), ds.database_iterator(bs=65536), k=100)
vecs_io.ivecs_write("/tmp/tt.ivecs", I)
benchs/bench_all_ivf/parse_bench_all_ivf.py
+134 -117
View File
@@ -12,10 +12,10 @@ from faiss.contrib.factory_tools import get_code_size as unitsize
def dbsize_from_name(dbname):
sufs = {
'1B': 10**9,
'100M': 10**8,
'10M': 10**7,
'1M': 10**6,
"1B": 10**9,
"100M": 10**8,
"10M": 10**7,
"1M": 10**6,
}
for s in sufs:
if dbname.endswith(s):
@@ -25,7 +25,7 @@ def dbsize_from_name(dbname):
def keep_latest_stdout(fnames):
fnames = [fname for fname in fnames if fname.endswith('.stdout')]
fnames = [fname for fname in fnames if fname.endswith(".stdout")]
fnames.sort()
n = len(fnames)
fnames2 = []
@@ -42,7 +42,7 @@ def parse_result_file(fname):
res = []
keys = []
stats = {}
stats['run_version'] = fname[-8]
stats["run_version"] = fname[-8]
indexkey = None
for l in open(fname):
if l.startswith("srun:"):
@@ -57,42 +57,42 @@ def parse_result_file(fname):
stats["nq"] = int(fi[9])
stats["nb"] = int(fi[11])
stats["nt"] = int(fi[13])
if l.startswith('index size on disk:'):
stats['index_size'] = int(l.split()[-1])
if l.startswith('current RSS:'):
stats['RSS'] = int(l.split()[-1])
if l.startswith('precomputed tables size:'):
stats['tables_size'] = int(l.split()[-1])
if l.startswith('Setting nb of threads to'):
stats['n_threads'] = int(l.split()[-1])
if l.startswith(' add in'):
stats['add_time'] = float(l.split()[-2])
if l.startswith("index size on disk:"):
stats["index_size"] = int(l.split()[-1])
if l.startswith("current RSS:"):
stats["RSS"] = int(l.split()[-1])
if l.startswith("precomputed tables size:"):
stats["tables_size"] = int(l.split()[-1])
if l.startswith("Setting nb of threads to"):
stats["n_threads"] = int(l.split()[-1])
if l.startswith(" add in"):
stats["add_time"] = float(l.split()[-2])
if l.startswith("vector code_size"):
stats['code_size'] = float(l.split()[-1])
if l.startswith('args:'):
args = eval(l[l.find(' '):])
stats["code_size"] = float(l.split()[-1])
if l.startswith("args:"):
args = eval(l[l.find(" ") :])
indexkey = args.indexkey
elif "time(ms/q)" in l:
# result header
if 'R@1 R@10 R@100' in l:
if "R@1 R@10 R@100" in l:
stats["measure"] = "recall"
stats["ranks"] = [1, 10, 100]
elif 'I@1 I@10 I@100' in l:
elif "I@1 I@10 I@100" in l:
stats["measure"] = "inter"
stats["ranks"] = [1, 10, 100]
elif 'inter@' in l:
elif "inter@" in l:
stats["measure"] = "inter"
fi = l.split()
if fi[1] == "inter@":
rank = int(fi[2])
else:
rank = int(fi[1][len("inter@"):])
rank = int(fi[1][len("inter@") :])
stats["ranks"] = [rank]
else:
assert False
st = 1
elif 'index size on disk:' in l:
elif "index size on disk:" in l:
stats["index_size"] = int(l.split()[-1])
elif st == 1:
st = 2
@@ -105,12 +105,13 @@ def parse_result_file(fname):
res.append([float(x) for x in fi[1:]])
return indexkey, np.array(res), keys, stats
# the directory used in run_on_cluster.bash
basedir = "/checkpoint/matthijs/bench_all_ivf/"
logdir = basedir + 'logs/'
logdir = basedir + "logs/"
def collect_results_for(db='deep1M', prefix="autotune."):
def collect_results_for(db="deep1M", prefix="autotune."):
# run parsing
allres = {}
allstats = {}
@@ -124,10 +125,10 @@ def collect_results_for(db='deep1M', prefix="autotune."):
for fname in fnames:
if not (
'db' + db in fname and
fname.startswith(prefix) and
fname.endswith('.stdout')
):
"db" + db in fname
and fname.startswith(prefix)
and fname.endswith(".stdout")
):
continue
print("parse", fname, end=" ", flush=True)
try:
@@ -144,7 +145,7 @@ def collect_results_for(db='deep1M', prefix="autotune."):
missing.append(fname)
else:
if indexkey in allres:
if allstats[indexkey]['run_version'] > stats['run_version']:
if allstats[indexkey]["run_version"] > stats["run_version"]:
# don't use this run
continue
@@ -160,12 +161,7 @@ def extract_pareto_optimal(allres, keys, recall_idx=0, times_idx=3):
v = allres[k]
perf = v[:, recall_idx]
times = v[:, times_idx]
bigtab.append(
np.vstack((
np.ones(times.size) * i,
perf, times
))
)
bigtab.append(np.vstack((np.ones(times.size) * i, perf, times)))
if bigtab == []:
return [], np.zeros((3, 0))
@@ -175,10 +171,9 @@ def extract_pareto_optimal(allres, keys, recall_idx=0, times_idx=3):
perm = np.argsort(bigtab[1, :])
bigtab_sorted = bigtab[:, perm]
best_times = np.minimum.accumulate(bigtab_sorted[2, ::-1])[::-1]
selection, = np.where(bigtab_sorted[2, :] == best_times)
(selection,) = np.where(bigtab_sorted[2, :] == best_times)
selected_keys = [
keys[i] for i in
np.unique(bigtab_sorted[0, selection].astype(int))
keys[i] for i in np.unique(bigtab_sorted[0, selection].astype(int))
]
ops = bigtab_sorted[:, selection]
@@ -186,8 +181,13 @@ def extract_pareto_optimal(allres, keys, recall_idx=0, times_idx=3):
def plot_subset(
allres, allstats, selected_methods, recall_idx, times_idx=3,
report=["overhead", "build time"]):
allres,
allstats,
selected_methods,
recall_idx,
times_idx=3,
report=["overhead", "build time"],
):
# important methods
for k in selected_methods:
@@ -197,20 +197,20 @@ def plot_subset(
d = stats["d"]
dbsize = stats["nb"]
if "index_size" in stats and "tables_size" in stats:
tot_size = stats['index_size'] + stats['tables_size']
tot_size = stats["index_size"] + stats["tables_size"]
else:
tot_size = -1
id_size = 8 # 64 bit
id_size = 8 # 64 bit
addt = ''
if 'add_time' in stats:
add_time = stats['add_time']
addt = ""
if "add_time" in stats:
add_time = stats["add_time"]
if add_time > 7200:
add_min = add_time / 60
addt = ', %dh%02d' % (add_min / 60, add_min % 60)
addt = ", %dh%02d" % (add_min / 60, add_min % 60)
else:
add_sec = int(add_time)
addt = ', %dm%02d' % (add_sec / 60, add_sec % 60)
addt = ", %dm%02d" % (add_sec / 60, add_sec % 60)
code_size = unitsize(d, k)
@@ -222,33 +222,37 @@ def plot_subset(
tight_size = (code_size + id_size) * dbsize
if tot_size < 0 or "overhead" not in report:
pass # don't know what the index size is
pass # don't know what the index size is
elif tot_size > 10 * tight_size:
label += " overhead x%.1f" % (tot_size / tight_size)
else:
label += " overhead+%.1f%%" % (
tot_size / tight_size * 100 - 100)
label += " overhead+%.1f%%" % (tot_size / tight_size * 100 - 100)
if "build time" in report:
label += " " + addt
linestyle = (':' if 'Refine' in k or 'RFlat' in k else
'-.' if 'SQ' in k else
'-' if '4fs' in k else
'-')
linestyle = (
":"
if "Refine" in k or "RFlat" in k
else "-." if "SQ" in k else "-" if "4fs" in k else "-"
)
print(k, linestyle)
pyplot.semilogy(v[:, recall_idx], 1000 / v[:, times_idx], label=label,
linestyle=linestyle,
marker='o' if '4fs' in k else '+')
pyplot.semilogy(
v[:, recall_idx],
1000 / v[:, times_idx],
label=label,
linestyle=linestyle,
marker="o" if "4fs" in k else "+",
)
recall_rank = stats["ranks"][recall_idx]
if stats["measure"] == "recall":
pyplot.xlabel('1-recall at %d' % recall_rank)
pyplot.xlabel("1-recall at %d" % recall_rank)
elif stats["measure"] == "inter":
pyplot.xlabel('inter @ %d' % recall_rank)
pyplot.xlabel("inter @ %d" % recall_rank)
else:
assert False
pyplot.ylabel('QPS (%d threads)' % stats["n_threads"])
pyplot.ylabel("QPS (%d threads)" % stats["n_threads"])
def plot_tradeoffs(db, allres, allstats, code_size, recall_rank):
@@ -275,66 +279,76 @@ def plot_tradeoffs(db, allres, allstats, code_size, recall_rank):
for k in sorted(allres):
v = allres[k]
if v.ndim != 2: continue
if v.ndim != 2:
continue
us = unitsize(d, k)
if not code_size[0] <= us <= code_size[1]: continue
if not code_size[0] <= us <= code_size[1]:
continue
names_maxperf.append((v[-1, recall_idx], k))
# sort from lowest to highest topline accuracy
names_maxperf.sort()
names = [name for mp, name in names_maxperf]
selected_methods, optimal_points = \
extract_pareto_optimal(allres, names, recall_idx, times_idx)
selected_methods, optimal_points = extract_pareto_optimal(
allres, names, recall_idx, times_idx
)
not_selected = list(set(names) - set(selected_methods))
print("methods without an optimal OP: ", not_selected)
pyplot.title('database ' + db + ' ' + code_size_name)
pyplot.title("database " + db + " " + code_size_name)
# grayed out lines
for k in not_selected:
v = allres[k]
if v.ndim != 2: continue
if v.ndim != 2:
continue
us = unitsize(d, k)
if not code_size[0] <= us <= code_size[1]: continue
if not code_size[0] <= us <= code_size[1]:
continue
linestyle = (':' if 'PQ' in k else
'-.' if 'SQ4' in k else
'--' if 'SQ8' in k else '-')
linestyle = (
":"
if "PQ" in k
else "-." if "SQ4" in k else "--" if "SQ8" in k else "-"
)
pyplot.semilogy(v[:, recall_idx], 1000 / v[:, times_idx], label=None,
linestyle=linestyle,
marker='o' if 'HNSW' in k else '+',
color='#cccccc', linewidth=0.2)
pyplot.semilogy(
v[:, recall_idx],
1000 / v[:, times_idx],
label=None,
linestyle=linestyle,
marker="o" if "HNSW" in k else "+",
color="#cccccc",
linewidth=0.2,
)
plot_subset(allres, allstats, selected_methods, recall_idx, times_idx)
if len(not_selected) == 0:
om = ''
om = ""
else:
om = '\nomitted:'
om = "\nomitted:"
nc = len(om)
for m in not_selected:
if nc > 80:
om += '\n'
om += "\n"
nc = 0
om += ' ' + m
om += " " + m
nc += len(m) + 1
# pyplot.semilogy(optimal_points[1, :], optimal_points[2, :], marker="s")
# print(optimal_points[0, :])
pyplot.xlabel('1-recall at %d %s' % (recall_rank, om) )
pyplot.ylabel('QPS (%d threads)' % n_threads)
pyplot.xlabel("1-recall at %d %s" % (recall_rank, om))
pyplot.ylabel("QPS (%d threads)" % n_threads)
pyplot.legend()
pyplot.grid()
return selected_methods, not_selected
if __name__ == "__main__xx":
# tests on centroids indexing (v1)
@@ -342,16 +356,17 @@ if __name__ == "__main__xx":
pyplot.gcf().set_size_inches(15, 10)
i = 1
for ncent in 65536, 262144, 1048576, 4194304:
db = f'deep_centroids_{ncent}.k{k}.'
allres, allstats = collect_results_for(
db=db, prefix="cent_index.")
db = f"deep_centroids_{ncent}.k{k}."
allres, allstats = collect_results_for(db=db, prefix="cent_index.")
pyplot.subplot(2, 2, i)
plot_subset(
allres, allstats, list(allres.keys()),
allres,
allstats,
list(allres.keys()),
recall_idx=0,
times_idx=1,
report=["code_size"]
report=["code_size"],
)
i += 1
pyplot.title(f"{ncent} centroids")
@@ -359,7 +374,7 @@ if __name__ == "__main__xx":
pyplot.xlim([0.95, 1])
pyplot.grid()
pyplot.savefig('figs/deep1B_centroids_k%d.png' % k)
pyplot.savefig("figs/deep1B_centroids_k%d.png" % k)
if __name__ == "__main__xx":
@@ -367,18 +382,18 @@ if __name__ == "__main__xx":
pyplot.gcf().set_size_inches(15, 10)
i=1
i = 1
for ncent in 65536, 262144, 1048576, 4194304:
xyd = defaultdict(list)
for k in 1, 4, 8, 16, 32, 64, 128, 256:
db = f'deep_centroids_{ncent}.k{k}.'
db = f"deep_centroids_{ncent}.k{k}."
allres, allstats = collect_results_for(db=db, prefix="cent_index.")
for indexkey, res in allres.items():
idx, = np.where(res[:, 0] >= 0.99)
(idx,) = np.where(res[:, 0] >= 0.99)
if idx.size > 0:
xyd[indexkey].append((k, 1000 / res[idx[0], 1]))
@@ -386,43 +401,44 @@ if __name__ == "__main__xx":
i += 1
for indexkey, xy in xyd.items():
xy = np.array(xy)
pyplot.loglog(xy[:, 0], xy[:, 1], 'o-', label=indexkey)
pyplot.loglog(xy[:, 0], xy[:, 1], "o-", label=indexkey)
pyplot.title(f"{ncent} centroids")
pyplot.xlabel("k")
xt = 2**np.arange(9)
xt = 2 ** np.arange(9)
pyplot.xticks(xt, ["%d" % x for x in xt])
pyplot.ylabel("QPS (32 threads)")
pyplot.legend()
pyplot.grid()
pyplot.savefig('../plots/deep1B_centroids_min99.png')
pyplot.savefig("../plots/deep1B_centroids_min99.png")
if __name__ == "__main__xx":
# main indexing plots
i = 0
for db in 'bigann10M', 'deep10M', 'bigann100M', 'deep100M', 'deep1B', 'bigann1B':
allres, allstats = collect_results_for(
db=db, prefix="autotune.")
for db in (
"bigann10M",
"deep10M",
"bigann100M",
"deep100M",
"deep1B",
"bigann1B",
):
allres, allstats = collect_results_for(db=db, prefix="autotune.")
for cs in 8, 16, 32, 64:
pyplot.figure(i)
i += 1
pyplot.gcf().set_size_inches(15, 10)
cs_range = (
(0, 8) if cs == 8 else (cs // 2 + 1, cs)
)
cs_range = (0, 8) if cs == 8 else (cs // 2 + 1, cs)
plot_tradeoffs(
db, allres, allstats, code_size=cs_range, recall_rank=1)
pyplot.savefig('../plots/tradeoffs_%s_cs%d_r1.png' % (
db, cs))
db, allres, allstats, code_size=cs_range, recall_rank=1
)
pyplot.savefig("../plots/tradeoffs_%s_cs%d_r1.png" % (db, cs))
if __name__ == "__main__":
@@ -434,7 +450,7 @@ if __name__ == "__main__":
i += 1
allres, allstats = collect_results_for(db=db, prefix="autotune.")
plot_tradeoffs(db, allres, allstats, code_size=0, recall_rank=1)
pyplot.savefig('../plots/1M_tradeoffs_' + db + ".png")
pyplot.savefig("../plots/1M_tradeoffs_" + db + ".png")
for db in "sift1M", "deep1M":
allres, allstats = collect_results_for(db=db, prefix="autotune.")
@@ -442,18 +458,19 @@ if __name__ == "__main__":
pyplot.gcf().set_size_inches(15, 10)
i += 1
plot_tradeoffs(db, allres, allstats, code_size=(0, 64), recall_rank=1)
pyplot.savefig('../plots/1M_tradeoffs_' + db + "_small.png")
pyplot.savefig("../plots/1M_tradeoffs_" + db + "_small.png")
pyplot.figure(i)
pyplot.gcf().set_size_inches(15, 10)
i += 1
plot_tradeoffs(db, allres, allstats, code_size=(65, 10000), recall_rank=1)
pyplot.savefig('../plots/1M_tradeoffs_' + db + "_large.png")
plot_tradeoffs(
db, allres, allstats, code_size=(65, 10000), recall_rank=1
)
pyplot.savefig("../plots/1M_tradeoffs_" + db + "_large.png")
if __name__ == "__main__xx":
db = 'sift1M'
db = "sift1M"
allres, allstats = collect_results_for(db=db, prefix="autotune.")
pyplot.gcf().set_size_inches(15, 10)
@@ -463,7 +480,7 @@ if __name__ == "__main__xx":
"IVF1024,PQ64x4fs",
"IVF1024,PQ64x4fsr",
"IVF1024,SQ4",
"IVF1024,SQ8"
"IVF1024,SQ8",
]
plot_subset(allres, allstats, keys, recall_idx=0, report=["code_size"])
@@ -474,7 +491,7 @@ if __name__ == "__main__xx":
pyplot.ylabel("QPS (32 threads)")
pyplot.grid()
pyplot.savefig('../plots/ivf1024_variants.png')
pyplot.savefig("../plots/ivf1024_variants.png")
pyplot.figure(2)
pyplot.gcf().set_size_inches(15, 10)
@@ -496,4 +513,4 @@ if __name__ == "__main__xx":
pyplot.ylabel("QPS (32 threads)")
pyplot.grid()
pyplot.savefig('../plots/ivf1024_rerank.png')
pyplot.savefig("../plots/ivf1024_rerank.png")
benchs/bench_big_batch_ivf.py
+11 -12
View File
@@ -20,17 +20,17 @@ def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('dataset options')
aa('--dim', type=int, default=64)
aa('--size', default="S")
group = parser.add_argument_group("dataset options")
aa("--dim", type=int, default=64)
aa("--size", default="S")
group = parser.add_argument_group('index options')
aa('--nlist', type=int, default=100)
aa('--factory_string', default="", help="overrides nlist")
aa('--k', type=int, default=10)
aa('--nprobe', type=int, default=5)
aa('--nt', type=int, default=-1, help="nb search threads")
aa('--method', default="pairwise_distances", help="")
group = parser.add_argument_group("index options")
aa("--nlist", type=int, default=100)
aa("--factory_string", default="", help="overrides nlist")
aa("--k", type=int, default=10)
aa("--nprobe", type=int, default=5)
aa("--nt", type=int, default=-1, help="nb search threads")
aa("--method", default="pairwise_distances", help="")
args = parser.parse_args()
print("args:", args)
@@ -98,8 +98,7 @@ t_ref = toc()
tic("block search")
Dnew, Inew = big_batch_search(
index, ds.get_queries(),
k, method=args.method, verbose=10
index, ds.get_queries(), k, method=args.method, verbose=10
)
t_tot = toc()
benchs/bench_flat_l2_panorama.py
+7 -9
View File
@@ -17,9 +17,7 @@ except ImportError:
from faiss.contrib.datasets import DatasetSIFT1M, DatasetGIST1M
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset", default="gist1m", choices=["sift1m", "gist1m"]
)
parser.add_argument("--dataset", default="gist1m", choices=["sift1m", "gist1m"])
args = parser.parse_args()
if args.dataset == "sift1m":
@@ -91,15 +89,15 @@ for name in names:
qps_values.append(qps)
x = np.arange(len(names))
plt.bar(x, qps_values, color=['#1f77b4', '#ff7f0e'])
plt.bar(x, qps_values, color=["#1f77b4", "#ff7f0e"])
speedup = qps_values[1] / qps_values[0]
ax = plt.gca()
ax.text(
x[1],
qps_values[1] * 1.01,
f"{speedup:.2f}x",
ha="center",
va="bottom",
x[1],
qps_values[1] * 1.01,
f"{speedup:.2f}x",
ha="center",
va="bottom",
)
plt.xticks(x, labels, rotation=0)
plt.ylabel("QPS")
benchs/bench_for_interrupt.py
+49 -40
View File
@@ -18,15 +18,16 @@ parser = argparse.ArgumentParser()
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('dataset options')
aa('--dim', type=int, default=64)
aa('--nb', type=int, default=int(1e6))
aa('--subset_len', type=int, default=int(1e5))
aa('--key', default='IVF1000,Flat')
aa('--nprobe', type=int, default=640)
aa('--no_intcallback', default=False, action='store_true')
aa('--twostage', default=False, action='store_true')
aa('--nt', type=int, default=-1)
group = parser.add_argument_group("dataset options")
aa("--dim", type=int, default=64)
aa("--nb", type=int, default=int(1e6))
aa("--subset_len", type=int, default=int(1e5))
aa("--key", default="IVF1000,Flat")
aa("--nprobe", type=int, default=640)
aa("--no_intcallback", default=False, action="store_true")
aa("--twostage", default=False, action="store_true")
aa("--nt", type=int, default=-1)
args = parser.parse_args()
@@ -41,9 +42,9 @@ subset_len = args.subset_len
np.random.seed(1234) # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
xt = np.random.random((nt, d)).astype('float32')
xb = np.random.random((nb, d)).astype("float32")
xq = np.random.random((nq, d)).astype("float32")
xt = np.random.random((nt, d)).astype("float32")
k = 100
if args.no_intcallback:
@@ -54,41 +55,41 @@ if args.nt != -1:
nprobe = args.nprobe
key = args.key
#key = 'IVF1000,Flat'
# key = 'IVF1000,Flat'
# key = 'IVF1000,PQ64'
# key = 'IVF100_HNSW32,PQ64'
# faiss.omp_set_num_threads(1)
pf = 'dim%d_' % d
pf = "dim%d_" % d
if d == 64:
pf = ''
pf = ""
basename = '/tmp/base%s%s.index' % (pf, key)
basename = "/tmp/base%s%s.index" % (pf, key)
if os.path.exists(basename):
print('load', basename)
print("load", basename)
index_1 = faiss.read_index(basename)
else:
print('train + write', basename)
print("train + write", basename)
index_1 = faiss.index_factory(d, key)
index_1.train(xt)
faiss.write_index(index_1, basename)
print('add')
print("add")
index_1.add(xb)
print('set nprobe=', nprobe)
faiss.ParameterSpace().set_index_parameter(index_1, 'nprobe', nprobe)
print("set nprobe=", nprobe)
faiss.ParameterSpace().set_index_parameter(index_1, "nprobe", nprobe)
class ResultHeap:
""" Combine query results from a sliced dataset """
"""Combine query results from a sliced dataset"""
def __init__(self, nq, k):
" nq: number of query vectors, k: number of results per query "
self.I = np.zeros((nq, k), dtype='int64')
self.D = np.zeros((nq, k), dtype='float32')
"nq: number of query vectors, k: number of results per query"
self.I = np.zeros((nq, k), dtype="int64")
self.D = np.zeros((nq, k), dtype="float32")
self.nq, self.k = nq, k
heaps = faiss.float_maxheap_array_t()
heaps.k = k
@@ -103,30 +104,37 @@ class ResultHeap:
assert I.shape == (self.nq, self.k)
I += i0
self.heaps.addn_with_ids(
self.k, faiss.swig_ptr(D),
faiss.swig_ptr(I), self.k)
self.k, faiss.swig_ptr(D), faiss.swig_ptr(I), self.k
)
def finalize(self):
self.heaps.reorder()
stats = faiss.cvar.indexIVF_stats
stats.reset()
print('index size', index_1.ntotal,
'imbalance', index_1.invlists.imbalance_factor())
print(
"index size",
index_1.ntotal,
"imbalance",
index_1.invlists.imbalance_factor(),
)
start = time.time()
Dref, Iref = index_1.search(xq, k)
print('time of searching: %.3f s = %.3f + %.3f ms' % (
time.time() - start, stats.quantization_time, stats.search_time))
print(
"time of searching: %.3f s = %.3f + %.3f ms"
% (time.time() - start, stats.quantization_time, stats.search_time)
)
indexes = {}
if args.twostage:
for i in range(0, nb, subset_len):
index = faiss.read_index(basename)
faiss.ParameterSpace().set_index_parameter(index, 'nprobe', nprobe)
print("add %d:%d" %(i, i+subset_len))
index.add(xb[i:i + subset_len])
faiss.ParameterSpace().set_index_parameter(index, "nprobe", nprobe)
print("add %d:%d" % (i, i + subset_len))
index.add(xb[i : i + subset_len])
indexes[i] = index
rh = ResultHeap(nq, k)
@@ -134,9 +142,9 @@ sum_time = tq = ts = 0
for i in range(0, nb, subset_len):
if not args.twostage:
index = faiss.read_index(basename)
faiss.ParameterSpace().set_index_parameter(index, 'nprobe', nprobe)
print("add %d:%d" %(i, i+subset_len))
index.add(xb[i:i + subset_len])
faiss.ParameterSpace().set_index_parameter(index, "nprobe", nprobe)
print("add %d:%d" % (i, i + subset_len))
index.add(xb[i : i + subset_len])
else:
index = indexes[i]
@@ -148,9 +156,10 @@ for i in range(0, nb, subset_len):
ts += stats.search_time
rh.add_batch_result(Di, Ii, i)
print('time of searching separately: %.3f s = %.3f + %.3f ms' %
(sum_time, tq, ts))
print(
"time of searching separately: %.3f s = %.3f + %.3f ms" % (sum_time, tq, ts)
)
rh.finalize()
print('diffs: %d / %d' % ((Iref != rh.I).sum(), Iref.size))
print("diffs: %d / %d" % ((Iref != rh.I).sum(), Iref.size))
benchs/bench_fw/benchmark.py
+48 -23
View File
@@ -239,7 +239,8 @@ class TrainOperator(IndexOperator):
if codec_desc.factory is not None:
assert (
codec_desc.factory == "Flat" or codec_desc.training_vectors is not None
codec_desc.factory == "Flat"
or codec_desc.training_vectors is not None
)
index = IndexFromFactory(
num_threads=self.num_threads,
@@ -257,7 +258,10 @@ class TrainOperator(IndexOperator):
assert codec_desc.is_trained()
def train_one(
self, codec_desc: CodecDescriptor, results: Dict[str, Any], dry_run=False
self,
codec_desc: CodecDescriptor,
results: Dict[str, Any],
dry_run=False,
):
faiss.omp_set_num_threads(codec_desc.num_threads)
self.build_index_wrapper(codec_desc)
@@ -390,7 +394,9 @@ class SearchOperator(IndexOperator):
knn_desc.index.get_index()
def range_search_reference(self, index, parameters, range_metric, query_dataset):
def range_search_reference(
self, index, parameters, range_metric, query_dataset
):
logger.info("range_search_reference: begin")
if isinstance(range_metric, list):
assert len(range_metric) > 0
@@ -428,7 +434,9 @@ class SearchOperator(IndexOperator):
coefficients_training_data,
)
def estimate_range(self, index, parameters, range_scoring_radius, query_dataset):
def estimate_range(
self, index, parameters, range_scoring_radius, query_dataset
):
D, I, R, P, _ = index.knn_search(
False,
parameters,
@@ -477,7 +485,9 @@ class SearchOperator(IndexOperator):
return None, None, None, None, None, requires
if range_search_metric_function is not None:
range_search_metric = range_search_metric_function(R)
range_search_pr = range_search_pr_curve(D, range_search_metric, gt_rsm)
range_search_pr = range_search_pr_curve(
D, range_search_metric, gt_rsm
)
range_score_sum = np.sum(range_search_metric).item()
P |= {
"range_score_sum": range_score_sum,
@@ -716,8 +726,8 @@ class SearchOperator(IndexOperator):
assert requires is None
if (
knn_desc.range_ref_index_desc is None or
not knn_desc.index.supports_range_search()
knn_desc.range_ref_index_desc is None
or not knn_desc.index.supports_range_search()
):
return results, None
@@ -778,15 +788,17 @@ class SearchOperator(IndexOperator):
return results, None
def search(
self,
results: Dict[str, Any],
dry_run: bool = False,):
self,
results: Dict[str, Any],
dry_run: bool = False,
):
for knn_desc in self.knn_descs:
results, requires = self.search_one(
knn_desc=knn_desc,
results=results,
dry_run=dry_run,
range=self.range)
range=self.range,
)
if dry_run:
if requires is None:
continue
@@ -901,24 +913,29 @@ class ExecutionOperator:
def create_range_ref_knn(self, knn_desc):
if (
knn_desc.range_ref_index_desc is None or
not knn_desc.index.supports_range_search()
knn_desc.range_ref_index_desc is None
or not knn_desc.index.supports_range_search()
):
return
if knn_desc.range_ref_index_desc is not None:
ref_index_desc = (
self.search_op.get_desc(knn_desc.range_ref_index_desc)
ref_index_desc = self.search_op.get_desc(
knn_desc.range_ref_index_desc
)
if ref_index_desc is None:
raise ValueError(f"Unknown range index {knn_desc.range_ref_index_desc}")
raise ValueError(
f"Unknown range index {knn_desc.range_ref_index_desc}"
)
if ref_index_desc.range_metrics is None:
raise ValueError(
f"Range index {knn_desc.get_name()} has no radius_score"
)
results["metrics"] = {}
self.build_index_wrapper(ref_index_desc)
for metric_key, range_metric in ref_index_desc.range_metrics.items():
for (
metric_key,
range_metric,
) in ref_index_desc.range_metrics.items():
(
knn_desc.gt_radius,
range_search_metric_function,
@@ -962,8 +979,8 @@ class ExecutionOperator:
def execute(self, results: Dict[str, Any], dry_run: bool = False):
faiss.omp_set_num_threads(self.num_threads)
if self.train_op is not None:
results, requires = (
self.train_op.train(results=results, dry_run=dry_run)
results, requires = self.train_op.train(
results=results, dry_run=dry_run
)
if dry_run and requires:
return results, requires
@@ -975,8 +992,8 @@ class ExecutionOperator:
if not dry_run and self.compute_gt:
self.prepare_gt_or_range_knn(results)
results, requires = (
self.search_op.search(results=results, dry_run=dry_run)
results, requires = self.search_op.search(
results=results, dry_run=dry_run
)
if dry_run and requires:
return results, requires
@@ -1025,7 +1042,13 @@ class Benchmark:
raise ValueError("Failed to determine dimension of dataset")
def create_descriptors(
self, ci_desc: IndexDescriptorClassic, train, build, knn, reconstruct, range
self,
ci_desc: IndexDescriptorClassic,
train,
build,
knn,
reconstruct,
range,
):
codec_desc = None
index_desc = None
@@ -1126,7 +1149,9 @@ class Benchmark:
database_vectors=self.database_vectors,
query_vectors=self.query_vectors,
# index_descs=[self.get_flat_desc("Flat"), index_desc],
index_descs=[index_desc], # Should automatically find flat descriptors
index_descs=[
index_desc
], # Should automatically find flat descriptors
range_ref_index_desc=self.range_ref_index_desc,
k=self.k,
distance_metric=self.distance_metric,
benchs/bench_fw/descriptors.py
+43 -14
View File
@@ -142,9 +142,9 @@ class DatasetDescriptor:
assert self.tablename is not None
filename += self.tablename
if self.partitions is not None:
filename += "_" + "_".join(
self.partitions
).replace("=", "_").replace("/", "_")
filename += "_" + "_".join(self.partitions).replace(
"=", "_"
).replace("/", "_")
if self.num_vectors is not None:
filename += f"_{self.num_vectors}"
if self.filename_suffix is not None:
@@ -278,8 +278,12 @@ class CodecDescriptor(IndexBaseDescriptor):
name += f"d_{self.d}.{self.metric.upper()}."
if self.factory != "Flat":
assert self.training_vectors is not None
name += self.training_vectors.get_filename(CodecDescriptor.FILENAME_PREFIX)
name += IndexBaseDescriptor.param_dict_list_to_name(self.construction_params)
name += self.training_vectors.get_filename(
CodecDescriptor.FILENAME_PREFIX
)
name += IndexBaseDescriptor.param_dict_list_to_name(
self.construction_params
)
return name
def name_from_path(self):
@@ -287,16 +291,23 @@ class CodecDescriptor(IndexBaseDescriptor):
filename = os.path.basename(self.path)
ext = filename.split(".")[-1]
if filename.endswith(ext):
name = filename[:-len(ext)]
else: # should never hit this rather raise value error
name = filename[: -len(ext)]
else: # should never hit this rather raise value error
name = filename
return name
def alias(self, benchmark_io: BenchmarkIO):
if hasattr(benchmark_io, "bucket"):
return CodecDescriptor(desc_name=self.get_name(), bucket=benchmark_io.bucket, path=self.get_path(benchmark_io), d=self.d, metric=self.metric)
return CodecDescriptor(desc_name=self.get_name(), d=self.d, metric=self.metric)
return CodecDescriptor(
desc_name=self.get_name(),
bucket=benchmark_io.bucket,
path=self.get_path(benchmark_io),
d=self.d,
metric=self.metric,
)
return CodecDescriptor(
desc_name=self.get_name(), d=self.d, metric=self.metric
)
@dataclass
@@ -316,21 +327,39 @@ class IndexDescriptor(IndexBaseDescriptor):
def get_name(self) -> str:
if self.desc_name is None:
self.desc_name = self.codec_desc.get_name() + self.database_desc.get_filename(prefix=IndexDescriptor.FILENAME_PREFIX)
self.desc_name = (
self.codec_desc.get_name()
+ self.database_desc.get_filename(
prefix=IndexDescriptor.FILENAME_PREFIX
)
)
return self.desc_name
def flat_name(self):
if self.flat_desc_name is not None:
return self.flat_desc_name
self.flat_desc_name = self.codec_desc.flat_name() + self.database_desc.get_filename(prefix=IndexDescriptor.FILENAME_PREFIX)
self.flat_desc_name = (
self.codec_desc.flat_name()
+ self.database_desc.get_filename(
prefix=IndexDescriptor.FILENAME_PREFIX
)
)
return self.flat_desc_name
# alias is used to refer when index is uploaded to blobstore and referred again
def alias(self, benchmark_io: BenchmarkIO):
if hasattr(benchmark_io, "bucket"):
return IndexDescriptor(desc_name=self.get_name(), bucket=benchmark_io.bucket, path=self.get_path(benchmark_io), d=self.d, metric=self.metric)
return IndexDescriptor(desc_name=self.get_name(), d=self.d, metric=self.metric)
return IndexDescriptor(
desc_name=self.get_name(),
bucket=benchmark_io.bucket,
path=self.get_path(benchmark_io),
d=self.d,
metric=self.metric,
)
return IndexDescriptor(
desc_name=self.get_name(), d=self.d, metric=self.metric
)
@dataclass
benchs/bench_fw/index.py
+23 -25
View File
@@ -73,9 +73,8 @@ class IndexBase:
def set_index_param(index, name, val, assert_same=False):
index = faiss.downcast_index(index)
val = int(val)
if (
isinstance(index, faiss.IndexPreTransform)
or isinstance(index, faiss.IndexIDMap)
if isinstance(index, faiss.IndexPreTransform) or isinstance(
index, faiss.IndexIDMap
):
Index.set_index_param(index.index, name, val)
return
@@ -278,8 +277,8 @@ class IndexBase:
):
logger.info("knn_search: begin")
if (
search_parameters is not None and
search_parameters.get("snap", 0) == 1
search_parameters is not None
and search_parameters.get("snap", 0) == 1
):
query_vectors = self.snap(query_vectors)
filename = (
@@ -329,9 +328,9 @@ class IndexBase:
xq = self.io.get_dataset(query_vectors)
(D, I), t, _ = timer("knn_search", lambda: index.search(xq, k))
if (
self.is_flat() or
not hasattr(self, "database_vectors") or
(self.database_vectors is None)
self.is_flat()
or not hasattr(self, "database_vectors")
or (self.database_vectors is None)
): # TODO
R = D
else:
@@ -350,9 +349,7 @@ class IndexBase:
"nlist": int(stats.nlist // repeat),
"ndis": int(stats.ndis // repeat),
"nheap_updates": int(stats.nheap_updates // repeat),
"quantization_time": int(
stats.quantization_time // repeat
),
"quantization_time": int(stats.quantization_time // repeat),
"search_time": int(stats.search_time // repeat),
}
self.io.write_file(filename, ["D", "I", "R", "P"], [D, I, R, P])
@@ -482,14 +479,12 @@ class IndexBase:
):
logger.info("range_search: begin")
if (
search_parameters is not None and
search_parameters.get("snap", 0) == 1
search_parameters is not None
and search_parameters.get("snap", 0) == 1
):
query_vectors = self.snap(query_vectors)
filename = (
self.get_range_search_name(
search_parameters, query_vectors, radius
)
self.get_range_search_name(search_parameters, query_vectors, radius)
+ "zip"
)
if self.io.file_exist(filename):
@@ -527,9 +522,7 @@ class IndexBase:
R = D
else:
xb = self.io.get_dataset(self.database_vectors)
R = refine_distances_range(
lims, D, I, xq, xb, self.metric_type
)
R = refine_distances_range(lims, D, I, xq, xb, self.metric_type)
P = {
"time": t,
"radius": radius,
@@ -543,9 +536,7 @@ class IndexBase:
"nlist": int(stats.nlist // repeat),
"ndis": int(stats.ndis // repeat),
"nheap_updates": int(stats.nheap_updates // repeat),
"quantization_time": int(
stats.quantization_time // repeat
),
"quantization_time": int(stats.quantization_time // repeat),
"search_time": int(stats.search_time // repeat),
}
self.io.write_file(
@@ -681,7 +672,8 @@ class Index(IndexBase):
_, t, _ = timer(
"add_with_ids",
lambda: index.add_with_ids(
xb, np.arange(len(xb), dtype='int32')),
xb, np.arange(len(xb), dtype="int32")
),
once=True,
)
else:
@@ -934,7 +926,9 @@ class IndexFromFactory(Index):
assert self.training_vectors is not None
codec_name += self.training_vectors.get_filename("xt")
if self.construction_params is not None:
codec_name += IndexBaseDescriptor.param_dict_list_to_name(self.construction_params)
codec_name += IndexBaseDescriptor.param_dict_list_to_name(
self.construction_params
)
self.codec_name = codec_name
return self.codec_name
@@ -969,7 +963,11 @@ class IndexFromFactory(Index):
assert codec_size is not None
meta = {
"training_time": training_time,
"training_size": self.training_vectors.num_vectors if self.training_vectors else 0,
"training_size": (
self.training_vectors.num_vectors
if self.training_vectors
else 0
),
"codec_size": codec_size,
"sa_code_size": self.get_sa_code_size(codec),
"code_size": self.get_code_size(codec),
benchs/bench_fw/optimize.py
+15 -9
View File
@@ -70,9 +70,9 @@ class Optimizer:
evaluation="knn",
accuracy_metric="knn_intersection",
min_accuracy=min_accuracy,
name_filter=None
if include_flat
else (lambda n: not n.startswith("Flat")),
name_filter=(
None if include_flat else (lambda n: not n.startswith("Flat"))
),
pareto_mode=ParetoMode.GLOBAL,
pareto_metric=pareto_metric,
)
@@ -103,7 +103,9 @@ class Optimizer:
dry_run=False,
)
descs = [IndexDescriptorClassic(factory="Flat"),] + [
descs = [
IndexDescriptorClassic(factory="Flat"),
] + [
IndexDescriptorClassic(
factory="HNSW32",
construction_params=[{"efConstruction": 2**i}],
@@ -229,7 +231,8 @@ class Optimizer:
(None, "SQbf16"),
(None, "SQ8"),
(None, "SQ8_direct_signed"),
] + [
]
+ [
(f"OPQ{M}_{M * dim}", f"PQ{M}x{b}")
for M in [8, 12, 16, 32, 48, 64, 96, 128, 192, 256]
if d % M == 0
@@ -237,7 +240,8 @@ class Optimizer:
if M * dim <= d
for b in range(4, 14, 2)
if M * b < d * 8 # smaller than SQ8
] + [
]
+ [
(None, f"PQ{M}x{b}")
for M in [8, 12, 16, 32, 48, 64, 96, 128, 192, 256]
if d % M == 0
@@ -257,9 +261,11 @@ class Optimizer:
_, filtered = self.benchmark_and_filter_candidates(
index_descs=[
IndexDescriptorClassic(
factory=f"IVF{nlist},{pq}"
if opq is None
else f"{opq},IVF{nlist},{pq}",
factory=(
f"IVF{nlist},{pq}"
if opq is None
else f"{opq},IVF{nlist},{pq}"
),
search_params={
"nprobe": nprobe,
},
benchs/bench_fw/utils.py
+11 -7
View File
@@ -91,14 +91,18 @@ def refine_distances_range(
with ThreadPool(32) as pool:
R = pool.map(
lambda i: (
np.sum(np.square(xq[i] - xb[I[lims[i] : lims[i + 1]]]), axis=1)
if metric == faiss.METRIC_L2
else np.tensordot(
xq[i], xb[I[lims[i] : lims[i + 1]]], axes=(0, 1)
(
np.sum(
np.square(xq[i] - xb[I[lims[i] : lims[i + 1]]]), axis=1
)
if metric == faiss.METRIC_L2
else np.tensordot(
xq[i], xb[I[lims[i] : lims[i + 1]]], axes=(0, 1)
)
)
)
if lims[i + 1] > lims[i]
else [],
if lims[i + 1] > lims[i]
else []
),
range(len(lims) - 1),
)
return np.hstack(R)
benchs/bench_fw_codecs.py
+123 -79
View File
@@ -9,97 +9,130 @@ import os
from faiss.benchs.bench_fw.benchmark import Benchmark
from faiss.benchs.bench_fw.benchmark_io import BenchmarkIO
from faiss.benchs.bench_fw.descriptors import DatasetDescriptor, IndexDescriptorClassic
from faiss.benchs.bench_fw.descriptors import (
DatasetDescriptor,
IndexDescriptorClassic,
)
logging.basicConfig(level=logging.INFO)
def factory_factory(d):
return [
("SQ4", None, 256 * (2 ** 10), None),
("SQ8", None, 256 * (2 ** 10), None),
("SQfp16", None, 256 * (2 ** 10), None),
("ITQ64,LSH", None, 256 * (2 ** 10), None),
("Pad128,ITQ128,LSH", None, 256 * (2 ** 10), None),
("Pad256,ITQ256,LSH", None, 256 * (2 ** 10), None),
] + [
(f"OPQ32_128,Residual2x14,PQ32x{b}", None, 256 * (2 ** 14), None)
for b in range(8, 16, 2)
] + [
(f"PCAR{2 ** d_out},SQ{b}", None, 256 * (2 ** 10), None)
for d_out in range(6, 11)
if 2 ** d_out <= d
for b in [4, 8]
] + [
(f"OPQ{M}_{M * dim},PQ{M}x{b}", None, 256 * (2 ** b), None)
for M in [8, 12, 16, 32, 64, 128]
for dim in [2, 4, 6, 8, 12, 16]
if M * dim <= d
for b in range(8, 16, 2)
] + [
(f"RQ{cs // b}x{b}", [{"max_beam_size": 32}], 256 * (2 ** b), {"max_beam_size": bs, "use_beam_LUT": bl})
for cs in [64, 128, 256, 512]
for b in [6, 8, 10, 12]
for bs in [1, 2, 4, 8, 16, 32]
for bl in [0, 1]
if cs // b > 1
if cs // b < 65
if cs < d * 8 * 2
] + [
(f"LSQ{cs // b}x{b}", [{"encode_ils_iters": 16}], 256 * (2 ** b), {"encode_ils_iters": eii, "lsq_gpu": lg})
for cs in [64, 128, 256, 512]
for b in [6, 8, 10, 12]
for eii in [2, 4, 8, 16]
for lg in [0, 1]
if cs // b > 1
if cs // b < 65
if cs < d * 8 * 2
] + [
(f"PRQ{sub}x{cs // sub // b}x{b}", [{"max_beam_size": 32}], 256 * (2 ** b), {"max_beam_size": bs, "use_beam_LUT": bl})
for sub in [2, 3, 4, 8, 16, 32]
for cs in [64, 96, 128, 192, 256, 384, 512, 768, 1024, 2048]
for b in [6, 8, 10, 12]
for bs in [1, 2, 4, 8, 16, 32]
for bl in [0, 1]
if cs // sub // b > 1
if cs // sub // b < 65
if cs < d * 8 * 2
if d % sub == 0
] + [
(f"PLSQ{sub}x{cs // sub // b}x{b}", [{"encode_ils_iters": 16}], 256 * (2 ** b), {"encode_ils_iters": eii, "lsq_gpu": lg})
for sub in [2, 3, 4, 8, 16, 32]
for cs in [64, 128, 256, 512, 1024, 2048]
for b in [6, 8, 10, 12]
for eii in [2, 4, 8, 16]
for lg in [0, 1]
if cs // sub // b > 1
if cs // sub // b < 65
if cs < d * 8 * 2
if d % sub == 0
]
return (
[
("SQ4", None, 256 * (2**10), None),
("SQ8", None, 256 * (2**10), None),
("SQfp16", None, 256 * (2**10), None),
("ITQ64,LSH", None, 256 * (2**10), None),
("Pad128,ITQ128,LSH", None, 256 * (2**10), None),
("Pad256,ITQ256,LSH", None, 256 * (2**10), None),
]
+ [
(f"OPQ32_128,Residual2x14,PQ32x{b}", None, 256 * (2**14), None)
for b in range(8, 16, 2)
]
+ [
(f"PCAR{2 ** d_out},SQ{b}", None, 256 * (2**10), None)
for d_out in range(6, 11)
if 2**d_out <= d
for b in [4, 8]
]
+ [
(f"OPQ{M}_{M * dim},PQ{M}x{b}", None, 256 * (2**b), None)
for M in [8, 12, 16, 32, 64, 128]
for dim in [2, 4, 6, 8, 12, 16]
if M * dim <= d
for b in range(8, 16, 2)
]
+ [
(
f"RQ{cs // b}x{b}",
[{"max_beam_size": 32}],
256 * (2**b),
{"max_beam_size": bs, "use_beam_LUT": bl},
)
for cs in [64, 128, 256, 512]
for b in [6, 8, 10, 12]
for bs in [1, 2, 4, 8, 16, 32]
for bl in [0, 1]
if cs // b > 1
if cs // b < 65
if cs < d * 8 * 2
]
+ [
(
f"LSQ{cs // b}x{b}",
[{"encode_ils_iters": 16}],
256 * (2**b),
{"encode_ils_iters": eii, "lsq_gpu": lg},
)
for cs in [64, 128, 256, 512]
for b in [6, 8, 10, 12]
for eii in [2, 4, 8, 16]
for lg in [0, 1]
if cs // b > 1
if cs // b < 65
if cs < d * 8 * 2
]
+ [
(
f"PRQ{sub}x{cs // sub // b}x{b}",
[{"max_beam_size": 32}],
256 * (2**b),
{"max_beam_size": bs, "use_beam_LUT": bl},
)
for sub in [2, 3, 4, 8, 16, 32]
for cs in [64, 96, 128, 192, 256, 384, 512, 768, 1024, 2048]
for b in [6, 8, 10, 12]
for bs in [1, 2, 4, 8, 16, 32]
for bl in [0, 1]
if cs // sub // b > 1
if cs // sub // b < 65
if cs < d * 8 * 2
if d % sub == 0
]
+ [
(
f"PLSQ{sub}x{cs // sub // b}x{b}",
[{"encode_ils_iters": 16}],
256 * (2**b),
{"encode_ils_iters": eii, "lsq_gpu": lg},
)
for sub in [2, 3, 4, 8, 16, 32]
for cs in [64, 128, 256, 512, 1024, 2048]
for b in [6, 8, 10, 12]
for eii in [2, 4, 8, 16]
for lg in [0, 1]
if cs // sub // b > 1
if cs // sub // b < 65
if cs < d * 8 * 2
if d % sub == 0
]
)
def run_local(rp):
bio, d, tablename, distance_metric = rp
if tablename == "contriever":
training_vectors=DatasetDescriptor(
tablename="training_set.npy"
)
database_vectors=DatasetDescriptor(
training_vectors = DatasetDescriptor(tablename="training_set.npy")
database_vectors = DatasetDescriptor(
tablename="database1M.npy",
)
query_vectors=DatasetDescriptor(
query_vectors = DatasetDescriptor(
tablename="queries.npy",
)
else:
training_vectors=DatasetDescriptor(
namespace="std_t", tablename=tablename,
training_vectors = DatasetDescriptor(
namespace="std_t",
tablename=tablename,
)
database_vectors=DatasetDescriptor(
namespace="std_d", tablename=tablename,
database_vectors = DatasetDescriptor(
namespace="std_d",
tablename=tablename,
)
query_vectors=DatasetDescriptor(
namespace="std_q", tablename=tablename,
query_vectors = DatasetDescriptor(
namespace="std_q",
tablename=tablename,
)
benchmark = Benchmark(
@@ -114,22 +147,33 @@ def run_local(rp):
training_size=training_size,
search_params=search_params,
)
for factory, construction_params, training_size, search_params in factory_factory(d)
for factory, construction_params, training_size, search_params in factory_factory(
d
)
],
k=1,
distance_metric=distance_metric,
)
benchmark.set_io(bio)
benchmark.benchmark(result_file="result.json", train=True, reconstruct=False, knn=False, range=False)
benchmark.benchmark(
result_file="result.json",
train=True,
reconstruct=False,
knn=False,
range=False,
)
def run(bio, d, tablename, distance_metric):
bio.launch_jobs(run_local, [(bio, d, tablename, distance_metric)], local=True)
bio.launch_jobs(
run_local, [(bio, d, tablename, distance_metric)], local=True
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('experiment')
parser.add_argument('path')
parser.add_argument("experiment")
parser.add_argument("path")
args = parser.parse_args()
assert os.path.exists(args.path)
path = os.path.join(args.path, args.experiment)
benchs/bench_fw_ivf.py
+49 -25
View File
@@ -26,9 +26,7 @@ def sift1M(bio):
database_vectors=DatasetDescriptor(
namespace="std_d", tablename="sift1M"
),
query_vectors=DatasetDescriptor(
namespace="std_q", tablename="sift1M"
),
query_vectors=DatasetDescriptor(namespace="std_q", tablename="sift1M"),
index_descs=[
IndexDescriptorClassic(
factory=f"IVF{2 ** nlist},Flat",
@@ -39,7 +37,14 @@ def sift1M(bio):
distance_metric="L2",
)
benchmark.io = bio
benchmark.benchmark(result_file="result.json", local=True, train=True, reconstruct=False, knn=True, range=False)
benchmark.benchmark(
result_file="result.json",
local=True,
train=True,
reconstruct=False,
knn=True,
range=False,
)
def bigann(bio):
@@ -58,41 +63,52 @@ def bigann(bio):
index_descs=[
IndexDescriptorClassic(
factory=f"IVF{2 ** nlist},Flat",
) for nlist in range(11, 19)
] + [
)
for nlist in range(11, 19)
]
+ [
IndexDescriptorClassic(
factory=f"IVF{2 ** nlist}_HNSW32,Flat",
construction_params=[None, {"efConstruction": 200, "efSearch": 40}],
) for nlist in range(11, 19)
construction_params=[
None,
{"efConstruction": 200, "efSearch": 40},
],
)
for nlist in range(11, 19)
],
k=1,
distance_metric="L2",
)
benchmark.set_io(bio)
benchmark.benchmark(f"result{scale}.json", local=False, train=True, reconstruct=False, knn=True, range=False)
benchmark.benchmark(
f"result{scale}.json",
local=False,
train=True,
reconstruct=False,
knn=True,
range=False,
)
def ssnpp(bio):
benchmark = Benchmark(
num_threads=32,
training_vectors=DatasetDescriptor(
tablename="ssnpp_training_5M.npy"
),
database_vectors=DatasetDescriptor(
tablename="ssnpp_database_5M.npy"
),
query_vectors=DatasetDescriptor(
tablename="ssnpp_queries_10K.npy"
),
training_vectors=DatasetDescriptor(tablename="ssnpp_training_5M.npy"),
database_vectors=DatasetDescriptor(tablename="ssnpp_database_5M.npy"),
query_vectors=DatasetDescriptor(tablename="ssnpp_queries_10K.npy"),
index_descs=[
IndexDescriptorClassic(
factory=f"IVF{2 ** nlist},PQ256x4fs,Refine(SQfp16)",
) for nlist in range(9, 16)
] + [
)
for nlist in range(9, 16)
]
+ [
IndexDescriptorClassic(
factory=f"IVF{2 ** nlist},Flat",
) for nlist in range(9, 16)
] + [
)
for nlist in range(9, 16)
]
+ [
IndexDescriptorClassic(
factory=f"PQ256x4fs,Refine(SQfp16)",
),
@@ -104,12 +120,20 @@ def ssnpp(bio):
distance_metric="L2",
)
benchmark.set_io(bio)
benchmark.benchmark("result.json", local=False, train=True, reconstruct=False, knn=True, range=False)
benchmark.benchmark(
"result.json",
local=False,
train=True,
reconstruct=False,
knn=True,
range=False,
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('experiment')
parser.add_argument('path')
parser.add_argument("experiment")
parser.add_argument("path")
args = parser.parse_args()
assert os.path.exists(args.path)
path = os.path.join(args.path, args.experiment)
benchs/bench_fw_range.py
+31 -21
View File
@@ -9,7 +9,10 @@ import os
from faiss.benchs.bench_fw.benchmark import Benchmark
from faiss.benchs.bench_fw.benchmark_io import BenchmarkIO
from faiss.benchs.bench_fw.descriptors import DatasetDescriptor, IndexDescriptorClassic
from faiss.benchs.bench_fw.descriptors import (
DatasetDescriptor,
IndexDescriptorClassic,
)
logging.basicConfig(level=logging.INFO)
@@ -34,25 +37,25 @@ def ssnpp(bio):
[0.15, 0.923],
[0.2, 0.887],
[0.25, 0.801],
[0.3, 0.729],
[0.35, 0.651],
[0.4, 0.55],
[0.45, 0.459],
[0.5, 0.372],
[0.55, 0.283],
[0.6, 0.189],
[0.65, 0.143],
[0.7, 0.106],
[0.75, 0.116],
[0.8, 0.088],
[0.3, 0.729],
[0.35, 0.651],
[0.4, 0.55],
[0.45, 0.459],
[0.5, 0.372],
[0.55, 0.283],
[0.6, 0.189],
[0.65, 0.143],
[0.7, 0.106],
[0.75, 0.116],
[0.8, 0.088],
[0.85, 0.064],
[0.9, 0.05],
[0.95, 0.04],
[1.0, 0.028],
[1.05, 0.02],
[0.9, 0.05],
[0.95, 0.04],
[1.0, 0.028],
[1.05, 0.02],
[1.1, 0.013],
[1.15, 0.007],
[1.2, 0.004],
[1.15, 0.007],
[1.2, 0.004],
[1.3, 0],
]
},
@@ -66,13 +69,20 @@ def ssnpp(bio):
range_ref_index_desc="Flat",
)
benchmark.set_io(bio)
benchmark.benchmark("result.json", local=False, train=True, reconstruct=False, knn=False, range=True)
benchmark.benchmark(
"result.json",
local=False,
train=True,
reconstruct=False,
knn=False,
range=True,
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('experiment')
parser.add_argument('path')
parser.add_argument("experiment")
parser.add_argument("path")
args = parser.parse_args()
assert os.path.exists(args.path)
path = os.path.join(args.path, args.experiment)
benchs/bench_fwht.py
+12 -12
View File
@@ -18,17 +18,17 @@ n = 10000
nq = 200
n_trials = 15
print(f'n={n} vectors, nq={nq} queries, {n_trials} speed trials')
print(f"n={n} vectors, nq={nq} queries, {n_trials} speed trials")
# --- Speed benchmark ---
print('\n=== Speed ===')
print("\n=== Speed ===")
print(f'{"dim":>6s} {"HR (ms)":>10s} {"RRM (ms)":>10s} {"speedup":>8s}')
print('-' * 42)
print("-" * 42)
dims = [64, 128, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096, 6144, 8192]
for d in dims:
np.random.seed(42)
x = np.random.randn(n, d).astype('float32')
x = np.random.randn(n, d).astype("float32")
hr = faiss.HadamardRotation(d, 42)
hr.apply(x[:100]) # warmup
@@ -52,17 +52,17 @@ for d in dims:
rr_ms = sorted(times_rr)[1]
speedup = rr_ms / hr_ms
print(f'{d:6d} {hr_ms:10.2f} {rr_ms:10.2f} {speedup:7.1f}x')
print(f"{d:6d} {hr_ms:10.2f} {rr_ms:10.2f} {speedup:7.1f}x")
# --- Recall benchmark with IVF (where rotation actually matters) ---
print(f'\n=== Recall@1 with IVF (L2, n={n}, nq={nq}, k=1) ===')
print(f"\n=== Recall@1 with IVF (L2, n={n}, nq={nq}, k=1) ===")
print(f'{"dim":>6s} {"HR R@1":>8s} {"RRM R@1":>8s} {"none R@1":>8s}')
print('-' * 40)
print("-" * 40)
for d in [64, 128, 256, 768, 1024, 2048, 4096]:
np.random.seed(42)
xb = np.random.randn(n, d).astype('float32')
xq = np.random.randn(nq, d).astype('float32')
xb = np.random.randn(n, d).astype("float32")
xq = np.random.randn(nq, d).astype("float32")
# Ground truth (brute force, no transform)
gt_index = faiss.IndexFlatL2(d)
@@ -72,9 +72,9 @@ for d in [64, 128, 256, 768, 1024, 2048, 4096]:
nlist = 64
recalls = {}
for label, factory_str in [
('HR', f'HR,IVF{nlist},Flat'),
('RRM', f'RR,IVF{nlist},Flat'),
('none', f'IVF{nlist},Flat'),
("HR", f"HR,IVF{nlist},Flat"),
("RRM", f"RR,IVF{nlist},Flat"),
("none", f"IVF{nlist},Flat"),
]:
index = faiss.index_factory(d, factory_str)
index.train(xb)
benchs/bench_gpu_1bn.py
+132 -113
View File
@@ -21,7 +21,8 @@ from datasets import ivecs_read
def usage():
print("""
print(
"""
Usage: bench_gpu_1bn.py dataset indextype [options]
@@ -60,7 +61,9 @@ indextype: any index type supported by index_factory that runs on GPU.
%d will be replaced with the nprobe
-oD xx%d.npy output the search result distances to this file
""", file=sys.stderr)
""",
file=sys.stderr,
)
sys.exit(1)
@@ -90,33 +93,33 @@ args = sys.argv[1:]
while args:
a = args.pop(0)
if a == '-h':
if a == "-h":
usage()
elif a == '-ngpu':
elif a == "-ngpu":
ngpu = int(args.pop(0))
elif a == '-R':
elif a == "-R":
replicas = int(args.pop(0))
elif a == '-noptables':
elif a == "-noptables":
use_precomputed_tables = False
elif a == '-abs':
elif a == "-abs":
add_batch_size = int(args.pop(0))
elif a == '-qbs':
elif a == "-qbs":
query_batch_size = int(args.pop(0))
elif a == '-nnn':
elif a == "-nnn":
nnn = int(args.pop(0))
elif a == '-tempmem':
elif a == "-tempmem":
tempmem = int(args.pop(0))
elif a == '-nocache':
elif a == "-nocache":
use_cache = False
elif a == '-knngraph':
elif a == "-knngraph":
knngraph = True
elif a == '-altadd':
elif a == "-altadd":
altadd = True
elif a == '-float16':
elif a == "-float16":
use_float16 = True
elif a == '-nprobe':
nprobes = [int(x) for x in args.pop(0).split(',')]
elif a == '-max_add':
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
@@ -126,7 +129,7 @@ while args:
print("argument %s unknown" % a, file=sys.stderr)
sys.exit(1)
cacheroot = '/tmp/bench_gpu_1bn'
cacheroot = "/tmp/bench_gpu_1bn"
if not os.path.isdir(cacheroot):
print("%s does not exist, creating it" % cacheroot)
@@ -141,14 +144,14 @@ if not os.path.isdir(cacheroot):
def mmap_fvecs(fname):
x = np.memmap(fname, dtype='int32', mode='r')
x = np.memmap(fname, dtype="int32", mode="r")
d = x[0]
return x.view('float32').reshape(-1, d + 1)[:, 1:]
return x.view("float32").reshape(-1, d + 1)[:, 1:]
def mmap_bvecs(fname):
x = np.memmap(fname, dtype='uint8', mode='r')
d = x[:4].view('int32')[0]
x = np.memmap(fname, dtype="uint8", mode="r")
d = x[:4].view("int32")[0]
return x.reshape(-1, d + 4)[:, 4:]
@@ -158,7 +161,7 @@ def rate_limited_imap(f, l):
pool = ThreadPool(1)
res = None
for i in l:
res_next = pool.apply_async(f, (i, ))
res_next = pool.apply_async(f, (i,))
if res:
yield res.get()
res = res_next
@@ -176,16 +179,15 @@ class IdentPreproc:
def sanitize(x):
""" convert array to a c-contiguous float array """
return np.ascontiguousarray(x.astype('float32'))
"""convert array to a c-contiguous float array"""
return np.ascontiguousarray(x.astype("float32"))
def dataset_iterator(x, preproc, bs):
""" iterate over the lines of x in blocks of size bs"""
"""iterate over the lines of x in blocks of size bs"""
nb = x.shape[0]
block_ranges = [(i0, min(nb, i0 + bs))
for i0 in range(0, nb, bs)]
block_ranges = [(i0, min(nb, i0 + bs)) for i0 in range(0, nb, bs)]
def prepare_block(i01):
i0, i1 = i01
@@ -196,14 +198,17 @@ def dataset_iterator(x, preproc, bs):
def eval_intersection_measure(gt_I, I):
""" measure intersection measure (used for knngraph)"""
"""measure intersection measure (used for knngraph)"""
inter = 0
rank = I.shape[1]
assert gt_I.shape[1] >= rank
for q in range(nq_gt):
inter += faiss.ranklist_intersection_size(
rank, faiss.swig_ptr(gt_I[q, :]),
rank, faiss.swig_ptr(I[q, :].astype('int64')))
rank,
faiss.swig_ptr(gt_I[q, :]),
rank,
faiss.swig_ptr(I[q, :].astype("int64")),
)
return inter / float(rank * nq_gt)
@@ -213,28 +218,28 @@ def eval_intersection_measure(gt_I, I):
print("Preparing dataset", dbname)
if dbname.startswith('SIFT'):
if dbname.startswith("SIFT"):
# SIFT1M to SIFT1000M
dbsize = int(dbname[4:-1])
xb = mmap_bvecs('bigann/bigann_base.bvecs')
xq = mmap_bvecs('bigann/bigann_query.bvecs')
xt = mmap_bvecs('bigann/bigann_learn.bvecs')
xb = mmap_bvecs("bigann/bigann_base.bvecs")
xq = mmap_bvecs("bigann/bigann_query.bvecs")
xt = mmap_bvecs("bigann/bigann_learn.bvecs")
# trim xb to correct size
xb = xb[:dbsize * 1000 * 1000]
xb = xb[: dbsize * 1000 * 1000]
gt_I = ivecs_read('bigann/gnd/idx_%dM.ivecs' % dbsize)
gt_I = ivecs_read("bigann/gnd/idx_%dM.ivecs" % dbsize)
elif dbname == 'Deep1B':
xb = mmap_fvecs('deep1b/base.fvecs')
xq = mmap_fvecs('deep1b/deep1B_queries.fvecs')
xt = mmap_fvecs('deep1b/learn.fvecs')
elif dbname == "Deep1B":
xb = mmap_fvecs("deep1b/base.fvecs")
xq = mmap_fvecs("deep1b/deep1B_queries.fvecs")
xt = mmap_fvecs("deep1b/learn.fvecs")
# deep1B's train is outrageously big
xt = xt[:10 * 1000 * 1000]
gt_I = ivecs_read('deep1b/deep1B_groundtruth.ivecs')
xt = xt[: 10 * 1000 * 1000]
gt_I = ivecs_read("deep1b/deep1B_groundtruth.ivecs")
else:
print('unknown dataset', dbname, file=sys.stderr)
print("unknown dataset", dbname, file=sys.stderr)
sys.exit(1)
@@ -250,9 +255,10 @@ if knngraph:
gt_I = None
print("sizes: B %s Q %s T %s gt %s" % (
xb.shape, xq.shape, xt.shape,
gt_I.shape if gt_I is not None else None))
print(
"sizes: B %s Q %s T %s gt %s"
% (xb.shape, xq.shape, xt.shape, gt_I.shape if gt_I is not None else None)
)
#################################################################
@@ -263,13 +269,13 @@ print("sizes: B %s Q %s T %s gt %s" % (
#################################################################
pat = re.compile('(OPQ[0-9]+(_[0-9]+)?,|PCAR[0-9]+,)?' +
'(IVF[0-9]+),' +
'(PQ[0-9]+|Flat)')
pat = re.compile(
"(OPQ[0-9]+(_[0-9]+)?,|PCAR[0-9]+,)?" + "(IVF[0-9]+)," + "(PQ[0-9]+|Flat)"
)
matchobject = pat.match(index_key)
assert matchobject, 'could not parse ' + index_key
assert matchobject, "could not parse " + index_key
mog = matchobject.groups()
@@ -279,26 +285,41 @@ pqflat_str = mog[3]
ncent = int(ivf_str[3:])
prefix = ''
prefix = ""
if knngraph:
gt_cachefile = '%s/BK_gt_%s.npy' % (cacheroot, dbname)
prefix = 'BK_'
gt_cachefile = "%s/BK_gt_%s.npy" % (cacheroot, dbname)
prefix = "BK_"
# files must be kept distinct because the training set is not the
# same for the knngraph
if preproc_str:
preproc_cachefile = '%s/%spreproc_%s_%s.vectrans' % (
cacheroot, prefix, dbname, preproc_str[:-1])
preproc_cachefile = "%s/%spreproc_%s_%s.vectrans" % (
cacheroot,
prefix,
dbname,
preproc_str[:-1],
)
else:
preproc_cachefile = None
preproc_str = ''
preproc_str = ""
cent_cachefile = '%s/%scent_%s_%s%s.npy' % (
cacheroot, prefix, dbname, preproc_str, ivf_str)
cent_cachefile = "%s/%scent_%s_%s%s.npy" % (
cacheroot,
prefix,
dbname,
preproc_str,
ivf_str,
)
index_cachefile = '%s/%s%s_%s%s,%s.index' % (
cacheroot, prefix, dbname, preproc_str, ivf_str, pqflat_str)
index_cachefile = "%s/%s%s_%s%s,%s.index" % (
cacheroot,
prefix,
dbname,
preproc_str,
ivf_str,
pqflat_str,
)
if not use_cache:
@@ -328,7 +349,7 @@ for _ in range(ngpu):
def make_vres_vdev(i0=0, i1=-1):
" return vectors of device ids and resources useful for gpu_multiple"
"return vectors of device ids and resources useful for gpu_multiple"
vres = faiss.GpuResourcesVector()
vdev = faiss.IntVector()
if i1 == -1:
@@ -348,33 +369,31 @@ def compute_GT():
print("compute GT")
t0 = time.time()
gt_I = np.zeros((nq_gt, gt_sl), dtype='int64')
gt_D = np.zeros((nq_gt, gt_sl), dtype='float32')
gt_I = np.zeros((nq_gt, gt_sl), dtype="int64")
gt_D = np.zeros((nq_gt, gt_sl), dtype="float32")
heaps = faiss.float_maxheap_array_t()
heaps.k = gt_sl
heaps.nh = nq_gt
heaps.val = faiss.swig_ptr(gt_D)
heaps.ids = faiss.swig_ptr(gt_I)
heaps.heapify()
bs = 10 ** 5
bs = 10**5
n, d = xb.shape
xqs = sanitize(xq[:nq_gt])
db_gt = faiss.IndexFlatL2(d)
vres, vdev = make_vres_vdev()
db_gt_gpu = faiss.index_cpu_to_gpu_multiple(
vres, vdev, db_gt)
db_gt_gpu = faiss.index_cpu_to_gpu_multiple(vres, vdev, db_gt)
# compute ground-truth by blocks of bs, and add to heaps
for i0, xsl in dataset_iterator(xb, IdentPreproc(d), bs):
db_gt_gpu.add(xsl)
D, I = db_gt_gpu.search(xqs, gt_sl)
I += i0
heaps.addn_with_ids(
gt_sl, faiss.swig_ptr(D), faiss.swig_ptr(I), gt_sl)
heaps.addn_with_ids(gt_sl, faiss.swig_ptr(D), faiss.swig_ptr(I), gt_sl)
db_gt_gpu.reset()
print("\r %d/%d, %.3f s" % (i0, n, time.time() - t0), end=' ')
print("\r %d/%d, %.3f s" % (i0, n, time.time() - t0), end=" ")
print()
heaps.reorder()
@@ -402,12 +421,12 @@ def train_preprocessor():
print("train preproc", preproc_str)
d = xt.shape[1]
t0 = time.time()
if preproc_str.startswith('OPQ'):
fi = preproc_str[3:-1].split('_')
if preproc_str.startswith("OPQ"):
fi = preproc_str[3:-1].split("_")
m = int(fi[0])
dout = int(fi[1]) if len(fi) == 2 else d
preproc = faiss.OPQMatrix(d, m, dout)
elif preproc_str.startswith('PCAR'):
elif preproc_str.startswith("PCAR"):
dout = int(preproc_str[4:-1])
preproc = faiss.PCAMatrix(d, dout, 0, True)
else:
@@ -445,15 +464,13 @@ def train_coarse_quantizer(x, k, preproc):
# clus.niter = 2
clus.max_points_per_centroid = 10000000
print("apply preproc on shape", x.shape, 'k=', k)
print("apply preproc on shape", x.shape, "k=", k)
t0 = time.time()
x = preproc.apply_py(sanitize(x))
print(" preproc %.3f s output shape %s" % (
time.time() - t0, x.shape))
print(" preproc %.3f s output shape %s" % (time.time() - t0, x.shape))
vres, vdev = make_vres_vdev()
index = faiss.index_cpu_to_gpu_multiple(
vres, vdev, faiss.IndexFlatL2(d))
index = faiss.index_cpu_to_gpu_multiple(vres, vdev, faiss.IndexFlatL2(d))
clus.train(x, index)
centroids = faiss.vector_float_to_array(clus.centroids)
@@ -491,10 +508,11 @@ def prepare_trained_index(preproc):
coarse_quantizer = prepare_coarse_quantizer(preproc)
d = preproc.d_out
if pqflat_str == 'Flat':
if pqflat_str == "Flat":
print("making an IVFFlat index")
idx_model = faiss.IndexIVFFlat(coarse_quantizer, d, ncent,
faiss.METRIC_L2)
idx_model = faiss.IndexIVFFlat(
coarse_quantizer, d, ncent, faiss.METRIC_L2
)
else:
m = int(pqflat_str[2:])
assert m < 56 or use_float16, "PQ%d will work only with -float16" % m
@@ -516,7 +534,7 @@ def prepare_trained_index(preproc):
def compute_populated_index(preproc):
"""Add elements to a sharded index. Return the index and if available
a sharded gpu_index that contains the same data. """
a sharded gpu_index that contains the same data."""
indexall = prepare_trained_index(preproc)
@@ -530,8 +548,7 @@ def compute_populated_index(preproc):
co.shard = True
assert co.shard_type in (0, 1, 2)
vres, vdev = make_vres_vdev()
gpu_index = faiss.index_cpu_to_gpu_multiple(
vres, vdev, indexall, co)
gpu_index = faiss.index_cpu_to_gpu_multiple(vres, vdev, indexall, co)
print("add...")
t0 = time.time()
@@ -550,15 +567,14 @@ def compute_populated_index(preproc):
index_src_gpu.reserveMemory(max_add)
gpu_index.sync_with_shard_indexes()
print('\r%d/%d (%.3f s) ' % (
i0, nb, time.time() - t0), end=' ')
print("\r%d/%d (%.3f s) " % (i0, nb, time.time() - t0), end=" ")
sys.stdout.flush()
print("Add time: %.3f s" % (time.time() - t0))
print("Aggregate indexes to CPU")
t0 = time.time()
if hasattr(gpu_index, 'at'):
if hasattr(gpu_index, "at"):
# it is a sharded index
for i in range(ngpu):
index_src = faiss.index_gpu_to_cpu(gpu_index.at(i))
@@ -591,7 +607,8 @@ def compute_populated_index_2(preproc):
vres, vdev = make_vres_vdev()
coarse_quantizer_gpu = faiss.index_cpu_to_gpu_multiple(
vres, vdev, indexall.quantizer)
vres, vdev, indexall.quantizer
)
def quantize(args):
(i0, xs) = args
@@ -607,16 +624,17 @@ def compute_populated_index_2(preproc):
for i0, xs, assign in stage2:
i1 = i0 + xs.shape[0]
if indexall.__class__ == faiss.IndexIVFPQ:
indexall.add_core_o(i1 - i0, faiss.swig_ptr(xs),
None, None, faiss.swig_ptr(assign))
indexall.add_core_o(
i1 - i0, faiss.swig_ptr(xs), None, None, faiss.swig_ptr(assign)
)
elif indexall.__class__ == faiss.IndexIVFFlat:
indexall.add_core(i1 - i0, faiss.swig_ptr(xs), None,
faiss.swig_ptr(assign))
indexall.add_core(
i1 - i0, faiss.swig_ptr(xs), None, faiss.swig_ptr(assign)
)
else:
assert False
print('\r%d/%d (%.3f s) ' % (
i0, nb, time.time() - t0), end=' ')
print("\r%d/%d (%.3f s) " % (i0, nb, time.time() - t0), end=" ")
sys.stdout.flush()
print("Add time: %.3f s" % (time.time() - t0))
@@ -644,7 +662,7 @@ def get_populated_index(preproc):
co.usePrecomputed = use_precomputed_tables
co.indicesOptions = 0
co.verbose = True
co.shard = True # the replicas will be made "manually"
co.shard = True # the replicas will be made "manually"
t0 = time.time()
print("CPU index contains %d vectors, move to GPU" % indexall.ntotal)
if replicas == 1:
@@ -652,8 +670,7 @@ def get_populated_index(preproc):
if not gpu_index:
print("copying loaded index to GPUs")
vres, vdev = make_vres_vdev()
index = faiss.index_cpu_to_gpu_multiple(
vres, vdev, indexall, co)
index = faiss.index_cpu_to_gpu_multiple(vres, vdev, indexall, co)
else:
index = gpu_index
@@ -671,8 +688,7 @@ def get_populated_index(preproc):
print(" dispatch to GPUs %d:%d" % (gpu0, gpu1))
index1 = faiss.index_cpu_to_gpu_multiple(
vres, vdev, indexall, co)
index1 = faiss.index_cpu_to_gpu_multiple(vres, vdev, indexall, co)
index1.this.disown()
index.addIndex(index1)
index.own_fields = True
@@ -681,7 +697,6 @@ def get_populated_index(preproc):
return index
#################################################################
# Perform search
#################################################################
@@ -697,20 +712,23 @@ def eval_dataset(index, preproc):
sl = query_batch_size
nq = xq.shape[0]
for nprobe in nprobes:
ps.set_index_parameter(index, 'nprobe', nprobe)
ps.set_index_parameter(index, "nprobe", nprobe)
t0 = time.time()
if sl == 0:
D, I = index.search(preproc.apply_py(sanitize(xq)), nnn)
else:
I = np.empty((nq, nnn), dtype='int32')
D = np.empty((nq, nnn), dtype='float32')
I = np.empty((nq, nnn), dtype="int32")
D = np.empty((nq, nnn), dtype="float32")
inter_res = ''
inter_res = ""
for i0, xs in dataset_iterator(xq, preproc, sl):
print('\r%d/%d (%.3f s%s) ' % (
i0, nq, time.time() - t0, inter_res), end=' ')
print(
"\r%d/%d (%.3f s%s) "
% (i0, nq, time.time() - t0, inter_res),
end=" ",
)
sys.stdout.flush()
i1 = i0 + xs.shape[0]
@@ -720,24 +738,25 @@ def eval_dataset(index, preproc):
D[i0:i1] = Di
if knngraph and not inter_res and i1 >= nq_gt:
ires = eval_intersection_measure(
gt_I[:, :nnn], I[:nq_gt])
inter_res = ', %.4f' % ires
ires = eval_intersection_measure(gt_I[:, :nnn], I[:nq_gt])
inter_res = ", %.4f" % ires
t1 = time.time()
if knngraph:
ires = eval_intersection_measure(gt_I[:, :nnn], I[:nq_gt])
print(" probe=%-3d: %.3f s rank-%d intersection results: %.4f" % (
nprobe, t1 - t0, nnn, ires))
print(
" probe=%-3d: %.3f s rank-%d intersection results: %.4f"
% (nprobe, t1 - t0, nnn, ires)
)
else:
print(" probe=%-3d: %.3f s" % (nprobe, t1 - t0), end=' ')
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
nok = (I[:, :rank] == gtc).sum()
print("1-R@%d: %.4f" % (rank, nok / float(nq)), end=' ')
print("1-R@%d: %.4f" % (rank, nok / float(nq)), end=" ")
print()
if I_fname:
I_fname_i = I_fname % nprobe
benchs/bench_gpu_sift1m.py
+4 -1
View File
@@ -84,4 +84,7 @@ for lnprobe in range(10):
index.nprobe = nprobe
t, r = evaluate(index, xq, gt, 100)
print("nprobe=%4d %.3f ms recalls= %.4f %.4f %.4f" % (nprobe, t, r[1], r[10], r[100]))
print(
"nprobe=%4d %.3f ms recalls= %.4f %.4f %.4f"
% (nprobe, t, r[1], r[10], r[100])
)
benchs/bench_hamming_knn.py
+8 -4
View File
@@ -13,9 +13,9 @@ if __name__ == "__main__":
for d in 4, 8, 16, 32, 64, 128, 256, 512, 1024:
nq = 10000
nb = 30000
print('Bits per vector = 8 *', d)
xq = faiss.randint((nq, d // 4), seed=1234, vmax=256**4).view('uint8')
xb = faiss.randint((nb, d // 4), seed=1234, vmax=256**4).view('uint8')
print("Bits per vector = 8 *", d)
xq = faiss.randint((nq, d // 4), seed=1234, vmax=256**4).view("uint8")
xb = faiss.randint((nb, d // 4), seed=1234, vmax=256**4).view("uint8")
for variant in "hc", "mc":
print(f"{variant=:}", end="\t")
for k in 1, 4, 16, 64, 256:
@@ -25,5 +25,9 @@ if __name__ == "__main__":
D, I = faiss.knn_hamming(xq, xb, k, variant=variant)
t1 = time.time()
times.append(t1 - t0)
print(f'| {k=:} t={np.mean(times):.3f} s ± {np.std(times):.3f} ', flush=True, end="")
print(
f"| {k=:} t={np.mean(times):.3f} s ± {np.std(times):.3f} ",
flush=True,
end="",
)
print()
benchs/bench_hnsw.py
+1 -1
View File
@@ -221,7 +221,7 @@ if "hnsw_locks" in todo:
t2 = None
for i in range(0, len(xb), batch_size):
t1 = time.time()
index.add(xb[i: i + batch_size])
index.add(xb[i : i + batch_size])
t2 = time.time()
if i > 2 and t2 - t0 > 2:
break
benchs/bench_hnsw_flat_panorama.py
+5 -3
View File
@@ -104,7 +104,9 @@ def benchmark_dataset(ds, dataset_name, k=10, nlevels=8, M=32):
plot_data = []
# HNSW Flat (baseline)
eval_and_plot(f"HNSW{M},Flat", ds, k=k, nlevels=nlevels, plot_data=plot_data)
eval_and_plot(
f"HNSW{M},Flat", ds, k=k, nlevels=nlevels, plot_data=plot_data
)
# HNSW Flat Panorama (with PCA to concentrate energy)
eval_and_plot(
@@ -153,6 +155,6 @@ if __name__ == "__main__":
print(f"{'='*60}")
benchmark_dataset(ds, name, k=k, nlevels=nlevels, M=M)
print("\n" + "="*60)
print("\n" + "=" * 60)
print("All benchmarks completed!")
print("="*60)
print("=" * 60)
benchs/bench_hnsw_prune_headroom.py
+2 -1
View File
@@ -108,7 +108,8 @@ def run_benchmark(
for headroom in headroom_values:
for rep in range(reps):
index, build_time = build_hnsw_index(
d, m, xb, ef_construction, headroom)
d, m, xb, ef_construction, headroom
)
results["build_times"][headroom] = build_time
faiss.cvar.hnsw_stats.reset()
benchs/bench_hybrid_cpu_gpu.py
+174 -128
View File
@@ -34,12 +34,13 @@ def search_preassigned(xq, k, index, quantizer, batch_size=0):
nprobe = index.nprobe
if batch_size == 0:
batch_size = n + 1
D = np.empty((n, k), dtype='float32')
I = np.empty((n, k), dtype='int64')
D = np.empty((n, k), dtype="float32")
I = np.empty((n, k), dtype="int64")
for i0 in range(0, n, batch_size):
Dq, Iq = quantizer.search(xq[i0:i0 + batch_size], nprobe)
D[i0:i0 + batch_size], I[i0:i0 + batch_size] = \
index.search_preassigned(xq[i0:i0 + batch_size], k, Iq, Dq)
Dq, Iq = quantizer.search(xq[i0 : i0 + batch_size], nprobe)
D[i0 : i0 + batch_size], I[i0 : i0 + batch_size] = (
index.search_preassigned(xq[i0 : i0 + batch_size], k, Iq, Dq)
)
return D, I
@@ -61,16 +62,17 @@ def tiled_search_preassigned(xq, k, index, quantizer, batch_size=32768):
i1 = min(i0 + batch_size, n)
return quantizer.search(xq[i0:i1], nprobe)
D = np.empty((n, k), dtype='float32')
I = np.empty((n, k), dtype='int64')
D = np.empty((n, k), dtype="float32")
I = np.empty((n, k), dtype="int64")
qq = coarse_quant(0)
for i0 in range(0, n, batch_size):
i1 = min(i0 + batch_size, n)
qq_next = qq_pool.apply_async(coarse_quant, (i0 + batch_size, ))
qq_next = qq_pool.apply_async(coarse_quant, (i0 + batch_size,))
Dq, Iq = qq
index.search_preassigned(
xq[i0:i1], k, Iq=Iq, Dq=Dq, I=I[i0:i1], D=D[i0:i1])
xq[i0:i1], k, Iq=Iq, Dq=Dq, I=I[i0:i1], D=D[i0:i1]
)
qq = qq_next.get()
qq_pool.close()
@@ -81,6 +83,7 @@ def tiled_search_preassigned(xq, k, index, quantizer, batch_size=32768):
# IVF index objects with a separate coarse quantizer
#################################################################
class SeparateCoarseQuantizationIndex:
"""
Separately manage the coarse quantizer and the IVF index.
@@ -120,10 +123,12 @@ class SeparateCoarseQuantizationIndex:
return self.index_ivf.search_preassigned(xq, k, Iq, Dq)
if self.seq_tiling:
return search_preassigned(
xq, k, self.index_ivf, self.quantizer, self.bs)
xq, k, self.index_ivf, self.quantizer, self.bs
)
else:
return tiled_search_preassigned(
xq, k, self.index_ivf, self.quantizer, self.bs)
xq, k, self.index_ivf, self.quantizer, self.bs
)
class ShardedGPUIndex:
@@ -131,6 +136,7 @@ class ShardedGPUIndex:
Multiple GPU indexes, each on its GPU, with a common coarse quantizer.
The Python version of IndexShardsIVF
"""
def __init__(self, quantizer, index, bs=-1, seq_tiling=False):
self.quantizer = quantizer
self.cpu_index = index
@@ -159,8 +165,8 @@ class ShardedGPUIndex:
sub_index_0 = faiss.downcast_index(index.at(0))
nprobe = sub_index_0.nprobe
Dall = np.empty((ngpu, nq, k), dtype='float32')
Iall = np.empty((ngpu, nq, k), dtype='int64')
Dall = np.empty((ngpu, nq, k), dtype="float32")
Iall = np.empty((ngpu, nq, k), dtype="int64")
bs = self.bs
if bs <= 0:
@@ -169,7 +175,9 @@ class ShardedGPUIndex:
def do_search(rank):
gpu_index = faiss.downcast_index(index.at(rank))
Dall[rank], Iall[rank] = gpu_index.search_preassigned(
xq, k, Iq, Dq)
xq, k, Iq, Dq
)
list(self.pool.map(do_search, range(ngpu)))
else:
qq_pool = self.q_pool
@@ -178,30 +186,32 @@ class ShardedGPUIndex:
def coarse_quant(i0):
if i0 >= nq:
return None
return self.quantizer.search(xq[i0:i0 + bs], nprobe)
return self.quantizer.search(xq[i0 : i0 + bs], nprobe)
def do_search(rank, i0, qq):
gpu_index = faiss.downcast_index(index.at(rank))
Dq, Iq = qq
Dall[rank, i0:i0 + bs], Iall[rank, i0:i0 + bs] = \
gpu_index.search_preassigned(xq[i0:i0 + bs], k, Iq, Dq)
Dall[rank, i0 : i0 + bs], Iall[rank, i0 : i0 + bs] = (
gpu_index.search_preassigned(xq[i0 : i0 + bs], k, Iq, Dq)
)
qq = coarse_quant(0)
for i0 in range(0, nq, bs):
qq_next = qq_pool.apply_async(coarse_quant, (i0 + bs, ))
list(self.pool.map(
lambda rank: do_search(rank, i0, qq),
range(ngpu)
))
qq_next = qq_pool.apply_async(coarse_quant, (i0 + bs,))
list(
self.pool.map(
lambda rank: do_search(rank, i0, qq), range(ngpu)
)
)
qq = qq_next.get()
return faiss.merge_knn_results(Dall, Iall)
def extract_index_ivf(index):
""" extract the IVF sub-index from the index, supporting GpuIndexes
as well """
"""extract the IVF sub-index from the index, supporting GpuIndexes
as well"""
try:
return faiss.extract_index_ivf(index)
except RuntimeError:
@@ -221,7 +231,8 @@ def set_index_parameter(index, name, val):
set_index_parameter(index.index_ivf, name, val)
elif name.startswith("quantizer_"):
set_index_parameter(
index.quantizer, name[name.find("_") + 1:], val)
index.quantizer, name[name.find("_") + 1 :], val
)
else:
raise RuntimeError()
return
@@ -231,7 +242,8 @@ def set_index_parameter(index, name, val):
set_index_parameter(index.cpu_index, name, val)
elif name.startswith("quantizer_"):
set_index_parameter(
index.quantizer, name[name.find("_") + 1:], val)
index.quantizer, name[name.find("_") + 1 :], val
)
else:
raise RuntimeError()
return
@@ -244,20 +256,23 @@ def set_index_parameter(index, name, val):
elif isinstance(index, faiss.IndexShardsIVF):
if name != "nprobe" and name.startswith("quantizer_"):
set_index_parameter(
index.quantizer, name[name.find("_") + 1:], val)
index.quantizer, name[name.find("_") + 1 :], val
)
else:
for i in range(index.count()):
sub_index = index.at(i)
set_index_parameter(sub_index, name, val)
elif (isinstance(index, faiss.IndexShards) or
isinstance(index, faiss.IndexReplicas)):
elif isinstance(index, faiss.IndexShards) or isinstance(
index, faiss.IndexReplicas
):
for i in range(index.count()):
sub_index = index.at(i)
set_index_parameter(sub_index, name, val)
elif name.startswith("quantizer_"):
index_ivf = extract_index_ivf(index)
set_index_parameter(
index_ivf.quantizer, name[name.find("_") + 1:], val)
index_ivf.quantizer, name[name.find("_") + 1 :], val
)
elif name == "efSearch":
index.hnsw.efSearch
index.hnsw.efSearch = int(val)
@@ -280,47 +295,85 @@ def main():
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('dataset options')
aa('--nq', type=int, default=int(10e5),
help="nb queries (queries will be duplicated if below that number)")
aa('--db', default='bigann10M', help='dataset')
group = parser.add_argument_group("dataset options")
aa(
"--nq",
type=int,
default=int(10e5),
help="nb queries (queries will be duplicated if below that number)",
)
aa("--db", default="bigann10M", help="dataset")
group = parser.add_argument_group('index options')
aa('--indexname', default="", help="override index name")
aa('--mmap', default=False, action='store_true', help='mmap index')
aa('--shard_type', default=1, type=int, help="set type of sharding")
aa('--useFloat16', default=False, action='store_true',
help='GPU cloner options')
aa('--useFloat16CoarseQuantizer', default=False, action='store_true',
help='GPU cloner options')
aa('--usePrecomputed', default=False, action='store_true',
help='GPU cloner options')
group = parser.add_argument_group('search options')
aa('--k', type=int, default=100)
aa('--search_type', default="cpu",
group = parser.add_argument_group("index options")
aa("--indexname", default="", help="override index name")
aa("--mmap", default=False, action="store_true", help="mmap index")
aa("--shard_type", default=1, type=int, help="set type of sharding")
aa(
"--useFloat16",
default=False,
action="store_true",
help="GPU cloner options",
)
aa(
"--useFloat16CoarseQuantizer",
default=False,
action="store_true",
help="GPU cloner options",
)
aa(
"--usePrecomputed",
default=False,
action="store_true",
help="GPU cloner options",
)
group = parser.add_argument_group("search options")
aa("--k", type=int, default=100)
aa(
"--search_type",
default="cpu",
choices=[
"cpu", "gpu", "gpu_flat_quantizer",
"cpu_flat_gpu_quantizer", "gpu_tiled", "gpu_ivf_quantizer",
"multi_gpu", "multi_gpu_flat_quantizer",
"multi_gpu_sharded", "multi_gpu_flat_quantizer_sharded",
"multi_gpu_sharded1", "multi_gpu_sharded1_flat",
"cpu",
"gpu",
"gpu_flat_quantizer",
"cpu_flat_gpu_quantizer",
"gpu_tiled",
"gpu_ivf_quantizer",
"multi_gpu",
"multi_gpu_flat_quantizer",
"multi_gpu_sharded",
"multi_gpu_flat_quantizer_sharded",
"multi_gpu_sharded1",
"multi_gpu_sharded1_flat",
"multi_gpu_sharded1_ivf",
"multi_gpu_Csharded1", "multi_gpu_Csharded1_flat",
"multi_gpu_Csharded1",
"multi_gpu_Csharded1_flat",
"multi_gpu_Csharded1_ivf",
],
help="how to search"
help="how to search",
)
aa('--ivf_quant_nlist', type=int, default=1024,
help="nb of inverted lists for IVF quantizer")
aa('--batch_size', type=int, default=-1,
help="batch size for tiled CPU / GPU computation (-1= no tiling)")
aa('--n_autotune', type=int, default=300,
help="max nb of auto-tuning steps")
aa('--nt', type=int, default=-1, help="force number of CPU threads to this")
aa(
"--ivf_quant_nlist",
type=int,
default=1024,
help="nb of inverted lists for IVF quantizer",
)
aa(
"--batch_size",
type=int,
default=-1,
help="batch size for tiled CPU / GPU computation (-1= no tiling)",
)
aa(
"--n_autotune",
type=int,
default=300,
help="max nb of auto-tuning steps",
)
aa("--nt", type=int, default=-1, help="force number of CPU threads to this")
group = parser.add_argument_group('output options')
aa('--quiet', default=False, action="store_true")
aa('--stats', default="", help="pickle to store output stats")
group = parser.add_argument_group("output options")
aa("--quiet", default=False, action="store_true")
aa("--stats", default="", help="pickle to store output stats")
args = parser.parse_args()
print("args:", args)
@@ -342,7 +395,7 @@ def main():
xqx = []
n = 0
while n < args.nq:
xqx.append(xq[:args.nq - n])
xqx.append(xq[: args.nq - n])
n += len(xqx[-1])
print(f"increased nb queries from {len(xq)} to {n}")
xq = np.vstack(xqx)
@@ -365,10 +418,12 @@ def main():
print("prepare index")
op = OperatingPointsWithRanges()
op.add_range(
"nprobe", [
2 ** i for i in range(20)
if 2 ** i < index_ivf.nlist * 0.1 and 2 ** i <= 4096
]
"nprobe",
[
2**i
for i in range(20)
if 2**i < index_ivf.nlist * 0.1 and 2**i <= 4096
],
)
# prepare options for GPU clone
@@ -380,18 +435,12 @@ def main():
co.shard_type = args.shard_type
if args.search_type == "cpu":
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
elif args.search_type == "gpu":
print("move index to 1 GPU")
res = faiss.StandardGpuResources()
index = faiss.index_cpu_to_gpu(res, 0, index, co)
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
op.restrict_range("nprobe", 2049)
elif args.search_type == "gpu_tiled":
print("move index to 1 GPU")
@@ -399,31 +448,28 @@ def main():
quantizer_hnsw = replace_ivf_quantizer(index_ivf, new_quantizer)
res = faiss.StandardGpuResources()
index = faiss.index_cpu_to_gpu(res, 0, index, co)
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
op.restrict_range("nprobe", 2049)
index = SeparateCoarseQuantizationIndex(
quantizer_hnsw, index, bs=args.batch_size)
quantizer_hnsw, index, bs=args.batch_size
)
elif args.search_type == "gpu_ivf_quantizer":
index_ivf = faiss.extract_index_ivf(index)
centroids = index_ivf.quantizer.reconstruct_n()
replace_ivf_quantizer(index_ivf, faiss.IndexFlatL2(index_ivf.d))
res = faiss.StandardGpuResources()
new_quantizer = faiss.index_factory(
index_ivf.d, f"IVF{args.ivf_quant_nlist},Flat")
index_ivf.d, f"IVF{args.ivf_quant_nlist},Flat"
)
new_quantizer.train(centroids)
new_quantizer.add(centroids)
index = SeparateCoarseQuantizationIndex(
faiss.index_cpu_to_gpu(res, 0, new_quantizer, co),
faiss.index_cpu_to_gpu(res, 0, index, co),
bs=args.batch_size, seq_tiling=True
)
op.add_range(
"quantizer_nprobe",
[2 ** i for i in range(9)]
bs=args.batch_size,
seq_tiling=True,
)
op.add_range("quantizer_nprobe", [2**i for i in range(9)])
op.restrict_range("nprobe", 1025)
elif args.search_type == "gpu_flat_quantizer":
index_ivf = faiss.extract_index_ivf(index)
@@ -438,27 +484,29 @@ def main():
res = faiss.StandardGpuResources()
quantizer = faiss.index_cpu_to_gpu(res, 0, quantizer, co)
index = SeparateCoarseQuantizationIndex(
quantizer, index, bs=args.batch_size)
quantizer, index, bs=args.batch_size
)
op.restrict_range("nprobe", 2049)
elif args.search_type in ("multi_gpu", "multi_gpu_sharded"):
print(f"move index to {faiss.get_num_gpus()} GPU")
co.shard = "sharded" in args.search_type
index = faiss.index_cpu_to_all_gpus(index, co=co)
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
op.restrict_range("nprobe", 2049)
elif args.search_type in (
"multi_gpu_flat_quantizer", "multi_gpu_flat_quantizer_sharded"):
"multi_gpu_flat_quantizer",
"multi_gpu_flat_quantizer_sharded",
):
index_ivf = faiss.extract_index_ivf(index)
new_quantizer = faiss.IndexFlatL2(ds.d)
replace_ivf_quantizer(index_ivf, new_quantizer)
index = faiss.index_cpu_to_all_gpus(index, co=co)
op.restrict_range("nprobe", 2049)
elif args.search_type in (
"multi_gpu_sharded1", "multi_gpu_sharded1_flat",
"multi_gpu_sharded1_ivf"):
"multi_gpu_sharded1",
"multi_gpu_sharded1_flat",
"multi_gpu_sharded1_ivf",
):
print(f"move index to {faiss.get_num_gpus()} GPU")
new_quantizer = faiss.IndexFlatL2(index_ivf.d)
hnsw_quantizer = replace_ivf_quantizer(index_ivf, new_quantizer)
@@ -469,47 +517,43 @@ def main():
index = faiss.index_cpu_to_gpu_multiple_py(res, index, co, gpus)
op.restrict_range("nprobe", 2049)
if args.search_type == "multi_gpu_sharded1":
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
index = ShardedGPUIndex(hnsw_quantizer, index, bs=args.batch_size)
elif args.search_type == "multi_gpu_sharded1_ivf":
centroids = hnsw_quantizer.storage.reconstruct_n()
quantizer = faiss.index_factory(
centroids.shape[1], f"IVF{args.ivf_quant_nlist},Flat")
centroids.shape[1], f"IVF{args.ivf_quant_nlist},Flat"
)
quantizer.train(centroids)
quantizer.add(centroids)
co.shard = False
quantizer = faiss.index_cpu_to_gpu_multiple_py(
res, quantizer, co, gpus)
res, quantizer, co, gpus
)
index = ShardedGPUIndex(quantizer, index, bs=args.batch_size)
op.add_range(
"quantizer_nprobe",
[2 ** i for i in range(9)]
)
op.add_range("quantizer_nprobe", [2**i for i in range(9)])
op.restrict_range("nprobe", 1025)
elif args.search_type == "multi_gpu_sharded1_flat":
quantizer = hnsw_quantizer.storage
quantizer = faiss.index_cpu_to_gpu_multiple_py(
res, quantizer, co, gpus)
res, quantizer, co, gpus
)
index = ShardedGPUIndex(quantizer, index, bs=args.batch_size)
else:
raise RuntimeError()
elif args.search_type in (
"multi_gpu_Csharded1", "multi_gpu_Csharded1_flat",
"multi_gpu_Csharded1_ivf"):
"multi_gpu_Csharded1",
"multi_gpu_Csharded1_flat",
"multi_gpu_Csharded1_ivf",
):
print(f"move index to {faiss.get_num_gpus()} GPU")
co.shard = True
co.common_ivf_quantizer
co.common_ivf_quantizer = True
op.restrict_range("nprobe", 2049)
if args.search_type == "multi_gpu_Csharded1":
op.add_range(
"quantizer_efSearch",
[2 ** i for i in range(10)]
)
op.add_range("quantizer_efSearch", [2**i for i in range(10)])
index = faiss.index_cpu_to_all_gpus(index, co)
elif args.search_type == "multi_gpu_Csharded1_flat":
new_quantizer = faiss.IndexFlatL2(index_ivf.d)
@@ -517,12 +561,10 @@ def main():
index = faiss.index_cpu_to_all_gpus(index, co)
elif args.search_type == "multi_gpu_Csharded1_ivf":
quantizer = faiss.index_factory(
index_ivf.d, f"IVF{args.ivf_quant_nlist},Flat")
quantizer_hnsw = replace_ivf_quantizer(index_ivf, quantizer)
op.add_range(
"quantizer_nprobe",
[2 ** i for i in range(9)]
index_ivf.d, f"IVF{args.ivf_quant_nlist},Flat"
)
quantizer_hnsw = replace_ivf_quantizer(index_ivf, quantizer)
op.add_range("quantizer_nprobe", [2**i for i in range(9)])
index = faiss.index_cpu_to_all_gpus(index, co)
else:
raise RuntimeError()
@@ -543,7 +585,7 @@ def main():
"processor": [l for l in open("/proc/cpuinfo") if "model name" in l][0],
"GPU": list(os.popen("nvidia-smi", "r")),
"mem": open("/proc/meminfo", "r").readlines(),
"pid": os.getpid()
"pid": os.getpid(),
}
op.args = args
if args.stats:
@@ -556,7 +598,9 @@ def main():
(max_perf, min_time) = op.predict_bounds(key)
if not op.is_pareto_optimal(max_perf, min_time):
print(f"SKIP, {max_perf=:.3f} {min_time=:.3f}", )
print(
f"SKIP, {max_perf=:.3f} {min_time=:.3f}",
)
continue
for name, val in parameters.items():
@@ -577,19 +621,21 @@ def main():
recalls = {}
for rank in 1, 10, 100:
recall = (gt[:, :1] == I[:ds.nq, :rank]).sum() / ds.nq
recall = (gt[:, :1] == I[: ds.nq, :rank]).sum() / ds.nq
recalls[rank] = recall
print(f"time={t1 - t0:.3f} s recalls={recalls}")
perf = recalls[1]
op.add_operating_point(key, perf, t1 - t0)
op.all_experiments.append({
"cno": cno,
"key": key,
"parameters": parameters,
"time": t1 - t0,
"recalls": recalls
})
op.all_experiments.append(
{
"cno": cno,
"key": key,
"parameters": parameters,
"time": t1 - t0,
"recalls": recalls,
}
)
if args.stats:
pickle.dump(op, open(args.stats, "wb"))
benchs/bench_index_flat.py
+10 -6
View File
@@ -22,7 +22,7 @@ def format_tab(x):
faiss.cvar.distance_compute_min_k_reservoir = 5
# for have_threads in True, False:
for have_threads in False, :
for have_threads in (False,):
if have_threads:
# good config for Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz
@@ -70,17 +70,21 @@ for have_threads in False, :
t0 = time.time()
index.search(ds.get_queries(), k)
t1 = time.time()
if run >= nrun // 5: # the rest is considered warmup
if run >= nrun // 5: # the rest is considered warmup
times.append((t1 - t0))
times = np.array(times)
if unit == "ms":
times *= 1000
print("search k=%3d t=%.3f ms (± %.4f)" % (
k, np.mean(times), np.std(times)))
print(
"search k=%3d t=%.3f ms (± %.4f)"
% (k, np.mean(times), np.std(times))
)
else:
print("search k=%3d t=%.3f s (± %.4f)" % (
k, np.mean(times), np.std(times)))
print(
"search k=%3d t=%.3f s (± %.4f)"
% (k, np.mean(times), np.std(times))
)
restab1.append(np.mean(times))
print("restab=\n", format_tab(restab))
benchs/bench_ivf_fastscan.py
+7 -9
View File
@@ -11,11 +11,9 @@ import numpy as np
import matplotlib.pyplot as plt
try:
from faiss.contrib.datasets_fb import \
DatasetSIFT1M
from faiss.contrib.datasets_fb import DatasetSIFT1M
except ImportError:
from faiss.contrib.datasets import \
DatasetSIFT1M
from faiss.contrib.datasets import DatasetSIFT1M
# ds = DatasetDeep1B(10**6)
@@ -46,8 +44,8 @@ def eval_recall(index, name):
corrects = (gt == I).sum()
recall = corrects / nq
print(
f'\tnprobe {index.nprobe:3d}, Recall@{k}: '
f'{recall:.6f}, speed: {speed:.6f} ms/query'
f"\tnprobe {index.nprobe:3d}, Recall@{k}: "
f"{recall:.6f}, speed: {speed:.6f} ms/query"
)
return recall, qps
@@ -66,7 +64,7 @@ def eval_and_plot(name, rescale_norm=True, plot=True):
faiss.write_index(index, index_path)
# search params
if hasattr(index, 'rescale_norm'):
if hasattr(index, "rescale_norm"):
index.rescale_norm = rescale_norm
name += f"(rescale_norm={rescale_norm})"
faiss.omp_set_num_threads(1)
@@ -108,5 +106,5 @@ eval_and_plot(f"IVF{nlist},RQ{M-2}x4fs_Nrq2x4")
plt.title("Indices on SIFT1M")
plt.xlabel("Recall@1")
plt.ylabel("QPS")
plt.legend(bbox_to_anchor=(1.02, 0.1), loc='upper left', borderaxespad=0)
plt.savefig("bench_ivf_fastscan.png", bbox_inches='tight')
plt.legend(bbox_to_anchor=(1.02, 0.1), loc="upper left", borderaxespad=0)
plt.savefig("bench_ivf_fastscan.png", bbox_inches="tight")
benchs/bench_ivf_fastscan_single_query.py
+24 -13
View File
@@ -30,7 +30,7 @@ nb, d = xb.shape
nq, d = xq.shape
nt, d = xt.shape
print('the dimension is {}, {}'.format(nb, d))
print("the dimension is {}, {}".format(nb, d))
k = 64
@@ -43,7 +43,7 @@ def eval_recall(index, name, single_query=False):
if single_query:
t0 = time.time()
for row in range(nq):
Ds, Is = index.search(xq[row:row + 1], k=k)
Ds, Is = index.search(xq[row : row + 1], k=k)
D[row, :] = Ds
I[row, :] = Is
t = time.time() - t0
@@ -53,16 +53,21 @@ def eval_recall(index, name, single_query=False):
corrects = (gt[:, :1] == I[:, :k]).sum()
recall = corrects / nq
print(
f'\tnprobe {index.nprobe:3d}, 1Recall@{k}: '
f'{recall:.6f}, speed: {speed:.6f} ms/query'
f"\tnprobe {index.nprobe:3d}, 1Recall@{k}: "
f"{recall:.6f}, speed: {speed:.6f} ms/query"
)
return recall, qps
def eval_and_plot(
name, rescale_norm=True, plot=True, single_query=False,
implem=None, num_threads=1):
name,
rescale_norm=True,
plot=True,
single_query=False,
implem=None,
num_threads=1,
):
index = faiss.index_factory(d, name)
index_path = f"indices/{name}.faissindex"
@@ -75,10 +80,10 @@ def eval_and_plot(
faiss.write_index(index, index_path)
# search params
if hasattr(index, 'rescale_norm'):
if hasattr(index, "rescale_norm"):
index.rescale_norm = rescale_norm
name += f"(rescale_norm={rescale_norm})"
if implem is not None and hasattr(index, 'implem'):
if implem is not None and hasattr(index, "implem"):
index.implem = implem
name += f"(implem={implem})"
if single_query:
@@ -109,12 +114,18 @@ M, nlist = 64, 4096
plt.figure(figsize=(8, 6), dpi=80)
eval_and_plot(f"IVF{nlist},PQ{M}x4fs", num_threads=8)
eval_and_plot(f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=0, num_threads=8)
eval_and_plot(f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=14, num_threads=8)
eval_and_plot(f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=15, num_threads=8)
eval_and_plot(
f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=0, num_threads=8
)
eval_and_plot(
f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=14, num_threads=8
)
eval_and_plot(
f"IVF{nlist},PQ{M}x4fs", single_query=True, implem=15, num_threads=8
)
plt.title("Indices on Bigann50M")
plt.xlabel("1Recall@{}".format(k))
plt.ylabel("QPS")
plt.legend(bbox_to_anchor=(1.02, 0.1), loc='upper left', borderaxespad=0)
plt.savefig("bench_ivf_fastscan.png", bbox_inches='tight')
plt.legend(bbox_to_anchor=(1.02, 0.1), loc="upper left", borderaxespad=0)
plt.savefig("bench_ivf_fastscan.png", bbox_inches="tight")
benchs/bench_ivf_flat_panorama.py
+1 -3
View File
@@ -17,9 +17,7 @@ except ImportError:
from faiss.contrib.datasets import DatasetSIFT1M, DatasetGIST1M
parser = argparse.ArgumentParser()
parser.add_argument(
"--dataset", default="gist1m", choices=["sift1m", "gist1m"]
)
parser.add_argument("--dataset", default="gist1m", choices=["sift1m", "gist1m"])
args = parser.parse_args()
if args.dataset == "sift1m":
benchs/bench_ivfflat_cuvs.py
+60 -34
View File
@@ -26,11 +26,9 @@ import argparse
import rmm
try:
from faiss.contrib.datasets_fb import \
DatasetSIFT1M
from faiss.contrib.datasets_fb import DatasetSIFT1M
except ImportError:
from faiss.contrib.datasets import \
DatasetSIFT1M
from faiss.contrib.datasets import DatasetSIFT1M
# ds = DatasetDeep1B(10**6)
@@ -58,25 +56,33 @@ def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('benchmarking options')
group = parser.add_argument_group("benchmarking options")
aa('--bm_train', default=True,
help='whether to benchmark train operation on GPU index')
aa('--bm_add', default=True,
help='whether to benchmark add operation on GPU index')
aa('--bm_search', default=True,
help='whether to benchmark search operation on GPU index')
aa(
"--bm_train",
default=True,
help="whether to benchmark train operation on GPU index",
)
aa(
"--bm_add",
default=True,
help="whether to benchmark add operation on GPU index",
)
aa(
"--bm_search",
default=True,
help="whether to benchmark search operation on GPU index",
)
group = parser.add_argument_group('IVF options')
aa('--nlist', default=1024, type=int,
help="number of IVF centroids")
group = parser.add_argument_group("IVF options")
aa("--nlist", default=1024, type=int, help="number of IVF centroids")
group = parser.add_argument_group('searching')
group = parser.add_argument_group("searching")
aa('--k', default=10, type=int, help='nb of nearest neighbors')
aa('--nprobe', default=10, help='nb of IVF lists to probe')
aa("--k", default=10, type=int, help="nb of nearest neighbors")
aa("--nprobe", default=10, help="nb of IVF lists to probe")
args = parser.parse_args()
@@ -97,10 +103,10 @@ def bench_train_milliseconds(trainVecs, ncols, nlist, use_cuvs):
index = faiss.GpuIndexIVFFlat(res, ncols, nlist, faiss.METRIC_L2, config)
t0 = time.time()
index.train(trainVecs)
return 1000*(time.time() - t0)
return 1000 * (time.time() - t0)
#warmup
# warmup
xw = rs.rand(nt, d)
bench_train_milliseconds(xw, d, args.nlist, True)
@@ -111,20 +117,26 @@ if args.bm_train:
print("=" * 40)
cuvs_gpu_train_time = bench_train_milliseconds(xt, d, args.nlist, True)
classical_gpu_train_time = bench_train_milliseconds(xt, d, args.nlist, False)
print("Method: IVFFlat, Operation: TRAIN, dim: %d, nlist %d, numTrain: %d, classical GPU train time: %.3f milliseconds, cuVS enabled GPU train time: %.3f milliseconds" % (
d, args.nlist, nt, classical_gpu_train_time, cuvs_gpu_train_time))
classical_gpu_train_time = bench_train_milliseconds(
xt, d, args.nlist, False
)
print(
"Method: IVFFlat, Operation: TRAIN, dim: %d, nlist %d, numTrain: %d, classical GPU train time: %.3f milliseconds, cuVS enabled GPU train time: %.3f milliseconds"
% (d, args.nlist, nt, classical_gpu_train_time, cuvs_gpu_train_time)
)
def bench_add_milliseconds(addVecs, q, use_cuvs):
# construct a GPU index using the same trained coarse quantizer
config = faiss.GpuIndexIVFFlatConfig()
config.use_cuvs = use_cuvs
index_gpu = faiss.GpuIndexIVFFlat(res, q, d, args.nlist, faiss.METRIC_L2, config)
assert(index_gpu.is_trained)
index_gpu = faiss.GpuIndexIVFFlat(
res, q, d, args.nlist, faiss.METRIC_L2, config
)
assert index_gpu.is_trained
t0 = time.time()
index_gpu.add(addVecs)
return 1000*(time.time() - t0)
return 1000 * (time.time() - t0)
if args.bm_add:
@@ -136,8 +148,10 @@ if args.bm_add:
idx_cpu.train(xt)
cuvs_gpu_add_time = bench_add_milliseconds(xb, quantizer, True)
classical_gpu_add_time = bench_add_milliseconds(xb, quantizer, False)
print("Method: IVFFlat, Operation: ADD, dim: %d, nlist %d, numAdd: %d, classical GPU add time: %.3f milliseconds, cuVS enabled GPU add time: %.3f milliseconds" % (
d, args.nlist, nb, classical_gpu_add_time, cuvs_gpu_add_time))
print(
"Method: IVFFlat, Operation: ADD, dim: %d, nlist %d, numAdd: %d, classical GPU add time: %.3f milliseconds, cuVS enabled GPU add time: %.3f milliseconds"
% (d, args.nlist, nb, classical_gpu_add_time, cuvs_gpu_add_time)
)
def bench_search_milliseconds(index, queryVecs, nprobe, k, use_cuvs):
@@ -147,21 +161,33 @@ def bench_search_milliseconds(index, queryVecs, nprobe, k, use_cuvs):
index_gpu.nprobe = nprobe
t0 = time.time()
index_gpu.search(queryVecs, k)
return 1000*(time.time() - t0)
return 1000 * (time.time() - t0)
if args.bm_search:
print("=" * 40)
print("GPU Search Benchmarks")
print("=" * 40)
idx_cpu = faiss.IndexIVFFlat(
faiss.IndexFlatL2(d), d, args.nlist)
idx_cpu = faiss.IndexIVFFlat(faiss.IndexFlatL2(d), d, args.nlist)
idx_cpu.train(xt)
idx_cpu.add(xb)
cuvs_gpu_search_time = bench_search_milliseconds(
idx_cpu, xq, args.nprobe, args.k, True)
idx_cpu, xq, args.nprobe, args.k, True
)
classical_gpu_search_time = bench_search_milliseconds(
idx_cpu, xq, args.nprobe, args.k, False)
print("Method: IVFFlat, Operation: SEARCH, dim: %d, nlist: %d, numVecs: %d, numQuery: %d, nprobe: %d, k: %d, classical GPU search time: %.3f milliseconds, cuVS enabled GPU search time: %.3f milliseconds" % (
d, args.nlist, nb, nq, args.nprobe, args.k, classical_gpu_search_time, cuvs_gpu_search_time))
idx_cpu, xq, args.nprobe, args.k, False
)
print(
"Method: IVFFlat, Operation: SEARCH, dim: %d, nlist: %d, numVecs: %d, numQuery: %d, nprobe: %d, k: %d, classical GPU search time: %.3f milliseconds, cuVS enabled GPU search time: %.3f milliseconds"
% (
d,
args.nlist,
nb,
nq,
args.nprobe,
args.k,
classical_gpu_search_time,
cuvs_gpu_search_time,
)
)
benchs/bench_ivfpq_cuvs.py
+56 -35
View File
@@ -26,11 +26,9 @@ import argparse
import rmm
try:
from faiss.contrib.datasets_fb import \
DatasetSIFT1M
from faiss.contrib.datasets_fb import DatasetSIFT1M
except ImportError:
from faiss.contrib.datasets import \
DatasetSIFT1M
from faiss.contrib.datasets import DatasetSIFT1M
# ds = DatasetDeep1B(10**6)
@@ -60,32 +58,45 @@ def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('benchmarking options')
group = parser.add_argument_group("benchmarking options")
aa('--bm_train', default=True,
help='whether to benchmark train operation on GPU index')
aa('--bm_add', default=True,
help='whether to benchmark add operation on GPU index')
aa('--bm_search', default=True,
help='whether to benchmark search operation on GPU index')
aa(
"--bm_train",
default=True,
help="whether to benchmark train operation on GPU index",
)
aa(
"--bm_add",
default=True,
help="whether to benchmark add operation on GPU index",
)
aa(
"--bm_search",
default=True,
help="whether to benchmark search operation on GPU index",
)
group = parser.add_argument_group('IVF options')
aa('--nlist', default=1024, type=np.int64,
help="number of IVF centroids")
aa('--bits_per_code', default=8, type=np.int64, help='bits per code. Note that < 8 is only supported when cuVS is enabled')
group = parser.add_argument_group("IVF options")
aa("--nlist", default=1024, type=np.int64, help="number of IVF centroids")
aa(
"--bits_per_code",
default=8,
type=np.int64,
help="bits per code. Note that < 8 is only supported when cuVS is enabled",
)
group = parser.add_argument_group('searching')
group = parser.add_argument_group("searching")
aa('--k', default=10, type=int, help='nb of nearest neighbors')
aa('--nprobe', default=10, help='nb of IVF lists to probe')
aa("--k", default=10, type=int, help="nb of nearest neighbors")
aa("--nprobe", default=10, help="nb of IVF lists to probe")
args = parser.parse_args()
print("args:", args)
gt = gt[:, :args.k]
gt = gt[:, : args.k]
nlist = args.nlist
bits_per_code = args.bits_per_code
@@ -114,10 +125,10 @@ def bench_train_milliseconds(trainVecs, use_cuvs):
index = faiss.GpuIndexIVFPQ(res, d, 1024, 32, 8, faiss.METRIC_L2, config)
t0 = time.time()
index.train(trainVecs)
return 1000*(time.time() - t0)
return 1000 * (time.time() - t0)
#warmup
# warmup
xw = rs.rand(nt, d)
bench_train_milliseconds(xw, True)
@@ -129,8 +140,10 @@ if args.bm_train:
cuvs_gpu_train_time = bench_train_milliseconds(xt, True)
classical_gpu_train_time = bench_train_milliseconds(xt, False)
print("TRAIN, dim: %d, nlist %d, numTrain: %d, classical GPU train time: %.3f milliseconds, cuVS enabled GPU train time: %.3f milliseconds" % (
d, nlist, nt, classical_gpu_train_time, cuvs_gpu_train_time))
print(
"TRAIN, dim: %d, nlist %d, numTrain: %d, classical GPU train time: %.3f milliseconds, cuVS enabled GPU train time: %.3f milliseconds"
% (d, nlist, nt, classical_gpu_train_time, cuvs_gpu_train_time)
)
def bench_add_milliseconds(addVecs, index_cpu, use_cuvs):
@@ -138,10 +151,10 @@ def bench_add_milliseconds(addVecs, index_cpu, use_cuvs):
config = faiss.GpuClonerOptions()
config.use_cuvs = use_cuvs
index_gpu = faiss.index_cpu_to_gpu(res, 0, index_cpu, config)
assert(index_gpu.is_trained)
assert index_gpu.is_trained
t0 = time.time()
index_gpu.add(addVecs)
return 1000*(time.time() - t0)
return 1000 * (time.time() - t0)
if args.bm_add:
@@ -153,8 +166,10 @@ if args.bm_add:
index_cpu.train(xt)
cuvs_gpu_add_time = bench_add_milliseconds(xb, index_cpu, True)
classical_gpu_add_time = bench_add_milliseconds(xb, index_cpu, False)
print("ADD, dim: %d, nlist %d, numAdd: %d, classical GPU add time: %.3f milliseconds, cuVS enabled GPU add time: %.3f milliseconds" % (
d, nlist, nb, classical_gpu_add_time, cuvs_gpu_add_time))
print(
"ADD, dim: %d, nlist %d, numAdd: %d, classical GPU add time: %.3f milliseconds, cuVS enabled GPU add time: %.3f milliseconds"
% (d, nlist, nb, classical_gpu_add_time, cuvs_gpu_add_time)
)
def bench_search_milliseconds(index, queryVecs, nprobe, k, use_cuvs):
@@ -164,10 +179,10 @@ def bench_search_milliseconds(index, queryVecs, nprobe, k, use_cuvs):
index_gpu.nprobe = nprobe
t0 = time.time()
_, I = index_gpu.search(queryVecs, k)
return I, 1000*(time.time() - t0)
return I, 1000 * (time.time() - t0)
# Search benchmarks: both indexes have identical IVF centroids and lists.
# Search benchmarks: both indexes have identical IVF centroids and lists.
if args.bm_search:
print("=" * 40)
print("GPU Search Benchmarks")
@@ -177,10 +192,16 @@ if args.bm_search:
index_cpu.add(xb)
cuvs_indices, cuvs_gpu_search_time = bench_search_milliseconds(
index_cpu, xq, args.nprobe, args.k, True)
classical_gpu_indices, classical_gpu_search_time = bench_search_milliseconds(
index_cpu, xq, args.nprobe, args.k, False)
index_cpu, xq, args.nprobe, args.k, True
)
classical_gpu_indices, classical_gpu_search_time = (
bench_search_milliseconds(index_cpu, xq, args.nprobe, args.k, False)
)
cuvs_recall, cuvs_qps = eval_recall(cuvs_indices, cuvs_gpu_search_time)
classical_recall, classical_qps = eval_recall(classical_gpu_indices, classical_gpu_search_time)
print("SEARCH, dim: %d, nlist: %d, numVecs: %d, numQuery: %d, nprobe: %d, k: %d, classical GPU qps: %.3f, cuVS enabled GPU qps: %.3f" % (
d, nlist, nb, nq, args.nprobe, args.k, classical_qps, cuvs_qps))
classical_recall, classical_qps = eval_recall(
classical_gpu_indices, classical_gpu_search_time
)
print(
"SEARCH, dim: %d, nlist: %d, numVecs: %d, numQuery: %d, nprobe: %d, k: %d, classical GPU qps: %.3f, cuVS enabled GPU qps: %.3f"
% (d, nlist, nb, nq, args.nprobe, args.k, classical_qps, cuvs_qps)
)
benchs/bench_partition.py
+15 -14
View File
@@ -8,15 +8,16 @@ import faiss
import numpy as np
def do_partition(n, qin, maxval=65536, seed=123, id_type='int64'):
def do_partition(n, qin, maxval=65536, seed=123, id_type="int64"):
print(
f"n={n} qin={qin} maxval={maxval} id_type={id_type} ",
end="\t", flush=True
end="\t",
flush=True,
)
# print("seed=", seed)
rs = np.random.RandomState(seed)
vals = rs.randint(maxval, size=n).astype('uint16')
vals = rs.randint(maxval, size=n).astype("uint16")
ids = (rs.permutation(n) + 12345).astype(id_type)
sp = faiss.swig_ptr
@@ -37,13 +38,13 @@ def do_partition(n, qin, maxval=65536, seed=123, id_type='int64'):
if isinstance(qin, int):
q = qin
faiss.CMax_uint16_partition_fuzzy(
tab_a.get(), sp(ids), n, q, q, None)
tab_a.get(), sp(ids), n, q, q, None
)
else:
q_min, q_max = qin
q = np.array([-1], dtype='uint64')
q = np.array([-1], dtype="uint64")
faiss.CMax_uint16_partition_fuzzy(
tab_a.get(), sp(ids), n,
q_min, q_max, sp(q)
tab_a.get(), sp(ids), n, q_min, q_max, sp(q)
)
q = q[0]
@@ -56,13 +57,13 @@ def do_partition(n, qin, maxval=65536, seed=123, id_type='int64'):
times = np.array(times[100:]) * 1000000
print(
f"times {times.mean():.3f} µs (± {times.std():.4f} µs) nerr={nerr} "
f"bisect {stats.bissect_cycles / 1e6:.3f} Mcy "
f"compress {stats.compress_cycles / 1e6:.3f} Mcy"
)
do_partition(200, (100, 100))
do_partition(200, (100, 150))
do_partition(2000, (1000, 1000))
@@ -71,9 +72,9 @@ do_partition(20000, (10000, 10000))
do_partition(20000, (10000, 15000))
do_partition(200, (100, 100), id_type='int32')
do_partition(200, (100, 150), id_type='int32')
do_partition(2000, (1000, 1000), id_type='int32')
do_partition(2000, (1000, 1500), id_type='int32')
do_partition(20000, (10000, 10000), id_type='int32')
do_partition(20000, (10000, 15000), id_type='int32')
do_partition(200, (100, 100), id_type="int32")
do_partition(200, (100, 150), id_type="int32")
do_partition(2000, (1000, 1000), id_type="int32")
do_partition(2000, (1000, 1500), id_type="int32")
do_partition(20000, (10000, 10000), id_type="int32")
do_partition(20000, (10000, 15000), id_type="int32")
benchs/bench_polysemous_1bn.py
+46 -45
View File
@@ -18,14 +18,14 @@ from datasets import ivecs_read
def mmap_fvecs(fname):
x = np.memmap(fname, dtype='int32', mode='r')
x = np.memmap(fname, dtype="int32", mode="r")
d = x[0]
return x.view('float32').reshape(-1, d + 1)[:, 1:]
return x.view("float32").reshape(-1, d + 1)[:, 1:]
def mmap_bvecs(fname):
x = np.memmap(fname, dtype='uint8', mode='r')
d = x[:4].view('int32')[0]
x = np.memmap(fname, dtype="uint8", mode="r")
d = x[:4].view("int32")[0]
return x.reshape(-1, d + 4)[:, 4:]
@@ -34,12 +34,12 @@ def mmap_bvecs(fname):
#################################################################
dbname = sys.argv[1]
index_key = sys.argv[2]
dbname = sys.argv[1]
index_key = sys.argv[2]
parametersets = sys.argv[3:]
tmpdir = '/tmp/bench_polysemous'
tmpdir = "/tmp/bench_polysemous"
if not os.path.isdir(tmpdir):
print("%s does not exist, creating it" % tmpdir)
@@ -53,33 +53,32 @@ if not os.path.isdir(tmpdir):
print("Preparing dataset", dbname)
if dbname.startswith('SIFT'):
if dbname.startswith("SIFT"):
# SIFT1M to SIFT1000M
dbsize = int(dbname[4:-1])
xb = mmap_bvecs('bigann/bigann_base.bvecs')
xq = mmap_bvecs('bigann/bigann_query.bvecs')
xt = mmap_bvecs('bigann/bigann_learn.bvecs')
xb = mmap_bvecs("bigann/bigann_base.bvecs")
xq = mmap_bvecs("bigann/bigann_query.bvecs")
xt = mmap_bvecs("bigann/bigann_learn.bvecs")
# trim xb to correct size
xb = xb[:dbsize * 1000 * 1000]
xb = xb[: dbsize * 1000 * 1000]
gt = ivecs_read('bigann/gnd/idx_%dM.ivecs' % dbsize)
gt = ivecs_read("bigann/gnd/idx_%dM.ivecs" % dbsize)
elif dbname == 'Deep1B':
xb = mmap_fvecs('deep1b/base.fvecs')
xq = mmap_fvecs('deep1b/deep1B_queries.fvecs')
xt = mmap_fvecs('deep1b/learn.fvecs')
elif dbname == "Deep1B":
xb = mmap_fvecs("deep1b/base.fvecs")
xq = mmap_fvecs("deep1b/deep1B_queries.fvecs")
xt = mmap_fvecs("deep1b/learn.fvecs")
# deep1B's train is outrageously big
xt = xt[:10 * 1000 * 1000]
gt = ivecs_read('deep1b/deep1B_groundtruth.ivecs')
xt = xt[: 10 * 1000 * 1000]
gt = ivecs_read("deep1b/deep1B_groundtruth.ivecs")
else:
print('unknown dataset', dbname, file=sys.stderr)
print("unknown dataset", dbname, file=sys.stderr)
sys.exit(1)
print("sizes: B %s Q %s T %s gt %s" % (
xb.shape, xq.shape, xt.shape, gt.shape))
print("sizes: B %s Q %s T %s gt %s" % (xb.shape, xq.shape, xt.shape, gt.shape))
nq, d = xq.shape
nb, d = xb.shape
@@ -97,19 +96,18 @@ def choose_train_size(index_key):
n_train = 256 * 1000
if "IVF" in index_key:
matches = re.findall('IVF([0-9]+)', index_key)
matches = re.findall("IVF([0-9]+)", index_key)
ncentroids = int(matches[0])
n_train = max(n_train, 100 * ncentroids)
elif "IMI" in index_key:
matches = re.findall('IMI2x([0-9]+)', index_key)
matches = re.findall("IMI2x([0-9]+)", index_key)
nbit = int(matches[0])
n_train = max(n_train, 256 * (1 << nbit))
return n_train
def get_trained_index():
filename = "%s/%s_%s_trained.index" % (
tmpdir, dbname, index_key)
filename = "%s/%s_%s_trained.index" % (tmpdir, dbname, index_key)
if not os.path.exists(filename):
index = faiss.index_factory(d, index_key)
@@ -119,7 +117,7 @@ def get_trained_index():
xtsub = xt[:n_train]
print("Keeping %d train vectors" % xtsub.shape[0])
# make sure the data is actually in RAM and in float
xtsub = xtsub.astype('float32').copy()
xtsub = xtsub.astype("float32").copy()
index.verbose = True
t0 = time.time()
@@ -138,12 +136,13 @@ def get_trained_index():
# Adding vectors to dataset
#################################################################
def rate_limited_imap(f, l):
'a thread pre-processes the next element'
"a thread pre-processes the next element"
pool = ThreadPool(1)
res = None
for i in l:
res_next = pool.apply_async(f, (i, ))
res_next = pool.apply_async(f, (i,))
if res:
yield res.get()
res = res_next
@@ -151,20 +150,18 @@ def rate_limited_imap(f, l):
def matrix_slice_iterator(x, bs):
" iterate over the lines of x in blocks of size bs"
"iterate over the lines of x in blocks of size bs"
nb = x.shape[0]
block_ranges = [(i0, min(nb, i0 + bs))
for i0 in range(0, nb, bs)]
block_ranges = [(i0, min(nb, i0 + bs)) for i0 in range(0, nb, bs)]
return rate_limited_imap(
lambda i01: x[i01[0]:i01[1]].astype('float32').copy(),
block_ranges)
lambda i01: x[i01[0] : i01[1]].astype("float32").copy(), block_ranges
)
def get_populated_index():
filename = "%s/%s_%s_populated.index" % (
tmpdir, dbname, index_key)
filename = "%s/%s_%s_populated.index" % (tmpdir, dbname, index_key)
if not os.path.exists(filename):
index = get_trained_index()
@@ -172,7 +169,7 @@ def get_populated_index():
t0 = time.time()
for xs in matrix_slice_iterator(xb, 100000):
i1 = i0 + xs.shape[0]
print('\radd %d:%d, %.3f s' % (i0, i1, time.time() - t0), end=' ')
print("\radd %d:%d, %.3f s" % (i0, i1, time.time() - t0), end=" ")
sys.stdout.flush()
index.add(xs)
i0 = i1
@@ -196,23 +193,23 @@ ps = faiss.ParameterSpace()
ps.initialize(index)
# make sure queries are in RAM
xq = xq.astype('float32').copy()
xq = xq.astype("float32").copy()
# a static C++ object that collects statistics about searches
ivfpq_stats = faiss.cvar.indexIVFPQ_stats
ivf_stats = faiss.cvar.indexIVF_stats
if parametersets == ['autotune'] or parametersets == ['autotuneMT']:
if parametersets == ["autotune"] or parametersets == ["autotuneMT"]:
if parametersets == ['autotune']:
if parametersets == ["autotune"]:
faiss.omp_set_num_threads(1)
# setup the Criterion object: optimize for 1-R@1
crit = faiss.OneRecallAtRCriterion(nq, 1)
# by default, the criterion will request only 1 NN
crit.nnn = 100
crit.set_groundtruth(None, gt.astype('int64'))
crit.set_groundtruth(None, gt.astype("int64"))
# then we let Faiss find the optimal parameters by itself
print("exploring operating points")
@@ -233,10 +230,14 @@ else:
# we do queries in a single thread
faiss.omp_set_num_threads(1)
print(' ' * len(parametersets[0]), '\t', 'R@1 R@10 R@100 time %pass')
print(
" " * len(parametersets[0]),
"\t",
"R@1 R@10 R@100 time %pass",
)
for param in parametersets:
print(param, '\t', end=' ')
print(param, "\t", end=" ")
sys.stdout.flush()
ps.set_index_parameters(index, param)
t0 = time.time()
@@ -246,6 +247,6 @@ else:
t1 = time.time()
for rank in 1, 10, 100:
n_ok = (I[:, :rank] == gt[:, :1]).sum()
print("%.4f" % (n_ok / float(nq)), end=' ')
print("%8.3f " % ((t1 - t0) * 1000.0 / nq), end=' ')
print("%.4f" % (n_ok / float(nq)), end=" ")
print("%8.3f " % ((t1 - t0) * 1000.0 / nq), end=" ")
print("%5.2f" % (ivfpq_stats.n_hamming_pass * 100.0 / ivf_stats.ndis))
benchs/bench_polysemous_sift1m.py
+2 -2
View File
@@ -31,13 +31,13 @@ nt = 1
faiss.omp_set_num_threads(1)
print("PQ baseline", end=' ')
print("PQ baseline", end=" ")
index.search_type = faiss.IndexPQ.ST_PQ
t, r = evaluate(index, xq, gt, 1)
print("\t %7.3f ms per query, R@1 %.4f" % (t, r[1]))
for ht in 64, 62, 58, 54, 50, 46, 42, 38, 34, 30:
print("Polysemous", ht, end=' ')
print("Polysemous", ht, end=" ")
index.search_type = faiss.IndexPQ.ST_polysemous
index.polysemous_ht = ht
t, r = evaluate(index, xq, gt, 1)
benchs/bench_pq_tables.py
+6 -4
View File
@@ -21,7 +21,6 @@ def run_bench(d, dsub, nbit=8, metric=None):
pq = faiss.ProductQuantizer(d, M, nbit)
pq.train(faiss.randn((max(1000, pq.ksub * 50), d), 123))
sp = faiss.swig_ptr
times = []
@@ -43,16 +42,19 @@ def run_bench(d, dsub, nbit=8, metric=None):
else:
assert False
t1 = time.time()
if run >= nrun // 5: # the rest is considered warmup
if run >= nrun // 5: # the rest is considered warmup
times.append((t1 - t0))
times = np.array(times) * 1000
print(f"nx={nx}: {np.mean(times):.3f} ms (± {np.std(times):.4f})",
end="\t")
print(
f"nx={nx}: {np.mean(times):.3f} ms (± {np.std(times):.4f})",
end="\t",
)
res.append(times.mean())
print()
return res
# for have_threads in True, False:
for have_threads in False, True:
benchs/bench_pq_transposed_centroid_table.py
+83 -16
View File
@@ -45,11 +45,13 @@ def test_bigann10m(index_file, index_parameters):
index_ivf, vec_transform = unwind_index_ivf(index)
print('params regular transp_centroids regular R@1 R@10 R@100')
print(
"params regular transp_centroids regular R@1 R@10 R@100"
)
for index_parameter in index_parameters:
ps.set_index_parameters(index, index_parameter)
print(index_parameter.ljust(70), end=' ')
print(index_parameter.ljust(70), end=" ")
k = 100
@@ -76,28 +78,96 @@ def test_bigann10m(index_file, index_parameters):
D, I = index.search(xq, k)
t4_1 = time.time()
print(" %9.5f " % (t2_1 - t2_0), end=' ')
print(" %9.5f " % (t3_1 - t3_0), end=' ')
print(" %9.5f " % (t4_1 - t4_0), end=' ')
print(" %9.5f " % (t2_1 - t2_0), end=" ")
print(" %9.5f " % (t3_1 - t3_0), end=" ")
print(" %9.5f " % (t4_1 - t4_0), end=" ")
for rank in 1, 10, 100:
n_ok = (I[:, :rank] == gt[:, :1]).sum()
print("%.4f" % (n_ok / float(nq)), end=' ')
print("%.4f" % (n_ok / float(nq)), end=" ")
print()
if __name__ == "__main__":
faiss.contrib.datasets.dataset_basedir = '/home/aguzhva/ANN_SIFT1B/'
faiss.contrib.datasets.dataset_basedir = "/home/aguzhva/ANN_SIFT1B/"
# represents OPQ32_128,IVF65536_HNSW32,PQ32 index
index_file_1 = "/home/aguzhva/ANN_SIFT1B/run_tests/bench_ivf/indexes/hnsw32/.faissindex"
nprobe_values = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024]
quantizer_efsearch_values = [4, 8, 16, 32, 64, 128, 256, 512]
ht_values = [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54, 56, 58, 60, 62, 64, 66, 68, 70, 72, 74, 76, 78, 80, 82, 84, 86, 88, 90, 92, 94, 96, 98, 100, 102, 104, 106, 108, 110, 112, 114, 116, 118, 120, 122, 124, 126, 128, 256]
ht_values = [
2,
4,
6,
8,
10,
12,
14,
16,
18,
20,
22,
24,
26,
28,
30,
32,
34,
36,
38,
40,
42,
44,
46,
48,
50,
52,
54,
56,
58,
60,
62,
64,
66,
68,
70,
72,
74,
76,
78,
80,
82,
84,
86,
88,
90,
92,
94,
96,
98,
100,
102,
104,
106,
108,
110,
112,
114,
116,
118,
120,
122,
124,
126,
128,
256,
]
# represents OPQ32_128,IVF65536(IVF256,PQHDx4fs,RFlat),PQ32 index
index_file_2 = "/home/aguzhva/ANN_SIFT1B/run_tests/bench_ivf/indexes/pq4/.faissindex"
index_file_2 = (
"/home/aguzhva/ANN_SIFT1B/run_tests/bench_ivf/indexes/pq4/.faissindex"
)
quantizer_k_factor_rf_values = [1, 2, 4, 8, 16, 32, 64]
quantizer_nprobe_values = [1, 2, 4, 8, 16, 32, 64, 128]
@@ -110,9 +180,8 @@ if __name__ == "__main__":
ht = random.choice(ht_values)
index_parameters_1.append(
"nprobe={},quantizer_efSearch={},ht={}".format(
nprobe,
quantizer_efsearch,
ht)
nprobe, quantizer_efsearch, ht
)
)
test_bigann10m(index_file_1, index_parameters_1)
@@ -126,10 +195,8 @@ if __name__ == "__main__":
ht = random.choice(ht_values)
index_parameters_2.append(
"nprobe={},quantizer_k_factor_rf={},quantizer_nprobe={},ht={}".format(
nprobe,
quantizer_k_factor_rf,
quantizer_nprobe,
ht)
nprobe, quantizer_k_factor_rf, quantizer_nprobe, ht
)
)
test_bigann10m(index_file_2, index_parameters_2)
benchs/bench_quantizer.py
+34 -19
View File
@@ -9,11 +9,17 @@ import time
import numpy as np
try:
from faiss.contrib.datasets_fb import \
DatasetSIFT1M, DatasetDeep1B, DatasetBigANN
from faiss.contrib.datasets_fb import (
DatasetSIFT1M,
DatasetDeep1B,
DatasetBigANN,
)
except ImportError:
from faiss.contrib.datasets import \
DatasetSIFT1M, DatasetDeep1B, DatasetBigANN
from faiss.contrib.datasets import (
DatasetSIFT1M,
DatasetDeep1B,
DatasetBigANN,
)
def eval_codec(q, xq, xb, gt):
@@ -29,7 +35,8 @@ def eval_codec(q, xq, xb, gt):
recall = (I[:, 0] == gt[:, 0]).sum() / nq
print(
f"\tencode time: {t1 - t0:.3f} reconstruction error: {recons_err:.3f} "
f"1-recall@1: {recall:.4f} recons_err_compat {err_compat:.3f}")
f"1-recall@1: {recall:.4f} recons_err_compat {err_compat:.3f}"
)
def eval_quantizer(q, xq, xb, gt, xt, variants=None):
@@ -39,7 +46,7 @@ def eval_quantizer(q, xq, xb, gt, xt, variants=None):
q.train(xt)
t1 = time.time()
train_t = t1 - t0
print(f'\ttraining time: {train_t:.3f} s')
print(f"\ttraining time: {train_t:.3f} s")
for name, val in variants:
if name is not None:
print(f"{name}={val}")
@@ -92,19 +99,19 @@ nt, d = xt.shape
# fastest to slowest
if 'lsq-gpu' in todo:
if "lsq-gpu" in todo:
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
ngpus = faiss.get_num_gpus()
lsq.icm_encoder_factory = faiss.GpuIcmEncoderFactory(ngpus)
lsq.verbose = True
eval_quantizer(lsq, xb, xt, 'lsq-gpu')
eval_quantizer(lsq, xb, xt, "lsq-gpu")
if 'pq' in todo:
if "pq" in todo:
pq = faiss.ProductQuantizer(d, M, nbits)
print("===== PQ")
eval_quantizer(pq, xq, xb, gt, xt)
if 'opq' in todo:
if "opq" in todo:
d2 = ((d + M - 1) // M) * M
print("OPQ d2=", d2)
opq = faiss.OPQMatrix(d, M, d2)
@@ -116,33 +123,41 @@ if 'opq' in todo:
print("===== PQ")
eval_quantizer(pq, xq2, xb2, gt, xt2)
if 'prq' in todo:
if "prq" in todo:
print(f"===== PRQ{nsplits}x{Msub}x{nbits}")
prq = faiss.ProductResidualQuantizer(d, nsplits, Msub, nbits)
variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32)]
eval_quantizer(prq, xq, xb, gt, xt, variants=variants)
if 'plsq' in todo:
if "plsq" in todo:
print(f"===== PLSQ{nsplits}x{Msub}x{nbits}")
plsq = faiss.ProductLocalSearchQuantizer(d, nsplits, Msub, nbits)
variants = [("encode_ils_iters", i) for i in (2, 3, 4, 8, 16)]
eval_quantizer(plsq, xq, xb, gt, xt, variants=variants)
if 'rq' in todo:
if "rq" in todo:
print("===== RQ")
rq = faiss.ResidualQuantizer(d, M, nbits, )
rq = faiss.ResidualQuantizer(
d,
M,
nbits,
)
rq.max_beam_size
rq.max_beam_size = 30 # for compatibility with older runs
rq.max_beam_size = 30 # for compatibility with older runs
# rq.train_type = faiss.ResidualQuantizer.Train_default
# rq.verbose = True
variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32)]
eval_quantizer(rq, xq, xb, gt, xt, variants=variants)
if 'rq_lut' in todo:
if "rq_lut" in todo:
print("===== RQ")
rq = faiss.ResidualQuantizer(d, M, nbits, )
rq = faiss.ResidualQuantizer(
d,
M,
nbits,
)
rq.max_beam_size
rq.max_beam_size = 30 # for compatibility with older runs
rq.max_beam_size = 30 # for compatibility with older runs
rq.use_beam_LUT
rq.use_beam_LUT = 1
# rq.train_type = faiss.ResidualQuantizer.Train_default
@@ -150,7 +165,7 @@ if 'rq_lut' in todo:
variants = [("max_beam_size", i) for i in (1, 2, 4, 8, 16, 32, 64)]
eval_quantizer(rq, xq, xb, gt, xt, variants=variants)
if 'lsq' in todo:
if "lsq" in todo:
print("===== LSQ")
lsq = faiss.LocalSearchQuantizer(d, M, nbits)
variants = [("encode_ils_iters", i) for i in (2, 3, 4, 8, 16)]
benchs/bench_refine_panorama.py
+15 -5
View File
@@ -74,7 +74,9 @@ nprobe_list = [4, 16, 64, 256]
kfactor_list = [1, 8, 64, 256, 1024]
print(f"Benchmark on GIST1M with base '{factory}', k={k}, nq={nq}")
print("nprobe k_factor recall_flat qps_flat recall_pano qps_pano dims_scanned(%) speedup(x)")
print(
"nprobe k_factor recall_flat qps_flat recall_pano qps_pano dims_scanned(%) speedup(x)"
)
faiss.omp_set_num_threads(1)
@@ -111,10 +113,16 @@ for nprobe in nprobe_list:
qps_p_list.append(qps_p)
# Draw speedup and recall
plt.plot([kf, kf], [qps_f, qps_p], 'k--', linewidth=1, alpha=0.7)
plt.plot([kf, kf], [qps_f, qps_p], "k--", linewidth=1, alpha=0.7)
mid_y = (qps_f * qps_p) ** 0.5
plt.text(kf + 10, mid_y, f"{speedup:.2f}x\nr={recall_p:.2f}",
ha="left", va="center", fontsize=8)
plt.text(
kf + 10,
mid_y,
f"{speedup:.2f}x\nr={recall_p:.2f}",
ha="left",
va="center",
fontsize=8,
)
plt.plot(kfactor_list, qps_f_list, label="RefineFlat")
plt.plot(kfactor_list, qps_p_list, label=f"RefineFlatPanorama({nlevels})")
@@ -123,7 +131,9 @@ plt.ylim(bottom=100)
plt.xlim(right=kfactor_list[-1] * 1.075)
plt.xlabel("k_factor (k_base = k * k_factor)")
plt.ylabel("QPS")
plt.title(f"GIST1M, base={factory}, nprobe={fixed_nprobe}, nlevels={nlevels}, k={k}")
plt.title(
f"GIST1M, base={factory}, nprobe={fixed_nprobe}, nlevels={nlevels}, k={k}"
)
plt.legend()
plt.tight_layout()
plt.savefig("bench_refine_panorama.png", bbox_inches="tight")
benchs/bench_scalar_quantizer.py
+24 -21
View File
@@ -15,25 +15,27 @@ nq, d = xq.shape
ncent = 256
variants = [(name, getattr(faiss.ScalarQuantizer, name))
for name in dir(faiss.ScalarQuantizer)
if name.startswith('QT_') and name != 'QT_count']
variants = [
(name, getattr(faiss.ScalarQuantizer, name))
for name in dir(faiss.ScalarQuantizer)
if name.startswith("QT_") and name != "QT_count"
]
quantizer = faiss.IndexFlatL2(d)
# quantizer.add(np.zeros((1, d), dtype='float32'))
if False:
for name, qtype in [('flat', 0)] + variants:
for name, qtype in [("flat", 0)] + variants:
print("============== test", name)
t0 = time.time()
if name == 'flat':
index = faiss.IndexIVFFlat(quantizer, d, ncent,
faiss.METRIC_L2)
if name == "flat":
index = faiss.IndexIVFFlat(quantizer, d, ncent, faiss.METRIC_L2)
else:
index = faiss.IndexIVFScalarQuantizer(quantizer, d, ncent,
qtype, faiss.METRIC_L2)
index = faiss.IndexIVFScalarQuantizer(
quantizer, d, ncent, qtype, faiss.METRIC_L2
)
index.nprobe = 16
print("[%.3f s] train" % (time.time() - t0))
@@ -46,7 +48,7 @@ if False:
for rank in 1, 10, 100:
n_ok = (I[:, :rank] == gt[:, :1]).sum()
print("%.4f" % (n_ok / float(nq)), end=' ')
print("%.4f" % (n_ok / float(nq)), end=" ")
print()
if True:
@@ -54,18 +56,19 @@ if True:
print("============== test", name)
for rsname, vals in [('RS_minmax',
[-0.4, -0.2, -0.1, -0.05, 0.0, 0.1, 0.5]),
('RS_meanstd', [0.8, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0]),
('RS_quantiles', [0.02, 0.05, 0.1, 0.15]),
('RS_optim', [0.0])]:
for rsname, vals in [
("RS_minmax", [-0.4, -0.2, -0.1, -0.05, 0.0, 0.1, 0.5]),
("RS_meanstd", [0.8, 1.0, 1.5, 2.0, 3.0, 5.0, 10.0]),
("RS_quantiles", [0.02, 0.05, 0.1, 0.15]),
("RS_optim", [0.0]),
]:
for val in vals:
print("%-15s %5g " % (rsname, val), end=' ')
index = faiss.IndexIVFScalarQuantizer(quantizer, d, ncent,
qtype, faiss.METRIC_L2)
print("%-15s %5g " % (rsname, val), end=" ")
index = faiss.IndexIVFScalarQuantizer(
quantizer, d, ncent, qtype, faiss.METRIC_L2
)
index.nprobe = 16
index.sq.rangestat = getattr(faiss.ScalarQuantizer,
rsname)
index.sq.rangestat = getattr(faiss.ScalarQuantizer, rsname)
index.sq.rangestat_arg = val
@@ -77,5 +80,5 @@ if True:
for rank in 1, 10, 100:
n_ok = (I[:, :rank] == gt[:, :1]).sum()
print("%.4f" % (n_ok / float(nq)), end=' ')
print("%.4f" % (n_ok / float(nq)), end=" ")
print(" %.3f s" % (t1 - t0))
benchs/bench_super_kmeans.py
+1 -3
View File
@@ -87,9 +87,7 @@ def main():
n, d = args.n, args.d
x = gaussian_mixture(n, d, args.k, args.seed)
print("=== bench_super_kmeans ===")
print(
f"n={n} d={d} k={args.k} niter={args.niter} seed={args.seed}"
)
print(f"n={n} d={d} k={args.k} niter={args.niter} seed={args.seed}")
print("\n--- faiss.Clustering ---")
v_dt, v_obj, v_iters = run_vanilla(x, d, args.k, args.niter, args.seed)
benchs/bench_vector_ops.py
+17 -19
View File
@@ -12,12 +12,12 @@ import time
swig_ptr = faiss.swig_ptr
if False:
a = np.arange(10, 14).astype('float32')
b = np.arange(20, 24).astype('float32')
a = np.arange(10, 14).astype("float32")
b = np.arange(20, 24).astype("float32")
faiss.fvec_inner_product (swig_ptr(a), swig_ptr(b), 4)
faiss.fvec_inner_product(swig_ptr(a), swig_ptr(b), 4)
1/0
1 / 0
xd = 100
yd = 1000000
@@ -26,22 +26,21 @@ np.random.seed(1234)
faiss.omp_set_num_threads(1)
print('xd=%d yd=%d' % (xd, yd))
print("xd=%d yd=%d" % (xd, yd))
print('Running inner products test..')
print("Running inner products test..")
for d in 3, 4, 12, 36, 64:
x = faiss.rand(xd * d).reshape(xd, d)
y = faiss.rand(yd * d).reshape(yd, d)
distances = np.empty((xd, yd), dtype='float32')
distances = np.empty((xd, yd), dtype="float32")
t0 = time.time()
for i in range(xd):
faiss.fvec_inner_products_ny(swig_ptr(distances[i]),
swig_ptr(x[i]),
swig_ptr(y),
d, yd)
faiss.fvec_inner_products_ny(
swig_ptr(distances[i]), swig_ptr(x[i]), swig_ptr(y), d, yd
)
t1 = time.time()
# sparse verification
@@ -53,23 +52,22 @@ for d in 3, 4, 12, 36, 64:
num += abs(distances[xi, yi] - np.dot(x[xi], y[yi]))
denom += abs(distances[xi, yi])
print('d=%d t=%.3f s diff=%g' % (d, t1 - t0, num / denom))
print("d=%d t=%.3f s diff=%g" % (d, t1 - t0, num / denom))
print('Running L2sqr test..')
print("Running L2sqr test..")
for d in 3, 4, 12, 36, 64:
x = faiss.rand(xd * d).reshape(xd, d)
y = faiss.rand(yd * d).reshape(yd, d)
distances = np.empty((xd, yd), dtype='float32')
distances = np.empty((xd, yd), dtype="float32")
t0 = time.time()
for i in range(xd):
faiss.fvec_L2sqr_ny(swig_ptr(distances[i]),
swig_ptr(x[i]),
swig_ptr(y),
d, yd)
faiss.fvec_L2sqr_ny(
swig_ptr(distances[i]), swig_ptr(x[i]), swig_ptr(y), d, yd
)
t1 = time.time()
# sparse verification
@@ -81,4 +79,4 @@ for d in 3, 4, 12, 36, 64:
num += abs(distances[xi, yi] - np.sum((x[xi] - y[yi]) ** 2))
denom += abs(distances[xi, yi])
print('d=%d t=%.3f s diff=%g' % (d, t1 - t0, num / denom))
print("d=%d t=%.3f s diff=%g" % (d, t1 - t0, num / denom))
benchs/datasets.py
+3 -3
View File
@@ -10,17 +10,17 @@ import numpy as np
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
a = np.fromfile(fname, dtype="int32")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
def load_sift1M():
print("Loading sift1M...", end='', file=sys.stderr)
print("Loading sift1M...", end="", file=sys.stderr)
xt = fvecs_read("sift1M/sift_learn.fvecs")
xb = fvecs_read("sift1M/sift_base.fvecs")
xq = fvecs_read("sift1M/sift_query.fvecs")
benchs/distributed_ondisk/combined_index.py
+47 -34
View File
@@ -17,15 +17,16 @@ class CombinedIndex:
the info on how to perform searches
"""
def __init__(self, invlist_fnames, empty_index_fname,
masked_index_fname=None):
def __init__(
self, invlist_fnames, empty_index_fname, masked_index_fname=None
):
self.indexes = indexes = []
ilv = faiss.InvertedListsPtrVector()
for fname in invlist_fnames:
if os.path.exists(fname):
print('reading', fname, end='\r', flush=True)
print("reading", fname, end="\r", flush=True)
index = faiss.read_index(fname)
indexes.append(index)
il = faiss.extract_index_ivf(index).invlists
@@ -37,23 +38,24 @@ class CombinedIndex:
self.big_il = faiss.VStackInvertedLists(ilv.size(), ilv.data())
if masked_index_fname:
self.big_il_base = self.big_il
print('loading', masked_index_fname)
print("loading", masked_index_fname)
self.masked_index = faiss.read_index(
masked_index_fname,
faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)
masked_index_fname, faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY
)
self.big_il = faiss.MaskedInvertedLists(
faiss.extract_index_ivf(self.masked_index).invlists,
self.big_il_base)
self.big_il_base,
)
print('loading empty index', empty_index_fname)
print("loading empty index", empty_index_fname)
self.index = faiss.read_index(empty_index_fname)
ntotal = self.big_il.compute_ntotal()
print('replace invlists')
print("replace invlists")
index_ivf = faiss.extract_index_ivf(self.index)
index_ivf.replace_invlists(self.big_il, False)
index_ivf.ntotal = self.index.ntotal = ntotal
index_ivf.parallel_mode = 1 # seems reasonable to do this all the time
index_ivf.parallel_mode = 1 # seems reasonable to do this all the time
quantizer = faiss.downcast_index(index_ivf.quantizer)
quantizer.hnsw.efSearch = 1024
@@ -82,7 +84,6 @@ class CombinedIndex:
coarse_dis, list_nos = quantizer.search(xq, index_ivf.nprobe)
return xq, list_nos, coarse_dis
def ivf_search_preassigned(self, xq, list_nos, coarse_dis, k):
index_ivf = faiss.extract_index_ivf(self.index)
n, d = xq.shape
@@ -91,15 +92,20 @@ class CombinedIndex:
assert list_nos.shape == coarse_dis.shape
assert n2 == n
assert d2 == index_ivf.nprobe
D = np.empty((n, k), dtype='float32')
I = np.empty((n, k), dtype='int64')
D = np.empty((n, k), dtype="float32")
I = np.empty((n, k), dtype="int64")
index_ivf.search_preassigned(
n, faiss.swig_ptr(xq), k,
faiss.swig_ptr(list_nos), faiss.swig_ptr(coarse_dis),
faiss.swig_ptr(D), faiss.swig_ptr(I), False)
n,
faiss.swig_ptr(xq),
k,
faiss.swig_ptr(list_nos),
faiss.swig_ptr(coarse_dis),
faiss.swig_ptr(D),
faiss.swig_ptr(I),
False,
)
return D, I
def ivf_range_search_preassigned(self, xq, list_nos, coarse_dis, radius):
index_ivf = faiss.extract_index_ivf(self.index)
n, d = xq.shape
@@ -111,9 +117,13 @@ class CombinedIndex:
res = faiss.RangeSearchResult(n)
index_ivf.range_search_preassigned(
n, faiss.swig_ptr(xq), radius,
faiss.swig_ptr(list_nos), faiss.swig_ptr(coarse_dis),
res)
n,
faiss.swig_ptr(xq),
radius,
faiss.swig_ptr(list_nos),
faiss.swig_ptr(coarse_dis),
res,
)
lims = faiss.rev_swig_ptr(res.lims, n + 1).copy()
nd = int(lims[-1])
@@ -135,7 +145,8 @@ class CombinedIndex:
def set_prefetch_nthread(self, nt):
for idx in self.indexes:
il = faiss.downcast_InvertedLists(
faiss.extract_index_ivf(idx).invlists)
faiss.extract_index_ivf(idx).invlists
)
il.prefetch_nthread = nt
def set_omp_num_threads(self, nt):
@@ -143,38 +154,40 @@ class CombinedIndex:
class CombinedIndexDeep1B(CombinedIndex):
""" loads a CombinedIndex with the data from the big photodna index """
"""loads a CombinedIndex with the data from the big photodna index"""
def __init__(self):
# set some paths
workdir = "/checkpoint/matthijs/ondisk_distributed/"
# empty index with the proper quantizer
indexfname = workdir + 'trained.faissindex'
indexfname = workdir + "trained.faissindex"
# index that has some invlists that override the big one
masked_index_fname = None
invlist_fnames = [
'%s/hslices/slice%d.faissindex' % (workdir, i)
for i in range(50)
"%s/hslices/slice%d.faissindex" % (workdir, i) for i in range(50)
]
CombinedIndex.__init__(self, invlist_fnames, indexfname, masked_index_fname)
CombinedIndex.__init__(
self, invlist_fnames, indexfname, masked_index_fname
)
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
a = np.fromfile(fname, dtype="int32")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
if __name__ == '__main__':
if __name__ == "__main__":
import time
ci = CombinedIndexDeep1B()
print('loaded index of size ', ci.index.ntotal)
print("loaded index of size ", ci.index.ntotal)
deep1bdir = "/datasets01_101/simsearch/041218/deep1b/"
@@ -187,7 +200,7 @@ if __name__ == '__main__':
t0 = time.time()
D, I = ci.search(xq, 100)
t1 = time.time()
print('nprobe=%d 1-recall@1=%.4f t=%.2fs' % (
nprobe, (I[:, 0] == gt[:, 0]).sum() / len(xq),
t1 - t0
))
print(
"nprobe=%d 1-recall@1=%.4f t=%.2fs"
% (nprobe, (I[:, 0] == gt[:, 0]).sum() / len(xq), t1 - t0)
)
benchs/distributed_ondisk/distributed_kmeans.py
+82 -70
View File
@@ -45,8 +45,8 @@ class DatasetAssignDispatch:
return self.d
def get_subset(self, indices):
res = np.zeros((len(indices), self.d), dtype='float32')
nos = np.searchsorted(self.cs[1:], indices, side='right')
res = np.zeros((len(indices), self.d), dtype="float32")
nos = np.searchsorted(self.cs[1:], indices, side="right")
def handle(i):
mask = nos == i
@@ -58,10 +58,7 @@ class DatasetAssignDispatch:
return res
def assign_to(self, centroids, weights=None):
src = self.imap(
lambda x: x.assign_to(centroids, weights),
self.xes
)
src = self.imap(lambda x: x.assign_to(centroids, weights), self.xes)
I = []
D = []
sum_per_centroid = None
@@ -76,9 +73,9 @@ class DatasetAssignDispatch:
class AssignServer(rpc.Server):
""" Assign version that can be exposed via RPC """
"""Assign version that can be exposed via RPC"""
def __init__(self, s, assign, log_prefix=''):
def __init__(self, s, assign, log_prefix=""):
rpc.Server.__init__(self, s, log_prefix=log_prefix)
self.assign = assign
@@ -86,11 +83,9 @@ class AssignServer(rpc.Server):
return getattr(self.assign, f)
def do_test(todo):
testdata = '/datasets01_101/simsearch/041218/bigann/bigann_learn.bvecs'
testdata = "/datasets01_101/simsearch/041218/bigann/bigann_learn.bvecs"
if os.path.exists(testdata):
x = bvecs_mmap(testdata)
@@ -108,7 +103,7 @@ def do_test(todo):
if "0" in todo:
# reference C++ run
km = faiss.Kmeans(x.shape[1], 1000, niter=20, verbose=True)
km.train(xx.astype('float32'))
km.train(xx.astype("float32"))
if "1" in todo:
# using the Faiss c++ implementation
@@ -117,21 +112,24 @@ def do_test(todo):
if "2" in todo:
# use the dispatch object (on local datasets)
data = DatasetAssignDispatch([
DatasetAssign(xx[20000 * i : 20000 * (i + 1)])
for i in range(5)
], False
data = DatasetAssignDispatch(
[DatasetAssign(xx[20000 * i : 20000 * (i + 1)]) for i in range(5)],
False,
)
kmeans(1000, data, 20)
if "3" in todo:
# same, with GPU
ngpu = faiss.get_num_gpus()
print('using %d GPUs' % ngpu)
data = DatasetAssignDispatch([
DatasetAssignGPU(xx[100000 * i // ngpu: 100000 * (i + 1) // ngpu], i)
for i in range(ngpu)
], True
print("using %d GPUs" % ngpu)
data = DatasetAssignDispatch(
[
DatasetAssignGPU(
xx[100000 * i // ngpu : 100000 * (i + 1) // ngpu], i
)
for i in range(ngpu)
],
True,
)
kmeans(1000, data, 20)
@@ -142,98 +140,112 @@ def main():
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('general options')
aa('--test', default='', help='perform tests (comma-separated numbers)')
group = parser.add_argument_group("general options")
aa("--test", default="", help="perform tests (comma-separated numbers)")
aa('--k', default=0, type=int, help='nb centroids')
aa('--seed', default=1234, type=int, help='random seed')
aa('--niter', default=20, type=int, help='nb iterations')
aa('--gpu', default=-2, type=int, help='GPU to use (-2:none, -1: all)')
aa("--k", default=0, type=int, help="nb centroids")
aa("--seed", default=1234, type=int, help="random seed")
aa("--niter", default=20, type=int, help="nb iterations")
aa("--gpu", default=-2, type=int, help="GPU to use (-2:none, -1: all)")
group = parser.add_argument_group('I/O options')
aa('--indata', default='',
help='data file to load (supported formats fvecs, bvecs, npy')
aa('--i0', default=0, type=int, help='first vector to keep')
aa('--i1', default=-1, type=int, help='last vec to keep + 1')
aa('--out', default='', help='file to store centroids')
aa('--store_each_iteration', default=False, action='store_true',
help='store centroid checkpoints')
group = parser.add_argument_group("I/O options")
aa(
"--indata",
default="",
help="data file to load (supported formats fvecs, bvecs, npy",
)
aa("--i0", default=0, type=int, help="first vector to keep")
aa("--i1", default=-1, type=int, help="last vec to keep + 1")
aa("--out", default="", help="file to store centroids")
aa(
"--store_each_iteration",
default=False,
action="store_true",
help="store centroid checkpoints",
)
group = parser.add_argument_group('server options')
aa('--server', action='store_true', default=False, help='run server')
aa('--port', default=12345, type=int, help='server port')
aa('--when_ready', default=None, help='store host:port to this file when ready')
aa('--ipv4', default=False, action='store_true', help='force ipv4')
group = parser.add_argument_group("server options")
aa("--server", action="store_true", default=False, help="run server")
aa("--port", default=12345, type=int, help="server port")
aa(
"--when_ready",
default=None,
help="store host:port to this file when ready",
)
aa("--ipv4", default=False, action="store_true", help="force ipv4")
group = parser.add_argument_group('client options')
aa('--client', action='store_true', default=False, help='run client')
aa('--servers', default='', help='list of server:port separated by spaces')
group = parser.add_argument_group("client options")
aa("--client", action="store_true", default=False, help="run client")
aa("--servers", default="", help="list of server:port separated by spaces")
args = parser.parse_args()
if args.test:
do_test(args.test.split(','))
do_test(args.test.split(","))
return
# prepare data matrix (either local or remote)
if args.indata:
print('loading ', args.indata)
if args.indata.endswith('.bvecs'):
print("loading ", args.indata)
if args.indata.endswith(".bvecs"):
x = bvecs_mmap(args.indata)
elif args.indata.endswith('.fvecs'):
elif args.indata.endswith(".fvecs"):
x = fvecs_mmap(args.indata)
elif args.indata.endswith('.npy'):
x = np.load(args.indata, mmap_mode='r')
elif args.indata.endswith(".npy"):
x = np.load(args.indata, mmap_mode="r")
else:
raise AssertionError
if args.i1 == -1:
args.i1 = len(x)
x = x[args.i0:args.i1]
x = x[args.i0 : args.i1]
if args.gpu == -2:
data = DatasetAssign(x)
else:
print('moving to GPU')
print("moving to GPU")
data = DatasetAssignGPU(x, args.gpu)
elif args.client:
print('connecting to servers')
print("connecting to servers")
def connect_client(hostport):
host, port = hostport.split(':')
host, port = hostport.split(":")
port = int(port)
print('connecting %s:%d' % (host, port))
print("connecting %s:%d" % (host, port))
client = rpc.Client(host, port, v6=not args.ipv4)
print('client %s:%d ready' % (host, port))
print("client %s:%d ready" % (host, port))
return client
hostports = args.servers.strip().split(' ')
hostports = args.servers.strip().split(" ")
# pool = ThreadPool(len(hostports))
data = DatasetAssignDispatch(
list(map(connect_client, hostports)),
True
)
data = DatasetAssignDispatch(list(map(connect_client, hostports)), True)
else:
raise AssertionError
if args.server:
print('starting server')
print("starting server")
log_prefix = f"{rpc.socket.gethostname()}:{args.port}"
rpc.run_server(
lambda s: AssignServer(s, data, log_prefix=log_prefix),
args.port, report_to_file=args.when_ready,
v6=not args.ipv4)
args.port,
report_to_file=args.when_ready,
v6=not args.ipv4,
)
else:
print('running kmeans')
centroids = kmeans(args.k, data, niter=args.niter, seed=args.seed,
checkpoint=args.out if args.store_each_iteration else None)
if args.out != '':
print('writing centroids to', args.out)
print("running kmeans")
centroids = kmeans(
args.k,
data,
niter=args.niter,
seed=args.seed,
checkpoint=args.out if args.store_each_iteration else None,
)
if args.out != "":
print("writing centroids to", args.out)
np.save(args.out, centroids)
if __name__ == '__main__':
if __name__ == "__main__":
main()
benchs/distributed_ondisk/distributed_query_demo.py
+6 -6
View File
@@ -43,13 +43,13 @@ sindex.set_omp_num_threads(64)
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
a = np.fromfile(fname, dtype="int32")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
deep1bdir = "/datasets01_101/simsearch/041218/deep1b/"
@@ -64,7 +64,7 @@ for nprobe in 1, 10, 100, 1000:
t0 = time.time()
D, I = sindex.search(xq, 100)
t1 = time.time()
print('nprobe=%d 1-recall@1=%.4f t=%.2fs' % (
nprobe, (I[:, 0] == gt[:, 0]).sum() / len(xq),
t1 - t0
))
print(
"nprobe=%d 1-recall@1=%.4f t=%.2fs"
% (nprobe, (I[:, 0] == gt[:, 0]).sum() / len(xq), t1 - t0)
)
benchs/distributed_ondisk/make_index_vslice.py
+42 -30
View File
@@ -11,13 +11,13 @@ from multiprocessing.pool import ThreadPool
def ivecs_mmap(fname):
a = np.memmap(fname, dtype='int32', mode='r')
a = np.memmap(fname, dtype="int32", mode="r")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:]
def fvecs_mmap(fname):
return ivecs_mmap(fname).view('float32')
return ivecs_mmap(fname).view("float32")
def produce_batches(args):
@@ -27,8 +27,10 @@ def produce_batches(args):
if args.i1 == -1:
args.i1 = len(x)
print("Iterating on vectors %d:%d from %s by batches of size %d" % (
args.i0, args.i1, args.input, args.bs))
print(
"Iterating on vectors %d:%d from %s by batches of size %d"
% (args.i0, args.i1, args.input, args.bs)
)
for j0 in range(args.i0, args.i1, args.bs):
j1 = min(j0 + args.bs, args.i1)
@@ -36,7 +38,7 @@ def produce_batches(args):
def rate_limited_iter(l):
'a thread pre-processes the next element'
"a thread pre-processes the next element"
pool = ThreadPool(1)
res = None
@@ -55,42 +57,49 @@ def rate_limited_iter(l):
yield res
res = res_next
deep1bdir = "/datasets01_101/simsearch/041218/deep1b/"
workdir = "/checkpoint/matthijs/ondisk_distributed/"
def main():
parser = argparse.ArgumentParser(
description='make index for a subset of the data')
description="make index for a subset of the data"
)
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('index type')
aa('--inputindex',
default=workdir + 'trained.faissindex',
help='empty input index to fill in')
aa('--nt', default=-1, type=int, help='nb of openmp threads to use')
group = parser.add_argument_group("index type")
aa(
"--inputindex",
default=workdir + "trained.faissindex",
help="empty input index to fill in",
)
aa("--nt", default=-1, type=int, help="nb of openmp threads to use")
group = parser.add_argument_group('db options')
aa('--input', default=deep1bdir + "base.fvecs")
aa('--bs', default=2**18, type=int,
help='batch size for db access')
aa('--i0', default=0, type=int, help='lower bound to index')
aa('--i1', default=-1, type=int, help='upper bound of vectors to index')
group = parser.add_argument_group("db options")
aa("--input", default=deep1bdir + "base.fvecs")
aa("--bs", default=2**18, type=int, help="batch size for db access")
aa("--i0", default=0, type=int, help="lower bound to index")
aa("--i1", default=-1, type=int, help="upper bound of vectors to index")
group = parser.add_argument_group('output')
aa('-o', default='/tmp/x', help='output index')
aa('--keepquantizer', default=False, action='store_true',
help='by default we remove the data from the quantizer to save space')
group = parser.add_argument_group("output")
aa("-o", default="/tmp/x", help="output index")
aa(
"--keepquantizer",
default=False,
action="store_true",
help="by default we remove the data from the quantizer to save space",
)
args = parser.parse_args()
print('args=', args)
print("args=", args)
print('start accessing data')
print("start accessing data")
src = produce_batches(args)
print('loading index', args.inputindex)
print("loading index", args.inputindex)
index = faiss.read_index(args.inputindex)
if args.nt != -1:
@@ -99,21 +108,24 @@ def main():
t0 = time.time()
ntot = 0
for ids, x in rate_limited_iter(src):
print('add %d:%d (%.3f s)' % (ntot, ntot + ids.size, time.time() - t0))
index.add_with_ids(np.ascontiguousarray(x, dtype='float32'), ids)
print("add %d:%d (%.3f s)" % (ntot, ntot + ids.size, time.time() - t0))
index.add_with_ids(np.ascontiguousarray(x, dtype="float32"), ids)
ntot += ids.size
index_ivf = faiss.extract_index_ivf(index)
print('invlists stats: imbalance %.3f' % index_ivf.invlists.imbalance_factor())
print(
"invlists stats: imbalance %.3f" % index_ivf.invlists.imbalance_factor()
)
index_ivf.invlists.print_stats()
if not args.keepquantizer:
print('resetting quantizer content')
print("resetting quantizer content")
index_ivf = faiss.extract_index_ivf(index)
index_ivf.quantizer.reset()
print('store output', args.o)
print("store output", args.o)
faiss.write_index(index, args.o)
if __name__ == '__main__':
if __name__ == "__main__":
main()
benchs/distributed_ondisk/make_trained_index.py
+14 -14
View File
@@ -10,45 +10,45 @@ deep1bdir = "/datasets01_101/simsearch/041218/deep1b/"
workdir = "/checkpoint/matthijs/ondisk_distributed/"
print('Load centroids')
centroids = np.load(workdir + '1M_centroids.npy')
print("Load centroids")
centroids = np.load(workdir + "1M_centroids.npy")
ncent, d = centroids.shape
print('apply random rotation')
print("apply random rotation")
rrot = faiss.RandomRotationMatrix(d, d)
rrot.init(1234)
centroids = rrot.apply_py(centroids)
print('make HNSW index as quantizer')
print("make HNSW index as quantizer")
quantizer = faiss.IndexHNSWFlat(d, 32)
quantizer.hnsw.efSearch = 1024
quantizer.hnsw.efConstruction = 200
quantizer.add(centroids)
print('build index')
print("build index")
index = faiss.IndexPreTransform(
rrot,
faiss.IndexIVFScalarQuantizer(
quantizer, d, ncent, faiss.ScalarQuantizer.QT_6bit
)
)
),
)
def ivecs_mmap(fname):
a = np.memmap(fname, dtype='int32', mode='r')
a = np.memmap(fname, dtype="int32", mode="r")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:]
def fvecs_mmap(fname):
return ivecs_mmap(fname).view('float32')
return ivecs_mmap(fname).view("float32")
print('finish training index')
xt = fvecs_mmap(deep1bdir + 'learn.fvecs')
xt = np.ascontiguousarray(xt[:256 * 1000], dtype='float32')
print("finish training index")
xt = fvecs_mmap(deep1bdir + "learn.fvecs")
xt = np.ascontiguousarray(xt[: 256 * 1000], dtype="float32")
index.train(xt)
print('write output')
faiss.write_index(index, workdir + 'trained.faissindex')
print("write output")
faiss.write_index(index, workdir + "trained.faissindex")
benchs/distributed_ondisk/merge_to_ondisk.py
+17 -20
View File
@@ -7,28 +7,25 @@ import faiss
import argparse
from multiprocessing.pool import ThreadPool
if __name__ == '__main__':
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--inputs', nargs='*', required=True,
help='input indexes to merge')
parser.add_argument('--l0', type=int, default=0)
parser.add_argument('--l1', type=int, default=-1)
parser.add_argument(
"--inputs", nargs="*", required=True, help="input indexes to merge"
)
parser.add_argument("--l0", type=int, default=0)
parser.add_argument("--l1", type=int, default=-1)
parser.add_argument('--nt', default=-1,
help='nb threads')
parser.add_argument("--nt", default=-1, help="nb threads")
parser.add_argument('--output', required=True,
help='output index filename')
parser.add_argument('--outputIL',
help='output invfile filename')
parser.add_argument("--output", required=True, help="output index filename")
parser.add_argument("--outputIL", help="output invfile filename")
args = parser.parse_args()
if args.nt != -1:
print('set nb of threads to', args.nt)
print("set nb of threads to", args.nt)
ils = faiss.InvertedListsPtrVector()
ils_dont_dealloc = []
@@ -38,9 +35,11 @@ if __name__ == '__main__':
def load_index(fname):
print("loading", fname)
try:
index = faiss.read_index(fname, faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY)
index = faiss.read_index(
fname, faiss.IO_FLAG_MMAP | faiss.IO_FLAG_READ_ONLY
)
except RuntimeError as e:
print('could not load %s: %s' % (fname, e))
print("could not load %s: %s" % (fname, e))
return fname, None
print(" %d entries" % index.ntotal)
@@ -57,7 +56,7 @@ if __name__ == '__main__':
il.this.own()
ils_dont_dealloc.append(il)
if (args.l0, args.l1) != (0, -1):
print('restricting to lists %d:%d' % (args.l0, args.l1))
print("restricting to lists %d:%d" % (args.l0, args.l1))
# il = faiss.SliceInvertedLists(il, args.l0, args.l1)
il.crop_invlists(args.l0, args.l1)
@@ -70,13 +69,11 @@ if __name__ == '__main__':
print("loaded %d invlists" % ils.size())
if not args.outputIL:
args.outputIL = args.output + '_invlists'
args.outputIL = args.output + "_invlists"
il0 = ils.at(0)
il = faiss.OnDiskInvertedLists(
il0.nlist, il0.code_size,
args.outputIL)
il = faiss.OnDiskInvertedLists(il0.nlist, il0.code_size, args.outputIL)
print("perform merge")
benchs/distributed_ondisk/search_server.py
+52 -43
View File
@@ -21,7 +21,8 @@ import numpy as np
class MyServer(rpc.Server):
""" Assign version that can be exposed via RPC """
"""Assign version that can be exposed via RPC"""
def __init__(self, s, index):
rpc.Server.__init__(self, s)
self.index = index
@@ -36,32 +37,41 @@ def main():
def aa(*args, **kwargs):
group.add_argument(*args, **kwargs)
group = parser.add_argument_group('server options')
aa('--port', default=12012, type=int, help='server port')
aa('--when_ready_dir', default=None,
help='store host:port to this file when ready')
aa('--ipv4', default=False, action='store_true', help='force ipv4')
aa('--rank', default=0, type=int,
help='rank used as index in the client table')
group = parser.add_argument_group("server options")
aa("--port", default=12012, type=int, help="server port")
aa(
"--when_ready_dir",
default=None,
help="store host:port to this file when ready",
)
aa("--ipv4", default=False, action="store_true", help="force ipv4")
aa(
"--rank",
default=0,
type=int,
help="rank used as index in the client table",
)
args = parser.parse_args()
when_ready = None
if args.when_ready_dir:
when_ready = '%s/%d' % (args.when_ready_dir, args.rank)
when_ready = "%s/%d" % (args.when_ready_dir, args.rank)
print('loading index')
print("loading index")
index = combined_index.CombinedIndexDeep1B()
print('starting server')
print("starting server")
rpc.run_server(
lambda s: MyServer(s, index),
args.port, report_to_file=when_ready,
v6=not args.ipv4)
args.port,
report_to_file=when_ready,
v6=not args.ipv4,
)
if __name__ == '__main__':
if __name__ == "__main__":
main()
@@ -71,12 +81,12 @@ if __name__ == '__main__':
class ResultHeap:
""" Combine query results from a sliced dataset (for k-nn search) """
"""Combine query results from a sliced dataset (for k-nn search)"""
def __init__(self, nq, k):
" nq: number of query vectors, k: number of results per query "
self.I = np.zeros((nq, k), dtype='int64')
self.D = np.zeros((nq, k), dtype='float32')
"nq: number of query vectors, k: number of results per query"
self.I = np.zeros((nq, k), dtype="int64")
self.D = np.zeros((nq, k), dtype="float32")
self.nq, self.k = nq, k
heaps = faiss.float_maxheap_array_t()
heaps.k = k
@@ -91,15 +101,15 @@ class ResultHeap:
assert I.shape == (self.nq, self.k)
I += i0
self.heaps.addn_with_ids(
self.k, faiss.swig_ptr(D),
faiss.swig_ptr(I), self.k)
self.k, faiss.swig_ptr(D), faiss.swig_ptr(I), self.k
)
def finalize(self):
self.heaps.reorder()
def distribute_weights(weights, nbin):
""" assign a set of weights to a smaller set of bins to balance them """
"""assign a set of weights to a smaller set of bins to balance them"""
nw = weights.size
o = weights.argsort()
bins = np.zeros(nbin)
@@ -128,40 +138,37 @@ class SplitPerListIndex:
def set_nprobe(self, nprobe):
self.index.set_nprobe(nprobe)
self.pool.map(
lambda i: self.sub_indexes[i].set_nprobe(nprobe),
range(self.ni)
lambda i: self.sub_indexes[i].set_nprobe(nprobe), range(self.ni)
)
def set_omp_num_threads(self, nt):
faiss.omp_set_num_threads(nt)
self.pool.map(
lambda idx: idx.set_omp_num_threads(nt),
self.sub_indexes
)
self.pool.map(lambda idx: idx.set_omp_num_threads(nt), self.sub_indexes)
def set_parallel_mode(self, pm):
self.index.set_parallel_mode(pm)
self.pool.map(
lambda idx: idx.set_parallel_mode(pm),
self.sub_indexes
)
self.pool.map(lambda idx: idx.set_parallel_mode(pm), self.sub_indexes)
def set_prefetch_nthread(self, nt):
self.index.set_prefetch_nthread(nt)
self.pool.map(
lambda idx: idx.set_prefetch_nthread(nt),
self.sub_indexes
lambda idx: idx.set_prefetch_nthread(nt), self.sub_indexes
)
def balance_lists(self, list_nos):
big_il = self.index.big_il
weights = np.array([big_il.list_size(int(i))
for i in list_nos.ravel()])
weights = np.array([big_il.list_size(int(i)) for i in list_nos.ravel()])
bins, assign = distribute_weights(weights, self.ni)
if self.verbose:
print('bins weight range %d:%d total %d (%.2f MiB)' % (
bins.min(), bins.max(), bins.sum(),
bins.sum() * (self.code_size + 8) / 2 ** 20))
print(
"bins weight range %d:%d total %d (%.2f MiB)"
% (
bins.min(),
bins.max(),
bins.sum(),
bins.sum() * (self.code_size + 8) / 2**20,
)
)
self.nscan = bins.sum()
return assign.reshape(list_nos.shape)
@@ -175,10 +182,11 @@ class SplitPerListIndex:
list_nos_i[assign != i] = -1
t0 = time.time()
Di, Ii = sub_index.ivf_search_preassigned(
xqo, list_nos_i, coarse_dis, k)
xqo, list_nos_i, coarse_dis, k
)
# print(list_nos_i, Ii)
if self.verbose:
print('client %d: %.3f s' % (i, time.time() - t0))
print("client %d: %.3f s" % (i, time.time() - t0))
return Di, Ii
rh = ResultHeap(x.shape[0], k)
@@ -200,9 +208,10 @@ class SplitPerListIndex:
list_nos_i[assign != i] = -1
t0 = time.time()
limi, Di, Ii = sub_index.ivf_range_search_preassigned(
xqo, list_nos_i, coarse_dis, radius)
xqo, list_nos_i, coarse_dis, radius
)
if self.verbose:
print('slice %d: %.3f s' % (i, time.time() - t0))
print("slice %d: %.3f s" % (i, time.time() - t0))
return limi, Di, Ii
D = [[] for i in range(nq)]
@@ -211,7 +220,7 @@ class SplitPerListIndex:
sizes = np.zeros(nq, dtype=int)
for lims, Di, Ii in self.pool.imap(do_query, range(self.ni)):
for i in range(nq):
l0, l1 = lims[i:i + 2]
l0, l1 = lims[i : i + 2]
D[i].append(Di[l0:l1])
I[i].append(Ii[l0:l1])
sizes[i] += l1 - l0
benchs/kmeans_mnist.py
+10 -8
View File
@@ -23,25 +23,25 @@ def load_mnist(fname):
print("load", fname)
f = open(fname)
header = np.fromfile(f, dtype='int8', count=4*4)
header = header.reshape(4, 4)[:, ::-1].copy().view('int32')
header = np.fromfile(f, dtype="int8", count=4 * 4)
header = header.reshape(4, 4)[:, ::-1].copy().view("int32")
print(header)
nim, xd, yd = [int(x) for x in header[1:]]
data = np.fromfile(f, count=nim * xd * yd,
dtype='uint8')
data = np.fromfile(f, count=nim * xd * yd, dtype="uint8")
print(data.shape, nim, xd, yd)
data = data.reshape(nim, xd, yd)
return data
basedir = "/path/to/mnist/data"
x = load_mnist(basedir + 'mnist8m/mnist8m-patterns-idx3-ubyte')
x = load_mnist(basedir + "mnist8m/mnist8m-patterns-idx3-ubyte")
print("reshape")
x = x.reshape(x.shape[0], -1).astype('float32')
x = x.reshape(x.shape[0], -1).astype("float32")
def train_kmeans(x, k, ngpu):
@@ -66,8 +66,9 @@ def train_kmeans(x, k, ngpu):
if ngpu == 1:
index = faiss.GpuIndexFlatL2(res[0], d, flat_config[0])
else:
indexes = [faiss.GpuIndexFlatL2(res[i], d, flat_config[i])
for i in range(ngpu)]
indexes = [
faiss.GpuIndexFlatL2(res[i], d, flat_config[i]) for i in range(ngpu)
]
index = faiss.IndexReplicas()
for sub_index in indexes:
index.addIndex(sub_index)
@@ -81,6 +82,7 @@ def train_kmeans(x, k, ngpu):
return centroids.reshape(k, d)
print("run")
t0 = time.time()
train_kmeans(x, k, ngpu)
contrib/big_batch_search.py
+93 -67
View File
@@ -27,11 +27,7 @@ class BigBatchSearcher:
computation (parallel or not)
"""
def __init__(
self,
index, xq, k,
verbose=0,
use_float16=False):
def __init__(self, index, xq, k, verbose=0, use_float16=False):
# verbosity
self.verbose = verbose
@@ -65,9 +61,8 @@ class BigBatchSearcher:
def report(self, l):
if self.verbose == 1 or (
self.verbose == 2 and (
l > 1000 and time.time() < self.t_display + 1.0
)
self.verbose == 2
and (l > 1000 and time.time() < self.t_display + 1.0)
):
return
t = time.time() - self.t0
@@ -79,8 +74,8 @@ class BigBatchSearcher:
f"wait out {self.t_accu[5]:.3f} "
f"eta {datetime.timedelta(seconds=t*self.index.nlist/(l+1)-t)} "
f"mem {faiss.get_mem_usage_kb()}",
end="\r" if self.verbose <= 2 else "\n",
flush=True,
end="\r" if self.verbose <= 2 else "\n",
flush=True,
)
self.t_display = time.time()
@@ -88,11 +83,12 @@ class BigBatchSearcher:
self.tic("coarse quantization")
bs = 65536
nq = len(self.xq)
q_assign = np.empty((nq, self.index.nprobe), dtype='int32')
q_assign = np.empty((nq, self.index.nprobe), dtype="int32")
for i0 in range(0, nq, bs):
i1 = min(nq, i0 + bs)
q_dis_i, q_assign_i = self.index.quantizer.search(
self.xq[i0:i1], self.index.nprobe)
self.xq[i0:i1], self.index.nprobe
)
# q_dis[i0:i1] = q_dis_i
q_assign[i0:i1] = q_assign_i
self.toc()
@@ -101,18 +97,19 @@ class BigBatchSearcher:
def reorder_assign(self):
self.tic("bucket sort")
q_assign = self.q_assign
q_assign += 1 # move -1 -> 0
q_assign += 1 # move -1 -> 0
self.bucket_lims = faiss.matrix_bucket_sort_inplace(
self.q_assign, nbucket=self.index.nlist + 1, nt=16)
self.q_assign, nbucket=self.index.nlist + 1, nt=16
)
self.query_ids = self.q_assign.ravel()
if self.verbose > 0:
print(' number of -1s:', self.bucket_lims[1])
print(" number of -1s:", self.bucket_lims[1])
self.bucket_lims = self.bucket_lims[1:] # shift back to ignore -1s
del self.q_assign # inplace so let's forget about the old version...
del self.q_assign # inplace so let's forget about the old version...
self.toc()
def prepare_bucket(self, l):
""" prepare the queries and database items for bucket l"""
"""prepare the queries and database items for bucket l"""
t0 = time.time()
index = self.index
# prepare queries
@@ -131,8 +128,8 @@ class BigBatchSearcher:
xb_l = self.decode_func(xb_l)
if self.use_float16:
xb_l = xb_l.astype('float16')
xq_l = xq_l.astype('float16')
xb_l = xb_l.astype("float16")
xq_l = xq_l.astype("float16")
t2 = time.time()
self.t_accu[0] += t1 - t0
@@ -163,7 +160,10 @@ class BigBatchSearcher:
"sizes": self.sizes_in_checkpoint(),
"completed": completed,
"rh": (self.rh.D, self.rh.I),
}, f, -1)
},
f,
-1,
)
os.replace(tmpname, fname)
def read_checkpoint(self, fname):
@@ -176,14 +176,15 @@ class BigBatchSearcher:
class BlockComputer:
""" computation within one bucket """
"""computation within one bucket"""
def __init__(
self,
index,
method="knn_function",
pairwise_distances=faiss.pairwise_distances,
knn=faiss.knn):
self,
index,
method="knn_function",
pairwise_distances=faiss.pairwise_distances,
knn=faiss.knn,
):
self.index = index
if index.__class__ == faiss.IndexIVFFlat:
@@ -192,14 +193,16 @@ class BlockComputer:
by_residual = False
elif index.__class__ == faiss.IndexIVFPQ:
index_help = faiss.IndexPQ(
index.d, index.pq.M, index.pq.nbits, index.metric_type)
index.d, index.pq.M, index.pq.nbits, index.metric_type
)
index_help.pq = index.pq
decode_func = index_help.pq.decode
index_help.is_trained = True
by_residual = index.by_residual
elif index.__class__ == faiss.IndexIVFScalarQuantizer:
index_help = faiss.IndexScalarQuantizer(
index.d, index.sq.qtype, index.metric_type)
index.d, index.sq.qtype, index.metric_type
)
index_help.sq = index.sq
decode_func = index_help.sq.decode
index_help.is_trained = True
@@ -232,22 +235,24 @@ class BlockComputer:
def big_batch_search(
index, xq, k,
method="knn_function",
pairwise_distances=faiss.pairwise_distances,
knn=faiss.knn,
verbose=0,
threaded=0,
use_float16=False,
prefetch_threads=1,
computation_threads=1,
q_assign=None,
checkpoint=None,
checkpoint_freq=7200,
start_list=0,
end_list=None,
crash_at=-1
):
index,
xq,
k,
method="knn_function",
pairwise_distances=faiss.pairwise_distances,
knn=faiss.knn,
verbose=0,
threaded=0,
use_float16=False,
prefetch_threads=1,
computation_threads=1,
q_assign=None,
checkpoint=None,
checkpoint_freq=7200,
start_list=0,
end_list=None,
crash_at=-1,
):
"""
Search queries xq in the IVF index, with a search function that collects
batches of query vectors per inverted list. This can be faster than the
@@ -288,10 +293,11 @@ def big_batch_search(
assert method in ("index", "pairwise_distances", "knn_function")
mem_queries = xq.nbytes
mem_assign = len(xq) * nprobe * np.dtype('int32').itemsize
mem_res = len(xq) * k * (
np.dtype('int64').itemsize
+ np.dtype('float32').itemsize
mem_assign = len(xq) * nprobe * np.dtype("int32").itemsize
mem_res = (
len(xq)
* k
* (np.dtype("int64").itemsize + np.dtype("float32").itemsize)
)
mem_tot = mem_queries + mem_assign + mem_res
if verbose > 0:
@@ -301,16 +307,11 @@ def big_batch_search(
)
bbs = BigBatchSearcher(
index, xq, k,
verbose=verbose,
use_float16=use_float16
index, xq, k, verbose=verbose, use_float16=use_float16
)
comp = BlockComputer(
index,
method=method,
pairwise_distances=pairwise_distances,
knn=knn
index, method=method, pairwise_distances=pairwise_distances, knn=knn
)
bbs.decode_func = comp.decode_func
@@ -351,8 +352,8 @@ def big_batch_search(
# parallel version with granularity 1
def add_results_and_prefetch(to_add, l):
""" perform the addition for the previous bucket and
prefetch the next (if applicable) """
"""perform the addition for the previous bucket and
prefetch the next (if applicable)"""
if to_add is not None:
bbs.add_results_to_heap(*to_add)
if l < index.nlist:
@@ -365,7 +366,8 @@ def big_batch_search(
for l in range(start_list, end_list):
bbs.report(l)
prefetched_bucket_a = pool.apply_async(
add_results_and_prefetch, (to_add, l + 1))
add_results_and_prefetch, (to_add, l + 1)
)
q_subset, xq_l, list_ids, xb_l = prefetched_bucket
bbs.start_t_accu()
D, I = comp.block_search(xq_l, xb_l, list_ids, k)
@@ -390,11 +392,13 @@ def big_batch_search(
):
try:
with ThreadPool(pool_size) as pool:
res = [pool.apply_async(
task,
args=(i, output_queue, input_queue))
res = [
pool.apply_async(
task, args=(i, output_queue, input_queue)
)
for i in range(start_task, end_task)
if i not in completed]
if i not in completed
]
for r in res:
r.get()
pool.close()
@@ -431,7 +435,7 @@ def big_batch_search(
t_wait_out = 0
while True:
t0 = time.time()
logging.info(f'Compute input: task {task_id}')
logging.info(f"Compute input: task {task_id}")
input_value = input_queue.get()
t_wait_in = time.time() - t0
if input_value is None:
@@ -439,7 +443,9 @@ def big_batch_search(
input_queue.put(None)
break
centroid, q_subset, xq_l, list_ids, xb_l = input_value
logging.info(f'Compute work: task {task_id}, centroid {centroid}')
logging.info(
f"Compute work: task {task_id}, centroid {centroid}"
)
t0 = time.time()
if computation_threads > 1:
D, I = comp.block_search(
@@ -448,10 +454,21 @@ def big_batch_search(
else:
D, I = comp.block_search(xq_l, xb_l, list_ids, k)
t_compute = time.time() - t0
logging.info(f'Compute output: task {task_id}, centroid {centroid}')
logging.info(
f"Compute output: task {task_id}, centroid {centroid}"
)
t0 = time.time()
output_queue.put(
(centroid, t_wait_in, t_wait_out, t_compute, q_subset, D, list_ids, I)
(
centroid,
t_wait_in,
t_wait_out,
t_compute,
q_subset,
D,
list_ids,
I,
)
)
t_wait_out = time.time() - t0
logging.info(f"Compute end: {task_id}")
@@ -487,7 +504,16 @@ def big_batch_search(
value = compute_to_main_queue.get()
if not value:
break
centroid, t_wait_in, t_wait_out, t_compute, q_subset, D, list_ids, I = value
(
centroid,
t_wait_in,
t_wait_out,
t_compute,
q_subset,
D,
list_ids,
I,
) = value
# to test checkpointing
if centroid == crash_at:
1 / 0
contrib/client_server.py
+14 -20
View File
@@ -15,7 +15,7 @@ from . import rpc
class SearchServer(rpc.Server):
""" Assign version that can be exposed via RPC """
"""Assign version that can be exposed via RPC"""
def __init__(self, s: int, index: faiss.Index):
rpc.Server.__init__(self, s)
@@ -23,7 +23,7 @@ class SearchServer(rpc.Server):
self.index_ivf = faiss.extract_index_ivf(index)
def set_nprobe(self, nprobe: int) -> int:
""" set nprobe field """
"""set nprobe field"""
self.index_ivf.nprobe = nprobe
def get_ntotal(self) -> int:
@@ -35,23 +35,22 @@ class SearchServer(rpc.Server):
def run_index_server(index: faiss.Index, port: int, v6: bool = False):
""" serve requests for that index forever """
rpc.run_server(
lambda s: SearchServer(s, index),
port, v6=v6)
"""serve requests for that index forever"""
rpc.run_server(lambda s: SearchServer(s, index), port, v6=v6)
############################################################
# Client implementation
############################################################
class ClientIndex:
"""manages a set of distance sub-indexes. The sub_indexes search a
subset of the inverted lists. Searches are merged afterwards
"""
def __init__(self, machine_ports: List[Tuple[str, int]], v6: bool = False):
""" connect to a series of (host, port) pairs """
"""connect to a series of (host, port) pairs"""
self.sub_indexes = []
for machine, port in machine_ports:
self.sub_indexes.append(rpc.Client(machine, port, v6))
@@ -64,28 +63,23 @@ class ClientIndex:
self.verbose = False
def set_nprobe(self, nprobe: int) -> None:
self.pool.map(
lambda idx: idx.set_nprobe(nprobe),
self.sub_indexes
)
self.pool.map(lambda idx: idx.set_nprobe(nprobe), self.sub_indexes)
def set_omp_num_threads(self, nt: int) -> None:
self.pool.map(
lambda idx: idx.set_omp_num_threads(nt),
self.sub_indexes
)
self.pool.map(lambda idx: idx.set_omp_num_threads(nt), self.sub_indexes)
def get_ntotal(self) -> None:
return sum(self.pool.map(
lambda idx: idx.get_ntotal(),
self.sub_indexes
))
return sum(
self.pool.map(lambda idx: idx.get_ntotal(), self.sub_indexes)
)
def search(self, x, k: int):
rh = faiss.ResultHeap(x.shape[0], k)
for Di, Ii in self.pool.imap(lambda idx: idx.search(x, k), self.sub_indexes):
for Di, Ii in self.pool.imap(
lambda idx: idx.search(x, k), self.sub_indexes
):
rh.add_result(Di, Ii)
rh.finalize()
return rh.D, rh.I
contrib/clustering.py
+81 -46
View File
@@ -23,7 +23,9 @@ def print_nop(*arg, **kwargs):
pass
def two_level_clustering(xt, nc1, nc2, rebalance=True, clustering_niter=25, **args):
def two_level_clustering(
xt, nc1, nc2, rebalance=True, clustering_niter=25, **args
):
"""
perform 2-level clustering on a training set xt
nc1 and nc2 are the number of clusters at each level, the final number of
@@ -38,13 +40,13 @@ def two_level_clustering(xt, nc1, nc2, rebalance=True, clustering_niter=25, **ar
log = print if verbose else print_nop
log(f"2-level clustering of {xt.shape} nb 1st level clusters = {nc1} total {nc2}")
log(
f"2-level clustering of {xt.shape} nb 1st level clusters = {nc1} total {nc2}"
)
log("perform coarse training")
km = faiss.Kmeans(
d, nc1, niter=clustering_niter,
max_points_per_centroid=2000,
**args
d, nc1, niter=clustering_niter, max_points_per_centroid=2000, **args
)
km.train(xt)
@@ -58,7 +60,9 @@ def two_level_clustering(xt, nc1, nc2, rebalance=True, clustering_niter=25, **ar
t0 = time.time()
_, assign1 = km.assign(xt)
bc = np.bincount(assign1, minlength=nc1)
log(f"done in {time.time() - t0:.2f} s. Sizes of clusters {min(bc)}-{max(bc)}")
log(
f"done in {time.time() - t0:.2f} s. Sizes of clusters {min(bc)}-{max(bc)}"
)
o = assign1.argsort()
del km
@@ -79,7 +83,11 @@ def two_level_clustering(xt, nc1, nc2, rebalance=True, clustering_niter=25, **ar
t0 = time.time()
for c1 in range(nc1):
nc2 = int(all_nc2[c1])
log(f"[{time.time() - t0:.2f} s] training sub-cluster {c1}/{nc1} nc2={nc2}\r", end="", flush=True)
log(
f"[{time.time() - t0:.2f} s] training sub-cluster {c1}/{nc1} nc2={nc2}\r",
end="",
flush=True,
)
i1 = i0 + bc[c1]
subset = o[i0:i1]
assert np.all(assign1[subset] == c1)
@@ -121,7 +129,9 @@ def train_ivf_index_with_2level(index, xt, **args):
index.train(xt)
def balanced_assignment_with_penalties(x, centroids, alpha = 0.03, num_iter = 20, maxk = 100):
def balanced_assignment_with_penalties(
x, centroids, alpha=0.03, num_iter=20, maxk=100
):
"""
Assign vectors x to centroids with a balance constraint.
@@ -164,15 +174,15 @@ def balanced_assignment_with_penalties(x, centroids, alpha = 0.03, num_iter = 20
n = len(x)
nopt = n / nc # targed bin sizes
# we assign to the top-maxk clusters. The final assignment will pick among these clusters.
# we assign to the top-maxk clusters. The final assignment will pick among these clusters.
full_d2, full_assign = faiss.knn(x, centroids, maxk)
# scalar penalty for each cluster
penalties = np.ones(nc, dtype=np.float32)
for it in range(num_iter):
# compute penalized assignment
penalties2 = penalties ** 2
# compute penalized assignment
penalties2 = penalties**2
full_d2_penalized = full_d2 + penalties2[full_assign]
a0 = full_d2_penalized.argmin(axis=1)
assign = np.take_along_axis(full_assign, a0[:, None], axis=1).ravel()
@@ -183,13 +193,13 @@ def balanced_assignment_with_penalties(x, centroids, alpha = 0.03, num_iter = 20
stats = dict(
alpha=alpha,
imf=imbalance_factor(nc, assign),
mse = ((x - centroids[assign]) ** 2).sum(1).mean(), # recompute MSE
mse=((x - centroids[assign]) ** 2).sum(1).mean(), # recompute MSE
binsize_min=int(binsizes.min()),
binsize_max=int(binsizes.max()),
penalty_min=penalties.min(),
penalty_max=penalties.max(),
)
return assign, stats
@@ -206,7 +216,7 @@ class DatasetAssign:
to centroids. All other implementations offer the same interface"""
def __init__(self, x):
self.x = np.ascontiguousarray(x, dtype='float32')
self.x = np.ascontiguousarray(x, dtype="float32")
def count(self):
return self.x.shape[0]
@@ -226,7 +236,7 @@ class DatasetAssign:
I = I.ravel()
D = D.ravel()
nc, d = centroids.shape
sum_per_centroid = np.zeros((nc, d), dtype='float32')
sum_per_centroid = np.zeros((nc, d), dtype="float32")
if weights is None:
np.add.at(sum_per_centroid, I, self.x)
else:
@@ -236,15 +246,15 @@ class DatasetAssign:
class DatasetAssignGPU(DatasetAssign):
""" GPU version of the previous """
"""GPU version of the previous"""
def __init__(self, x, gpu_id, verbose=False):
DatasetAssign.__init__(self, x)
index = faiss.IndexFlatL2(x.shape[1])
if gpu_id >= 0:
self.index = faiss.index_cpu_to_gpu(
faiss.StandardGpuResources(),
gpu_id, index)
faiss.StandardGpuResources(), gpu_id, index
)
else:
# -1 -> assign to all GPUs
self.index = faiss.index_cpu_to_all_gpus(index)
@@ -256,14 +266,14 @@ class DatasetAssignGPU(DatasetAssign):
def sparse_assign_to_dense(xq, xb, xq_norms=None, xb_norms=None):
""" assignment function for xq is sparse, xb is dense
"""assignment function for xq is sparse, xb is dense
uses a matrix multiplication. The squared norms can be provided if
available.
"""
nq = xq.shape[0]
nb = xb.shape[0]
if xb_norms is None:
xb_norms = (xb ** 2).sum(1)
xb_norms = (xb**2).sum(1)
if xq_norms is None:
xq_norms = np.array(xq.power(2).sum(1))
d2 = xb_norms - 2 * xq @ xb.T
@@ -273,7 +283,8 @@ def sparse_assign_to_dense(xq, xb, xq_norms=None, xb_norms=None):
def sparse_assign_to_dense_blocks(
xq, xb, xq_norms=None, xb_norms=None, qbs=16384, bbs=16384, nt=None):
xq, xb, xq_norms=None, xb_norms=None, qbs=16384, bbs=16384, nt=None
):
"""
decomposes the sparse_assign_to_dense function into blocks to avoid a
possible memory blow up. Can be run in multithreaded mode, because scipy's
@@ -286,7 +297,7 @@ def sparse_assign_to_dense_blocks(
I = -np.ones(nq, dtype=int)
if xb_norms is None:
xb_norms = (xb ** 2).sum(1)
xb_norms = (xb**2).sum(1)
def handle_query_block(i):
xq_block = xq[i : i + qbs]
@@ -334,7 +345,8 @@ class DatasetAssignSparse(DatasetAssign):
def perform_search(self, centroids):
return sparse_assign_to_dense_blocks(
self.x, centroids, xq_norms=self.squared_norms)
self.x, centroids, xq_norms=self.squared_norms
)
def assign_to(self, centroids, weights=None):
D, I = self.perform_search(centroids)
@@ -343,19 +355,19 @@ class DatasetAssignSparse(DatasetAssign):
D = D.ravel()
n = self.x.shape[0]
if weights is None:
weights = np.ones(n, dtype='float32')
weights = np.ones(n, dtype="float32")
nc = len(centroids)
m = scipy.sparse.csc_matrix(
(weights, I, np.arange(n + 1)),
shape=(nc, n))
(weights, I, np.arange(n + 1)), shape=(nc, n)
)
sum_per_centroid = np.array((m * self.x).todense())
return I, D, sum_per_centroid
def imbalance_factor(k, assign):
assign = np.ascontiguousarray(assign, dtype='int64')
assign = np.ascontiguousarray(assign, dtype="int64")
return faiss.imbalance_factor(len(assign), k, faiss.swig_ptr(assign))
@@ -363,13 +375,14 @@ def check_if_torch(x):
if x.__class__ == np.ndarray:
return False
import torch
if isinstance(x, torch.Tensor):
return True
raise NotImplementedError(f"Unknown tensor type {type(x)}")
def reassign_centroids(hassign, centroids, rs=None):
""" reassign centroids when some of them collapse """
"""reassign centroids when some of them collapse"""
if rs is None:
rs = np.random
k, d = centroids.shape
@@ -383,17 +396,18 @@ def reassign_centroids(hassign, centroids, rs=None):
if is_torch:
import torch
fac = torch.ones_like(centroids[0])
else:
fac = np.ones_like(centroids[0])
fac[::2] += 1 / 1024.
fac[1::2] -= 1 / 1024.
fac[::2] += 1 / 1024.0
fac[1::2] -= 1 / 1024.0
# this is a single pass unless there are more than k/2
# empty centroids
while len(empty_cents) > 0:
# choose which centroids to split (numpy)
probas = hassign.astype('float') - 1
probas = hassign.astype("float") - 1
probas[probas < 0] = 0
probas /= probas.sum()
nnz = (probas > 0).sum()
@@ -416,9 +430,15 @@ def reassign_centroids(hassign, centroids, rs=None):
return nsplit
def kmeans(k, data, niter=25, seed=1234, checkpoint=None, verbose=True,
return_stats=False):
def kmeans(
k,
data,
niter=25,
seed=1234,
checkpoint=None,
verbose=True,
return_stats=False,
):
"""Pure python kmeans implementation. Follows the Faiss C++ version
quite closely, but takes a DatasetAssign instead of a training data
matrix. Also redo is not implemented.
@@ -429,8 +449,13 @@ def kmeans(k, data, niter=25, seed=1234, checkpoint=None, verbose=True,
n, d = data.count(), data.dim()
log = print if verbose else print_nop
log(("Clustering %d points in %dD to %d clusters, " +
"%d iterations seed %d") % (n, d, k, niter, seed))
log(
(
"Clustering %d points in %dD to %d clusters, "
+ "%d iterations seed %d"
)
% (n, d, k, niter, seed)
)
rs = np.random.RandomState(seed)
print("preproc...")
@@ -448,10 +473,10 @@ def kmeans(k, data, niter=25, seed=1234, checkpoint=None, verbose=True,
for i in range(niter):
t0s = time.time()
log('assigning', end='\r', flush=True)
log("assigning", end="\r", flush=True)
assign, D, sums = data.assign_to(centroids)
log('compute centroids', end='\r', flush=True)
log("compute centroids", end="\r", flush=True)
t_search_tot += time.time() - t0s
@@ -462,10 +487,11 @@ def kmeans(k, data, niter=25, seed=1234, checkpoint=None, verbose=True,
hassign = np.bincount(assign, minlength=k)
fac = hassign.reshape(-1, 1).astype('float32')
fac = hassign.reshape(-1, 1).astype("float32")
fac[fac == 0] = 1 # quiet warning
if is_torch:
import torch
fac = torch.from_numpy(fac).to(sums.device)
centroids = sums / fac
@@ -477,21 +503,30 @@ def kmeans(k, data, niter=25, seed=1234, checkpoint=None, verbose=True,
"time": (time.time() - t0),
"time_search": t_search_tot,
"imbalance_factor": imbalance_factor(k, assign),
"nsplit": nsplit
"nsplit": nsplit,
}
log((" Iteration %d (%.2f s, search %.2f s): "
"objective=%g imbalance=%.3f nsplit=%d") % (
i, s["time"], s["time_search"],
err, s["imbalance_factor"],
nsplit)
log(
(
" Iteration %d (%.2f s, search %.2f s): "
"objective=%g imbalance=%.3f nsplit=%d"
)
% (
i,
s["time"],
s["time_search"],
err,
s["imbalance_factor"],
nsplit,
)
)
iteration_stats.append(s)
if checkpoint is not None:
log('storing centroids in', checkpoint)
log("storing centroids in", checkpoint)
if is_torch:
import torch
torch.save(centroids, checkpoint)
else:
np.save(checkpoint, centroids)
contrib/datasets.py
+149 -82
View File
@@ -9,31 +9,38 @@ import faiss
import getpass
from .vecs_io import fvecs_read, ivecs_read, bvecs_mmap, fvecs_mmap, bvecs_iter, bvecs_iter_chunked
from .vecs_io import (
fvecs_read,
ivecs_read,
bvecs_mmap,
fvecs_mmap,
bvecs_iter,
bvecs_iter_chunked,
)
from .exhaustive_search import knn
class Dataset:
""" Generic abstract class for a test dataset """
"""Generic abstract class for a test dataset"""
def __init__(self):
""" the constructor should set the following fields: """
"""the constructor should set the following fields:"""
self.d = -1
self.metric = 'L2' # or IP
self.metric = "L2" # or IP
self.nq = -1
self.nb = -1
self.nt = -1
def get_queries(self):
""" return the queries as a (nq, d) array """
"""return the queries as a (nq, d) array"""
raise NotImplementedError()
def get_train(self, maxtrain=None):
""" return the queries as a (nt, d) array """
"""return the queries as a (nt, d) array"""
raise NotImplementedError()
def get_database(self):
""" return the queries as a (nb, d) array """
"""return the queries as a (nb, d) array"""
raise NotImplementedError()
def database_iterator(self, bs=128, split=(1, 0)):
@@ -48,26 +55,28 @@ class Dataset:
nsplit, rank = split
i0, i1 = self.nb * rank // nsplit, self.nb * (rank + 1) // nsplit
for j0 in range(i0, i1, bs):
yield xb[j0: min(j0 + bs, i1)]
yield xb[j0 : min(j0 + bs, i1)]
def get_groundtruth(self, k=None):
""" return the ground truth for k-nearest neighbor search """
"""return the ground truth for k-nearest neighbor search"""
raise NotImplementedError()
def get_groundtruth_range(self, thresh=None):
""" return the ground truth for range search """
"""return the ground truth for range search"""
raise NotImplementedError()
def __str__(self):
return (f"dataset in dimension {self.d}, with metric {self.metric}, "
f"size: Q {self.nq} B {self.nb} T {self.nt}")
return (
f"dataset in dimension {self.d}, with metric {self.metric}, "
f"size: Q {self.nq} B {self.nb} T {self.nt}"
)
def check_sizes(self):
""" runs the previous and checks the sizes of the matrices """
"""runs the previous and checks the sizes of the matrices"""
assert self.get_queries().shape == (self.nq, self.d)
if self.nt > 0:
xt = self.get_train(maxtrain=123)
assert xt.shape == (123, self.d), "shape=%s" % (xt.shape, )
assert xt.shape == (123, self.d), "shape=%s" % (xt.shape,)
assert self.get_database().shape == (self.nb, self.d)
assert self.get_groundtruth(k=13).shape == (self.nq, 13)
@@ -77,10 +86,10 @@ class SyntheticDataset(Dataset):
index
"""
def __init__(self, d, nt, nb, nq, metric='L2', seed=1338):
def __init__(self, d, nt, nb, nq, metric="L2", seed=1338):
Dataset.__init__(self)
self.d, self.nt, self.nb, self.nq = d, nt, nb, nq
d1 = 10 # intrinsic dimension (more or less)
d1 = 10 # intrinsic dimension (more or less)
n = nb + nt + nq
rs = np.random.RandomState(seed)
x = rs.normal(size=(n, d1))
@@ -89,11 +98,11 @@ class SyntheticDataset(Dataset):
# higher factor (>4) -> higher frequency -> less linear
x = x * (rs.rand(d) * 4 + 0.1)
x = np.sin(x)
x = x.astype('float32')
x = x.astype("float32")
self.metric = metric
self.xt = x[:nt]
self.xb = x[nt:nt + nb]
self.xq = x[nt + nb:]
self.xb = x[nt : nt + nb]
self.xq = x[nt + nb :]
def get_queries(self):
return self.xq
@@ -107,8 +116,14 @@ class SyntheticDataset(Dataset):
def get_groundtruth(self, k=100):
return knn(
self.xq, self.xb, k,
faiss.METRIC_L2 if self.metric == 'L2' else faiss.METRIC_INNER_PRODUCT
self.xq,
self.xb,
k,
(
faiss.METRIC_L2
if self.metric == "L2"
else faiss.METRIC_INNER_PRODUCT
),
)[1]
@@ -121,14 +136,15 @@ class SyntheticDataset(Dataset):
username = getpass.getuser()
for dataset_basedir in (
'/datasets01/simsearch/041218/',
'/mnt/vol/gfsai-flash3-east/ai-group/datasets/simsearch/',
f'/home/{username}/simsearch/data/'):
"/datasets01/simsearch/041218/",
"/mnt/vol/gfsai-flash3-east/ai-group/datasets/simsearch/",
f"/home/{username}/simsearch/data/",
):
if os.path.exists(dataset_basedir):
break
else:
# users can link their data directory to `./data`
dataset_basedir = 'data/'
dataset_basedir = "data/"
def set_dataset_basedir(path):
@@ -145,7 +161,7 @@ class DatasetSIFT1M(Dataset):
def __init__(self):
Dataset.__init__(self)
self.d, self.nt, self.nb, self.nq = 128, 100000, 1000000, 10000
self.basedir = dataset_basedir + 'sift1M/'
self.basedir = dataset_basedir + "sift1M/"
def get_queries(self):
return fvecs_read(self.basedir + "sift_query.fvecs")
@@ -166,7 +182,7 @@ class DatasetSIFT1M(Dataset):
def sanitize(x):
return np.ascontiguousarray(x, dtype='float32')
return np.ascontiguousarray(x, dtype="float32")
class DatasetBigANN(Dataset):
@@ -181,17 +197,19 @@ class DatasetBigANN(Dataset):
self.nb_M = nb_M
nb = nb_M * 10**6
self.d, self.nt, self.nb, self.nq = 128, 10**8, nb, 10000
self.basedir = dataset_basedir + 'bigann/'
self.basedir = dataset_basedir + "bigann/"
def get_queries(self):
return sanitize(bvecs_mmap(self.basedir + 'bigann_query.bvecs')[:])
return sanitize(bvecs_mmap(self.basedir + "bigann_query.bvecs")[:])
def get_train(self, maxtrain=None):
maxtrain = maxtrain if maxtrain is not None else self.nt
return sanitize(bvecs_mmap(self.basedir + 'bigann_learn.bvecs')[:maxtrain])
return sanitize(
bvecs_mmap(self.basedir + "bigann_learn.bvecs")[:maxtrain]
)
def get_groundtruth(self, k=None):
gt = ivecs_read(self.basedir + 'gnd/idx_%dM.ivecs' % self.nb_M)
gt = ivecs_read(self.basedir + "gnd/idx_%dM.ivecs" % self.nb_M)
if k is not None:
assert k <= 100
gt = gt[:, :k]
@@ -199,14 +217,16 @@ class DatasetBigANN(Dataset):
def get_database(self):
assert self.nb_M < 100, "dataset too large, use iterator"
return sanitize(bvecs_mmap(self.basedir + 'bigann_base.bvecs')[:self.nb])
return sanitize(
bvecs_mmap(self.basedir + "bigann_base.bvecs")[: self.nb]
)
def database_iterator(self, bs=128, split=(1, 0)):
xb = bvecs_mmap(self.basedir + 'bigann_base.bvecs')
xb = bvecs_mmap(self.basedir + "bigann_base.bvecs")
nsplit, rank = split
i0, i1 = self.nb * rank // nsplit, self.nb * (rank + 1) // nsplit
for j0 in range(i0, i1, bs):
yield sanitize(xb[j0: min(j0 + bs, i1)])
yield sanitize(xb[j0 : min(j0 + bs, i1)])
class DatasetDeep1B(Dataset):
@@ -219,17 +239,19 @@ class DatasetDeep1B(Dataset):
def __init__(self, nb=10**9):
Dataset.__init__(self)
nb_to_name = {
10**5: '100k',
10**6: '1M',
10**7: '10M',
10**8: '100M',
10**9: '1B'
10**5: "100k",
10**6: "1M",
10**7: "10M",
10**8: "100M",
10**9: "1B",
}
assert nb in nb_to_name
self.d, self.nt, self.nb, self.nq = 96, 358480000, nb, 10000
self.basedir = dataset_basedir + 'deep1b/'
self.basedir = dataset_basedir + "deep1b/"
self.gt_fname = "%sdeep%s_groundtruth.ivecs" % (
self.basedir, nb_to_name[self.nb])
self.basedir,
nb_to_name[self.nb],
)
def get_queries(self):
return sanitize(fvecs_read(self.basedir + "deep1B_queries.fvecs"))
@@ -247,14 +269,14 @@ class DatasetDeep1B(Dataset):
def get_database(self):
assert self.nb <= 10**8, "dataset too large, use iterator"
return sanitize(fvecs_mmap(self.basedir + "base.fvecs")[:self.nb])
return sanitize(fvecs_mmap(self.basedir + "base.fvecs")[: self.nb])
def database_iterator(self, bs=128, split=(1, 0)):
xb = fvecs_mmap(self.basedir + "base.fvecs")
nsplit, rank = split
i0, i1 = self.nb * rank // nsplit, self.nb * (rank + 1) // nsplit
for j0 in range(i0, i1, bs):
yield sanitize(xb[j0: min(j0 + bs, i1)])
yield sanitize(xb[j0 : min(j0 + bs, i1)])
class DatasetGlove(Dataset):
@@ -264,28 +286,29 @@ class DatasetGlove(Dataset):
def __init__(self, loc=None, download=False):
import h5py
assert not download, "not implemented"
if not loc:
loc = dataset_basedir + 'glove/glove-100-angular.hdf5'
self.glove_h5py = h5py.File(loc, 'r')
loc = dataset_basedir + "glove/glove-100-angular.hdf5"
self.glove_h5py = h5py.File(loc, "r")
# IP and L2 are equivalent in this case, but it is traditionally seen as an IP dataset
self.metric = 'IP'
self.metric = "IP"
self.d, self.nt = 100, 0
self.nb = self.glove_h5py['train'].shape[0]
self.nq = self.glove_h5py['test'].shape[0]
self.nb = self.glove_h5py["train"].shape[0]
self.nq = self.glove_h5py["test"].shape[0]
def get_queries(self):
xq = np.array(self.glove_h5py['test'])
xq = np.array(self.glove_h5py["test"])
faiss.normalize_L2(xq)
return xq
def get_database(self):
xb = np.array(self.glove_h5py['train'])
xb = np.array(self.glove_h5py["train"])
faiss.normalize_L2(xb)
return xb
def get_groundtruth(self, k=None):
gt = self.glove_h5py['neighbors']
gt = self.glove_h5py["neighbors"]
if k is not None:
assert k <= 100
gt = gt[:, :k]
@@ -301,21 +324,23 @@ class DatasetMusic100(Dataset):
def __init__(self):
Dataset.__init__(self)
self.d, self.nt, self.nb, self.nq = 100, 0, 10**6, 10000
self.metric = 'IP'
self.basedir = dataset_basedir + 'music-100/'
self.metric = "IP"
self.basedir = dataset_basedir + "music-100/"
def get_queries(self):
xq = np.fromfile(self.basedir + 'query_music100.bin', dtype='float32')
xq = np.fromfile(self.basedir + "query_music100.bin", dtype="float32")
xq = xq.reshape(-1, 100)
return xq
def get_database(self):
xb = np.fromfile(self.basedir + 'database_music100.bin', dtype='float32')
xb = np.fromfile(
self.basedir + "database_music100.bin", dtype="float32"
)
xb = xb.reshape(-1, 100)
return xb
def get_groundtruth(self, k=None):
gt = np.load(self.basedir + 'gt.npy')
gt = np.load(self.basedir + "gt.npy")
if k is not None:
assert k <= 100
gt = gt[:, :k]
@@ -331,7 +356,7 @@ class DatasetGIST1M(Dataset):
def __init__(self):
Dataset.__init__(self)
self.d, self.nt, self.nb, self.nq = 960, 100000, 1000000, 10000
self.basedir = dataset_basedir + 'gist1M/'
self.basedir = dataset_basedir + "gist1M/"
def get_queries(self):
return fvecs_read(self.basedir + "gist_query.fvecs")
@@ -350,6 +375,7 @@ class DatasetGIST1M(Dataset):
gt = gt[:, :k]
return gt
class DatasetDINO10B(Dataset):
"""
Data from https://dl.fbaipublicfiles.com/large_objects/dino_vitl_10B/
@@ -357,19 +383,52 @@ class DatasetDINO10B(Dataset):
The dataset is sharded in multiple chunked .bvecs files. Downloading instructions can be obtained with "wget https://dl.fbaipublicfiles.com/large_objects/dino_vitl_10B/README.md".
Supported sizes : 100k 200k 500k 1M ... 5B 10B listed in supported_nbs (see __init__).
"""
def __init__(self, nb, ignore_supported = False):
def __init__(self, nb, ignore_supported=False):
Dataset.__init__(self)
supported_nbs = [100_000, 200_000, 500_000, 1_000_000, 2_000_000, 5_000_000, 10_000_000, 20_000_000, 50_000_000, 100_000_000, 200_000_000, 500_000_000, 1_000_000_000, 2_000_000_000, 5_000_000_000, 10_000_000_000]
supported_nbs = [
100_000,
200_000,
500_000,
1_000_000,
2_000_000,
5_000_000,
10_000_000,
20_000_000,
50_000_000,
100_000_000,
200_000_000,
500_000_000,
1_000_000_000,
2_000_000_000,
5_000_000_000,
10_000_000_000,
]
if nb not in supported_nbs and not ignore_supported:
raise ValueError(f"Unsupported dataset size: {nb}, supported values are: {supported_nbs}")
raise ValueError(
f"Unsupported dataset size: {nb}, supported values are: {supported_nbs}"
)
if not os.path.exists(dataset_basedir):
raise ValueError(f"Provided dataset base directory does not exist: {dataset_basedir}")
raise ValueError(
f"Provided dataset base directory does not exist: {dataset_basedir}"
)
self.basedir = dataset_basedir + "dino_vitl_10B/"
self.indexdir = self.basedir + "chunked_base_10B"
assert os.path.exists(self.indexdir), f"Index path should exist, check your dataset path: {self.indexdir}"
assert os.path.exists(
self.indexdir
), f"Index path should exist, check your dataset path: {self.indexdir}"
self.queriesdir = self.basedir + "queries_clean.bvecs"
assert os.path.exists(self.queriesdir), f"Queries path should exist as dataset size {nb} is supported: {self.queriesdir}"
self.gtsdir = self.basedir + "gts/" + "gts_dino_patch_" + str(nb) + "_" + "k10.npy"
assert os.path.exists(
self.queriesdir
), f"Queries path should exist as dataset size {nb} is supported: {self.queriesdir}"
self.gtsdir = (
self.basedir
+ "gts/"
+ "gts_dino_patch_"
+ str(nb)
+ "_"
+ "k10.npy"
)
self.train_queriesdir = self.basedir + "train_queries_99M.bvecs"
self.nb = nb
self.d = 1024
@@ -377,31 +436,36 @@ class DatasetDINO10B(Dataset):
self.nt = 99_000_000
self.metric = "L2"
def get_queries(self):
"""Get all vectors as a single array"""
queries = bvecs_mmap(self.queriesdir)
return sanitize(queries)
def get_train(self, maxtrain = None):
def get_train(self, maxtrain=None):
"""Get training query vectors as a single array"""
if maxtrain is None or maxtrain > 10_000_000:
raise NotImplementedError("The training set is potentially too large to fit in RAM (400 GB of data). Please use train_iterator or use maxtrain parameter below 10_000_000 to get the first maxtrain training vectors.")
raise NotImplementedError(
"The training set is potentially too large to fit in RAM (400 GB of data). Please use train_iterator or use maxtrain parameter below 10_000_000 to get the first maxtrain training vectors."
)
return sanitize(bvecs_mmap(self.train_queriesdir)[:maxtrain])
def get_database(self):
"""Get all database vectors as a single array"""
if self.nb > 10_000_000:
raise NotImplementedError("The dataset is potentially too large to fit in RAM. Please use database_iterator or use a dataset size equal to or below 10_000_000.")
raise NotImplementedError(
"The dataset is potentially too large to fit in RAM. Please use database_iterator or use a dataset size equal to or below 10_000_000."
)
else:
return sanitize(bvecs_iter_chunked(self.indexdir, batch_size=self.nb).__next__())
return sanitize(
bvecs_iter_chunked(self.indexdir, batch_size=self.nb).__next__()
)
def database_iterator(self, bs=10_000):
"""Iterator over the database of size nb, corresponding to the first nb vectors in the .bvecs file"""
total_read = 0
for batch in bvecs_iter_chunked(self.indexdir, batch_size=bs):
if total_read + batch.shape[0] > self.nb:
batch = batch[:self.nb - total_read]
batch = batch[: self.nb - total_read]
yield sanitize(batch)
total_read += batch.shape[0]
if total_read >= self.nb:
@@ -412,10 +476,12 @@ class DatasetDINO10B(Dataset):
for batch in bvecs_iter(self.train_queriesdir, batch_size=bs):
yield sanitize(batch)
def get_groundtruth(self, k = 10):
def get_groundtruth(self, k=10):
"""Get ground truth from .npy file"""
if k > 10:
raise NotImplementedError("Ground truth files only available for k<=10")
raise NotImplementedError(
"Ground truth files only available for k<=10"
)
gts = np.load(self.gtsdir)
gts = gts[:, :k]
return gts
@@ -423,29 +489,30 @@ class DatasetDINO10B(Dataset):
def distance(self):
return "euclidean"
def dataset_from_name(dataset='deep1M', download=False):
""" converts a string describing a dataset to a Dataset object
def dataset_from_name(dataset="deep1M", download=False):
"""converts a string describing a dataset to a Dataset object
Supports sift1M, bigann1M..bigann1B, deep1M..deep1B, music-100 and glove
"""
if dataset == 'sift1M':
if dataset == "sift1M":
return DatasetSIFT1M()
elif dataset == 'gist1M':
elif dataset == "gist1M":
return DatasetGIST1M()
elif dataset.startswith('bigann'):
elif dataset.startswith("bigann"):
dbsize = 1000 if dataset == "bigann1B" else int(dataset[6:-1])
return DatasetBigANN(nb_M=dbsize)
elif dataset.startswith("deep"):
szsuf = dataset[4:]
if szsuf[-1] == 'M':
dbsize = 10 ** 6 * int(szsuf[:-1])
elif szsuf == '1B':
dbsize = 10 ** 9
elif szsuf[-1] == 'k':
if szsuf[-1] == "M":
dbsize = 10**6 * int(szsuf[:-1])
elif szsuf == "1B":
dbsize = 10**9
elif szsuf[-1] == "k":
dbsize = 1000 * int(szsuf[:-1])
else:
assert False, "did not recognize suffix " + szsuf
contrib/evaluation.py
+51 -49
View File
@@ -15,22 +15,19 @@ from multiprocessing.pool import ThreadPool
def knn_intersection_measure(I1, I2):
""" computes the intersection measure of two result tables
"""
"""computes the intersection measure of two result tables"""
nq, rank = I1.shape
assert I2.shape == (nq, rank)
ninter = sum(
np.intersect1d(I1[i], I2[i]).size
for i in range(nq)
)
ninter = sum(np.intersect1d(I1[i], I2[i]).size for i in range(nq))
return ninter / I1.size
###############################################################
# Range search results can be compared with Precision-Recall
def filter_range_results(lims, D, I, thresh):
""" select a set of results """
"""select a set of results"""
nq = lims.size - 1
mask = D < thresh
new_lims = np.zeros_like(lims)
@@ -41,13 +38,13 @@ def filter_range_results(lims, D, I, thresh):
def range_PR(lims_ref, Iref, lims_new, Inew, mode="overall"):
"""compute the precision and recall of range search results. The
function does not take the distances into account. """
function does not take the distances into account."""
def ref_result_for(i):
return Iref[lims_ref[i]:lims_ref[i + 1]]
return Iref[lims_ref[i] : lims_ref[i + 1]]
def new_result_for(i):
return Inew[lims_new[i]:lims_new[i + 1]]
return Inew[lims_new[i] : lims_new[i + 1]]
nq = lims_ref.size - 1
assert lims_new.size - 1 == nq
@@ -75,12 +72,12 @@ def range_PR(lims_ref, Iref, lims_new, Inew, mode="overall"):
lims_ref[1:] - lims_ref[:-1],
lims_new[1:] - lims_new[:-1],
ninter,
mode=mode
mode=mode,
)
def counts_to_PR(ngt, nres, ninter, mode="overall"):
""" computes a precision-recall for a set of queries.
"""computes a precision-recall for a set of queries.
ngt = nb of GT results per query
nres = nb of found results per query
ninter = nb of correct results per query (smaller than nres of course)
@@ -127,7 +124,7 @@ def counts_to_PR(ngt, nres, ninter, mode="overall"):
def sort_range_res_2(lims, D, I):
""" sort 2 arrays using the first as key """
"""sort 2 arrays using the first as key"""
I2 = np.empty_like(I)
D2 = np.empty_like(D)
nq = len(lims) - 1
@@ -152,12 +149,16 @@ def sort_range_res_1(lims, I):
def range_PR_multiple_thresholds(
lims_ref, Iref,
lims_new, Dnew, Inew,
thresholds,
mode="overall", do_sort="ref,new"
):
""" compute precision-recall values for range search results
lims_ref,
Iref,
lims_new,
Dnew,
Inew,
thresholds,
mode="overall",
do_sort="ref,new",
):
"""compute precision-recall values for range search results
for several thresholds on the "new" results.
This is to plot PR curves
"""
@@ -170,7 +171,7 @@ def range_PR_multiple_thresholds(
Inew, Dnew = sort_range_res_2(lims_new, Dnew, Inew)
def ref_result_for(i):
return Iref[lims_ref[i]:lims_ref[i + 1]]
return Iref[lims_ref[i] : lims_ref[i + 1]]
def new_result_for(i):
l0, l1 = lims_new[i], lims_new[i + 1]
@@ -193,7 +194,7 @@ def range_PR_multiple_thresholds(
return
# which offsets we are interested in
nres= np.searchsorted(res_dis, thresholds)
nres = np.searchsorted(res_dis, thresholds)
counts[q, :, 1] = nres
if gt_ids.size == 0:
@@ -216,8 +217,7 @@ def range_PR_multiple_thresholds(
recalls = np.zeros(nt)
for t in range(nt):
p, r = counts_to_PR(
counts[:, t, 0], counts[:, t, 1], counts[:, t, 2],
mode=mode
counts[:, t, 0], counts[:, t, 1], counts[:, t, 2], mode=mode
)
precisions[t] = p
recalls[t] = r
@@ -229,28 +229,29 @@ def range_PR_multiple_thresholds(
# Functions that compare search results with a reference result.
# They are intended for use in tests
def _cluster_tables_with_tolerance(tab1, tab2, thr):
""" for two tables, cluster them by merging values closer than thr.
Returns the cluster ids for each table element """
"""for two tables, cluster them by merging values closer than thr.
Returns the cluster ids for each table element"""
tab = np.hstack([tab1, tab2])
tab.sort()
n = len(tab)
diffs = np.ones(n)
diffs[1:] = tab[1:] - tab[:-1]
unique_vals = tab[diffs > thr]
idx1 = np.searchsorted(unique_vals, tab1, side='right') - 1
idx2 = np.searchsorted(unique_vals, tab2, side='right') - 1
idx1 = np.searchsorted(unique_vals, tab1, side="right") - 1
idx2 = np.searchsorted(unique_vals, tab2, side="right") - 1
return idx1, idx2
def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, rtol=1e-5):
""" test that knn search results are identical, with possible ties.
Raise if not. """
"""test that knn search results are identical, with possible ties.
Raise if not."""
np.testing.assert_allclose(Dref, Dnew, rtol=rtol)
# here we have to be careful because of draws
testcase = unittest.TestCase() # because it makes nice error messages
testcase = unittest.TestCase() # because it makes nice error messages
for i in range(len(Iref)):
if np.all(Iref[i] == Inew[i]): # easy case
if np.all(Iref[i] == Inew[i]): # easy case
continue
# otherwise collect elements per distance
@@ -265,10 +266,9 @@ def check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, rtol=1e-5):
testcase.assertEqual(set(Iref[i, mask]), set(Inew[i, mask]))
def check_ref_range_results(Lref, Dref, Iref,
Lnew, Dnew, Inew):
""" compare range search results wrt. a reference result,
throw if it fails """
def check_ref_range_results(Lref, Dref, Iref, Lnew, Dnew, Inew):
"""compare range search results wrt. a reference result,
throw if it fails"""
np.testing.assert_array_equal(Lref, Lnew)
nq = len(Lref) - 1
for i in range(nq):
@@ -277,12 +277,14 @@ def check_ref_range_results(Lref, Dref, Iref,
Ii_new = Inew[l0:l1]
Di_ref = Dref[l0:l1]
Di_new = Dnew[l0:l1]
if np.all(Ii_ref == Ii_new): # easy
if np.all(Ii_ref == Ii_new): # easy
pass
else:
def sort_by_ids(I, D):
o = I.argsort()
return I[o], D[o]
# sort both
(Ii_ref, Di_ref) = sort_by_ids(Ii_ref, Di_ref)
(Ii_new, Di_new) = sort_by_ids(Ii_new, Di_new)
@@ -309,11 +311,11 @@ class OperatingPoints:
self.suboptimal_points = []
def compare_keys(self, k1, k2):
""" return -1 if k1 > k2, 1 if k2 > k1, 0 otherwise """
"""return -1 if k1 > k2, 1 if k2 > k1, 0 otherwise"""
raise NotImplemented
def do_nothing_key(self):
""" parameters to say we do nothing, takes 0 time and has 0 performance"""
"""parameters to say we do nothing, takes 0 time and has 0 performance"""
raise NotImplemented
def is_pareto_optimal(self, perf_new, t_new):
@@ -323,15 +325,15 @@ class OperatingPoints:
return True
def predict_bounds(self, key):
""" predicts the bound on time and performance """
"""predicts the bound on time and performance"""
min_time = 0.0
max_perf = 1.0
for key2, perf, t in self.operating_points + self.suboptimal_points:
cmp = self.compare_keys(key, key2)
if cmp > 0: # key2 > key
if cmp > 0: # key2 > key
if t > min_time:
min_time = t
if cmp < 0: # key2 < key
if cmp < 0: # key2 < key
if perf < max_perf:
max_perf = perf
return max_perf, min_time
@@ -347,8 +349,7 @@ class OperatingPoints:
while i < len(self.operating_points):
op_Ls, perf2, t2 = self.operating_points[i]
if perf >= perf2 and t < t2:
self.suboptimal_points.append(
self.operating_points.pop(i))
self.suboptimal_points.append(self.operating_points.pop(i))
else:
i += 1
self.operating_points.append((key, perf, t))
@@ -388,7 +389,7 @@ class OperatingPointsWithRanges(OperatingPoints):
return int(np.prod([len(values) for name, values in self.ranges]))
def sample_experiments(self, n_autotune, rs=np.random):
""" sample a set of experiments of max size n_autotune
"""sample a set of experiments of max size n_autotune
(run all experiments in random order if n_autotune is 0)
"""
assert n_autotune == 0 or n_autotune >= 2
@@ -398,7 +399,8 @@ class OperatingPointsWithRanges(OperatingPoints):
experiments = rs.permutation(totex - 2)
else:
experiments = rs.choice(
totex - 2, size=n_autotune - 2, replace=False)
totex - 2, size=n_autotune - 2, replace=False
)
experiments = [0, totex - 1] + [int(cno) + 1 for cno in experiments]
return experiments
@@ -415,12 +417,11 @@ class OperatingPointsWithRanges(OperatingPoints):
def get_parameters(self, k):
"""Convert a key to a dictionary with parameter values"""
return {
name: values[k[i]]
for i, (name, values) in enumerate(self.ranges)
name: values[k[i]] for i, (name, values) in enumerate(self.ranges)
}
def restrict_range(self, name, max_val):
""" remove too large values from a range"""
"""remove too large values from a range"""
for name2, values in self.ranges:
if name == name2:
val2 = [v for v in values if v < max_val]
@@ -476,6 +477,7 @@ class RepeatTimer:
enters a loop. It focuses on ms-scale times because for second scale
it's usually less relevant to repeat the operation.
"""
def __init__(self, warmup=0, nt=-1, runs=1, max_secs=np.inf):
assert warmup < runs
self.warmup = warmup
@@ -494,5 +496,5 @@ class RepeatTimer:
return np.std(self.times) * 1000 if len(self.times) > 1 else 0.0
def nruns(self):
""" effective number of runs (may be lower than runs - warmup due to timeout)"""
"""effective number of runs (may be lower than runs - warmup due to timeout)"""
return len(self.times)
contrib/exhaustive_search.py
+70 -45
View File
@@ -12,7 +12,9 @@ import logging
LOG = logging.getLogger(__name__)
def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, shard=False, ngpu=-1):
def knn_ground_truth(
xq, db_iterator, k, metric_type=faiss.METRIC_L2, shard=False, ngpu=-1
):
"""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.
@@ -28,7 +30,7 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, shard=Fals
ngpu = faiss.get_num_gpus()
if ngpu:
LOG.info('running on %d GPUs' % ngpu)
LOG.info("running on %d GPUs" % ngpu)
co = faiss.GpuMultipleClonerOptions()
co.shard = shard
index = faiss.index_cpu_to_all_gpus(index, co=co, ngpu=ngpu)
@@ -50,12 +52,11 @@ def knn_ground_truth(xq, db_iterator, k, metric_type=faiss.METRIC_L2, shard=Fals
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.
@@ -71,13 +72,14 @@ def range_search_gpu(xq, r2, index_gpu, index_cpu, gpu_k=1024):
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=:}")
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')
D = D.astype("int16")
t2 = 0
lim_remain = None
if index_cpu is not None:
@@ -96,16 +98,19 @@ def range_search_gpu(xq, r2, index_gpu, index_cpu, gpu_k=1024):
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)
lim_remain, D_remain, I_remain = index_cpu.range_search(
xq[mask], r2
)
if is_binary_index:
D_remain = D_remain.astype('int16')
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
faiss.CombinerRangeKNNint16
if is_binary_index
else faiss.CombinerRangeKNNfloat
)
combiner = CombinerRangeKNN(nq, k, r2, keep_max)
@@ -120,11 +125,11 @@ def range_search_gpu(xq, r2, index_gpu, index_cpu, gpu_k=1024):
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')
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')
I_res = np.empty(nres, dtype="int64")
combiner.write_result(sp(D_res), sp(I_res))
else:
D_res, I_res = [], []
@@ -150,21 +155,27 @@ def range_search_gpu(xq, r2, index_gpu, index_cpu, gpu_k=1024):
return L_res, D_res, I_res
def range_ground_truth(xq, db_iterator, threshold, metric_type=faiss.METRIC_L2,
shard=False, ngpu=-1):
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')
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)
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)
@@ -192,8 +203,8 @@ def range_ground_truth(xq, db_iterator, threshold, metric_type=faiss.METRIC_L2,
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')
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]
@@ -205,7 +216,7 @@ def range_ground_truth(xq, db_iterator, threshold, metric_type=faiss.METRIC_L2,
def threshold_radius_nres(nres, dis, ids, thresh, keep_max=False):
""" select a set of results """
"""select a set of results"""
if keep_max:
mask = dis > thresh
else:
@@ -213,14 +224,14 @@ def threshold_radius_nres(nres, dis, ids, thresh, keep_max=False):
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()
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 """
"""restrict range-search results to those below a given radius"""
if keep_max:
mask = dis > thresh
else:
@@ -246,24 +257,32 @@ def apply_maxres(res_batches, target_nres, keep_max=False):
alldis.partition(target_nres)
radius = alldis[target_nres]
if alldis.dtype == 'float32':
if alldis.dtype == "float32":
radius = float(radius)
else:
radius = int(radius)
LOG.debug(' setting radius to %s' % 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)
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)
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):
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.
@@ -287,7 +306,7 @@ def range_search_max_results(index, query_iterator, radius,
ngpu = faiss.get_num_gpus()
if ngpu:
LOG.info('running on %d GPUs' % 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)
@@ -314,54 +333,60 @@ def range_search_max_results(index, query_iterator, radius,
t1 = time.time()
if is_binary_index:
# weird Faiss quirk that returns floats for Hamming distances
Di = Di.astype('int16')
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))
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
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.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)
"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
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 = 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
"""produces batches of progressively increasing sizes. This is useful to
adjust the search radius progressively without overflowing with
intermediate results """
intermediate results"""
nq = len(xq)
bs = start_bs
i = 0
while i < nq:
xqi = xq[i:i + bs]
xqi = xq[i : i + bs]
yield xqi
if bs < max_bs:
bs *= 2
contrib/factory_tools.py
+21 -17
View File
@@ -8,14 +8,14 @@ import re
def get_code_size(d, indexkey):
""" size of one vector in an index in dimension d
"""size of one vector in an index in dimension d
constructed with factory string indexkey"""
if indexkey == "Flat":
return d * 4
if indexkey.endswith(",RFlat"):
return d * 4 + get_code_size(d, indexkey[:-len(",RFlat")])
return d * 4 + get_code_size(d, indexkey[: -len(",RFlat")])
mo = re.match("IVF\\d+(_HNSW32)?,(.*)$", indexkey)
if mo:
@@ -33,42 +33,42 @@ def get_code_size(d, indexkey):
if mo:
return get_code_size(d, mo.group(1)) + get_code_size(d, mo.group(2))
mo = re.match('PQ(\\d+)x(\\d+)(fs|fsr)?$', indexkey)
mo = re.match("PQ(\\d+)x(\\d+)(fs|fsr)?$", indexkey)
if mo:
return (int(mo.group(1)) * int(mo.group(2)) + 7) // 8
mo = re.match('PQ(\\d+)\\+(\\d+)$', indexkey)
mo = re.match("PQ(\\d+)\\+(\\d+)$", indexkey)
if mo:
return (int(mo.group(1)) + int(mo.group(2)))
return int(mo.group(1)) + int(mo.group(2))
mo = re.match('PQ(\\d+)$', indexkey)
mo = re.match("PQ(\\d+)$", indexkey)
if mo:
return int(mo.group(1))
if indexkey == "HNSW32" or indexkey == "HNSW32,Flat":
return d * 4 + 64 * 4 # roughly
return d * 4 + 64 * 4 # roughly
if indexkey == 'SQ8':
if indexkey == "SQ8":
return d
elif indexkey == 'SQ4':
elif indexkey == "SQ4":
return (d + 1) // 2
elif indexkey == 'SQ6':
elif indexkey == "SQ6":
return (d * 6 + 7) // 8
elif indexkey == 'SQfp16':
elif indexkey == "SQfp16":
return d * 2
elif indexkey == 'SQbf16':
elif indexkey == "SQbf16":
return d * 2
mo = re.match('PCAR?(\\d+),(.*)$', indexkey)
mo = re.match("PCAR?(\\d+),(.*)$", indexkey)
if mo:
return get_code_size(int(mo.group(1)), mo.group(2))
mo = re.match('OPQ\\d+_(\\d+),(.*)$', indexkey)
mo = re.match("OPQ\\d+_(\\d+),(.*)$", indexkey)
if mo:
return get_code_size(int(mo.group(1)), mo.group(2))
mo = re.match('OPQ\\d+,(.*)$', indexkey)
mo = re.match("OPQ\\d+,(.*)$", indexkey)
if mo:
return get_code_size(d, mo.group(1))
mo = re.match('RR(\\d+),(.*)$', indexkey)
mo = re.match("RR(\\d+),(.*)$", indexkey)
if mo:
return get_code_size(int(mo.group(1)), mo.group(2))
raise RuntimeError("cannot parse " + indexkey)
@@ -158,7 +158,11 @@ def reverse_index_factory(index):
return f"PQ{index.pq.M}x{index.pq.nbits}"
elif isinstance(index, faiss.IndexLSH):
return "LSH" + ("r" if index.rotate_data else "") + ("t" if index.train_thresholds else "")
return (
"LSH"
+ ("r" if index.rotate_data else "")
+ ("t" if index.train_thresholds else "")
)
elif isinstance(index, faiss.IndexScalarQuantizer):
return sq_names[index.sq.qtype]
contrib/inspect_tools.py
+24 -31
View File
@@ -8,12 +8,12 @@ import faiss
def get_invlist(invlists, l):
""" returns the inverted lists content as a pair of (list_ids, list_codes).
"""returns the inverted lists content as a pair of (list_ids, list_codes).
The codes are reshaped to a proper size
"""
invlists = faiss.downcast_InvertedLists(invlists)
ls = invlists.list_size(l)
list_ids = np.zeros(ls, dtype='int64')
list_ids = np.zeros(ls, dtype="int64")
ids = codes = None
try:
ids = invlists.get_ids(l)
@@ -21,13 +21,13 @@ def get_invlist(invlists, l):
faiss.memcpy(faiss.swig_ptr(list_ids), ids, list_ids.nbytes)
codes = invlists.get_codes(l)
if invlists.code_size != faiss.InvertedLists.INVALID_CODE_SIZE:
list_codes = np.zeros((ls, invlists.code_size), dtype='uint8')
list_codes = np.zeros((ls, invlists.code_size), dtype="uint8")
else:
# it's a BlockInvertedLists
npb = invlists.n_per_block
bs = invlists.block_size
ls_round = (ls + npb - 1) // npb
list_codes = np.zeros((ls_round, bs // npb, npb), dtype='uint8')
list_codes = np.zeros((ls_round, bs // npb, npb), dtype="uint8")
if ls > 0:
faiss.memcpy(faiss.swig_ptr(list_codes), codes, list_codes.nbytes)
finally:
@@ -39,28 +39,27 @@ def get_invlist(invlists, l):
def get_invlist_sizes(invlists):
""" return the array of sizes of the inverted lists """
return np.array([
invlists.list_size(i)
for i in range(invlists.nlist)
], dtype='int64')
"""return the array of sizes of the inverted lists"""
return np.array(
[invlists.list_size(i) for i in range(invlists.nlist)], dtype="int64"
)
def print_object_fields(obj):
""" list values all fields of an object known to SWIG """
"""list values all fields of an object known to SWIG"""
for name in obj.__class__.__swig_getmethods__:
print(f"{name} = {getattr(obj, name)}")
def get_pq_centroids(pq):
""" return the PQ centroids as an array """
"""return the PQ centroids as an array"""
cen = faiss.vector_to_array(pq.centroids)
return cen.reshape(pq.M, pq.ksub, pq.dsub)
def get_LinearTransform_matrix(pca):
""" extract matrix + bias from the PCA object
"""extract matrix + bias from the PCA object
works for any linear transform (OPQ, random rotation, etc.)
"""
b = faiss.vector_to_array(pca.b)
@@ -69,10 +68,10 @@ def get_LinearTransform_matrix(pca):
def make_LinearTransform_matrix(A, b=None):
""" make a linear transform from a matrix and a bias term (optional)"""
"""make a linear transform from a matrix and a bias term (optional)"""
d_out, d_in = A.shape
if b is not None:
assert b.shape == (d_out, )
assert b.shape == (d_out,)
lt = faiss.LinearTransform(d_in, d_out, b is not None)
faiss.copy_array_to_vector(A.ravel(), lt.A)
if b is not None:
@@ -83,35 +82,29 @@ def make_LinearTransform_matrix(A, b=None):
def get_additive_quantizer_codebooks(aq):
""" return to codebooks of an additive quantizer """
"""return to codebooks of an additive quantizer"""
codebooks = faiss.vector_to_array(aq.codebooks).reshape(-1, aq.d)
co = faiss.vector_to_array(aq.codebook_offsets)
return [
codebooks[co[i]:co[i + 1]]
for i in range(aq.M)
]
return [codebooks[co[i] : co[i + 1]] for i in range(aq.M)]
def get_flat_data(index):
""" copy and return the data matrix in an IndexFlat """
"""copy and return the data matrix in an IndexFlat"""
xb = faiss.vector_to_array(index.codes).view("float32")
return xb.reshape(index.ntotal, index.d)
def get_flat_codes(index_flat):
""" get the codes from an indexFlatCodes as an array """
def get_flat_codes(index_flat):
"""get the codes from an indexFlatCodes as an array"""
return faiss.vector_to_array(index_flat.codes).reshape(
index_flat.ntotal, index_flat.code_size)
index_flat.ntotal, index_flat.code_size
)
def get_NSG_neighbors(nsg):
""" get the neighbor list for the vectors stored in the NSG structure, as
a N-by-K matrix of indices """
"""get the neighbor list for the vectors stored in the NSG structure, as
a N-by-K matrix of indices"""
graph = nsg.get_final_graph()
neighbors = np.zeros((graph.N, graph.K), dtype='int32')
faiss.memcpy(
faiss.swig_ptr(neighbors),
graph.data,
neighbors.nbytes
)
neighbors = np.zeros((graph.N, graph.K), dtype="int32")
faiss.memcpy(faiss.swig_ptr(neighbors), graph.data, neighbors.nbytes)
return neighbors
contrib/ivf_tools.py
+8 -12
View File
@@ -14,16 +14,14 @@ def add_preassigned(index_ivf, x, a, ids=None):
Add elements to an IVF index, where the assignment is already computed
"""
n, d = x.shape
assert a.shape == (n, )
assert a.shape == (n,)
if isinstance(index_ivf, faiss.IndexBinaryIVF):
d *= 8
assert d == index_ivf.d
if ids is not None:
assert ids.shape == (n, )
assert ids.shape == (n,)
ids = faiss.swig_ptr(ids)
index_ivf.add_core(
n, faiss.swig_ptr(x), ids, faiss.swig_ptr(a)
)
index_ivf.add_core(n, faiss.swig_ptr(x), ids, faiss.swig_ptr(a))
def search_preassigned(index_ivf, xq, k, list_nos, coarse_dis=None):
@@ -83,9 +81,7 @@ def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None):
sp = faiss.swig_ptr
index_ivf.range_search_preassigned_c(
n, sp(x), radius,
sp(list_nos), sp(coarse_dis),
res
n, sp(x), radius, sp(list_nos), sp(coarse_dis), res
)
# get pointers and copy them
lims = faiss.rev_swig_ptr(res.lims, n + 1).copy()
@@ -96,7 +92,7 @@ def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None):
def replace_ivf_quantizer(index_ivf, new_quantizer):
""" replace the IVF quantizer with a flat quantizer and return the
"""replace the IVF quantizer with a flat quantizer and return the
old quantizer"""
if new_quantizer.ntotal == 0:
centroids = index_ivf.quantizer.reconstruct_n()
@@ -120,15 +116,15 @@ def replace_ivf_quantizer(index_ivf, new_quantizer):
def permute_invlists(index_ivf, perm):
""" Apply some permutation to the inverted lists, and modify the quantizer
"""Apply some permutation to the inverted lists, and modify the quantizer
entries accordingly.
Perm is an array of size nlist, where old_index = perm[new_index]
"""
nlist, = perm.shape
(nlist,) = perm.shape
assert index_ivf.nlist == nlist
quantizer = faiss.downcast_index(index_ivf.quantizer)
assert quantizer.ntotal == index_ivf.nlist
perm = np.ascontiguousarray(perm, dtype='int64')
perm = np.ascontiguousarray(perm, dtype="int64")
# just make sure it's a permutation...
bc = np.bincount(perm, minlength=nlist)
contrib/ondisk.py
+7 -2
View File
@@ -11,7 +11,10 @@ LOG = logging.getLogger(__name__)
def merge_ondisk(
trained_index: faiss.Index, shard_fnames: List[str], ivfdata_fname: str, shift_ids=False
trained_index: faiss.Index,
shard_fnames: List[str],
ivfdata_fname: str,
shift_ids=False,
) -> None:
"""Add the contents of the indexes stored in shard_fnames into the index
trained_index. The on-disk data is stored in ivfdata_fname"""
@@ -51,7 +54,9 @@ def merge_ondisk(
ivf_vector.push_back(ivf)
LOG.info("merge %d inverted lists " % ivf_vector.size())
ntotal = invlists.merge_from_multiple(ivf_vector.data(), ivf_vector.size(), shift_ids)
ntotal = invlists.merge_from_multiple(
ivf_vector.data(), ivf_vector.size(), shift_ids
)
# now replace the inverted lists in the output index
index.ntotal = index_ivf.ntotal = ntotal
contrib/rpc.py
+55 -50
View File
@@ -27,8 +27,8 @@ LOG = logging.getLogger(__name__)
PORT = 12032
safe_modules = {
'numpy',
'numpy.core.multiarray',
"numpy",
"numpy.core.multiarray",
}
@@ -39,49 +39,51 @@ class RestrictedUnpickler(pickle.Unpickler):
if module in safe_modules:
return getattr(importlib.import_module(module), name)
# Forbid everything else.
raise pickle.UnpicklingError("global '%s.%s' is forbidden" %
(module, name))
raise pickle.UnpicklingError(
"global '%s.%s' is forbidden" % (module, name)
)
class FileSock:
" wraps a socket so that it is usable by pickle/cPickle "
"wraps a socket so that it is usable by pickle/cPickle"
def __init__(self,sock):
def __init__(self, sock):
self.sock = sock
self.nr=0
self.nr = 0
def write(self, buf):
# print("sending %d bytes"%len(buf))
#self.sock.sendall(buf)
# self.sock.sendall(buf)
# print("...done")
bs = 512 * 1024
ns = 0
while ns < len(buf):
sent = self.sock.send(buf[ns:ns + bs])
sent = self.sock.send(buf[ns : ns + bs])
ns += sent
def read(self,bs=512*1024):
#if self.nr==10000: pdb.set_trace()
self.nr+=1
def read(self, bs=512 * 1024):
# if self.nr==10000: pdb.set_trace()
self.nr += 1
# print("read bs=%d"%bs)
b = []
nb = 0
while len(b)<bs:
while len(b) < bs:
# print(' loop')
rb = self.sock.recv(bs - nb)
if not rb: break
if not rb:
break
b.append(rb)
nb += len(rb)
return b''.join(b)
return b"".join(b)
def readline(self):
# print("readline!")
"""may be optimized..."""
s=bytes()
s = bytes()
while True:
c=self.read(1)
s+=c
if len(c)==0 or chr(c[0])=='\n':
c = self.read(1)
s += c
if len(c) == 0 or chr(c[0]) == "\n":
return s
@@ -99,7 +101,7 @@ class Server:
transparently from a client
"""
def __init__(self, s, logf=sys.stderr, log_prefix=''):
def __init__(self, s, logf=sys.stderr, log_prefix=""):
self.logf = logf
self.log_prefix = log_prefix
@@ -108,7 +110,6 @@ class Server:
self.conn = s
self.fs = FileSock(s)
def log(self, s):
self.logf.write("Server log %s: %s\n" % (self.log_prefix, s))
@@ -130,13 +131,13 @@ class Server:
(fname, args) = RestrictedUnpickler(self.fs).load()
except EOFError:
raise ClientExit("read args")
self.log("executing method %s"%(fname))
self.log("executing method %s" % (fname))
st = None
ret = None
try:
f=getattr(self,fname)
f = getattr(self, fname)
except AttributeError:
st = AttributeError("unknown method "+fname)
st = AttributeError("unknown method " + fname)
self.log("unknown method")
try:
@@ -145,38 +146,38 @@ class Server:
# due to a bug (in mod_python?), ServerException cannot be
# unpickled, so send the string and make the exception on the client side
#st=ServerException(
# st=ServerException(
# "".join(traceback.format_tb(sys.exc_info()[2]))+
# str(e))
st="".join(traceback.format_tb(sys.exc_info()[2]))+str(e)
st = "".join(traceback.format_tb(sys.exc_info()[2])) + str(e)
self.log("exception in method")
traceback.print_exc(50,self.logf)
traceback.print_exc(50, self.logf)
self.logf.flush()
LOG.info("return")
try:
pickle.dump((st ,ret), self.fs, protocol=4)
pickle.dump((st, ret), self.fs, protocol=4)
except EOFError:
raise ClientExit("function return")
def exec_loop(self):
""" main execution loop. Loops and handles exit states"""
"""main execution loop. Loops and handles exit states"""
self.log("in exec_loop")
try:
while True:
self.one_function()
except ClientExit as e:
self.log("ClientExit %s"%e)
self.log("ClientExit %s" % e)
except socket.error as e:
self.log("socket error %s"%e)
traceback.print_exc(50,self.logf)
self.log("socket error %s" % e)
traceback.print_exc(50, self.logf)
except EOFError:
self.log("EOF during communication")
traceback.print_exc(50,self.logf)
traceback.print_exc(50, self.logf)
except BaseException:
# unexpected
traceback.print_exc(50,sys.stderr)
traceback.print_exc(50, sys.stderr)
sys.exit(1)
LOG.info("exit server")
@@ -188,14 +189,16 @@ class Server:
# spying stuff
def get_ps_stats(self):
ret=''
f=os.popen("echo ============ `hostname` uptime:; uptime;"+
"echo ============ self:; "+
"ps -p %d -o pid,vsize,rss,%%cpu,nlwp,psr; "%os.getpid()+
"echo ============ run queue:;"+
"ps ar -o user,pid,%cpu,%mem,ni,nlwp,psr,vsz,rss,cputime,command")
ret = ""
f = os.popen(
"echo ============ `hostname` uptime:; uptime;"
+ "echo ============ self:; "
+ "ps -p %d -o pid,vsize,rss,%%cpu,nlwp,psr; " % os.getpid()
+ "echo ============ run queue:;"
+ "ps ar -o user,pid,%cpu,%mem,ni,nlwp,psr,vsz,rss,cputime,command"
)
for l in f:
ret+=l
ret += l
return ret
@@ -204,6 +207,7 @@ class Client:
Methods of the server object can be called transparently. Exceptions are
re-raised.
"""
def __init__(self, HOST, port=PORT, v6=False):
socktype = socket.AF_INET6 if v6 else socket.AF_INET
@@ -220,18 +224,18 @@ class Client:
def get_result(self):
(st, ret) = RestrictedUnpickler(self.fs).load()
if st!=None:
if st != None:
raise ServerException(st)
else:
return ret
def __getattr__(self,name):
return lambda *x: self.generic_fun(name,x)
def __getattr__(self, name):
return lambda *x: self.generic_fun(name, x)
def run_server(new_handler, port=PORT, report_to_file=None, v6=False):
HOST = '' # Symbolic name meaning the local host
HOST = "" # Symbolic name meaning the local host
socktype = socket.AF_INET6 if v6 else socket.AF_INET
s = socket.socket(socktype, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -242,20 +246,21 @@ def run_server(new_handler, port=PORT, report_to_file=None, v6=False):
LOG.info("accepting connections")
if report_to_file is not None:
LOG.info('storing host+port in %s', report_to_file)
open(report_to_file, 'w').write('%s:%d ' % (socket.gethostname(), port))
LOG.info("storing host+port in %s", report_to_file)
open(report_to_file, "w").write("%s:%d " % (socket.gethostname(), port))
while True:
try:
conn, addr = s.accept()
except socket.error as e:
if e[1]=='Interrupted system call': continue
if e[1] == "Interrupted system call":
continue
raise
LOG.info('Connected to %s', addr)
LOG.info("Connected to %s", addr)
ibs = new_handler(conn)
tid = _thread.start_new_thread(ibs.exec_loop,())
tid = _thread.start_new_thread(ibs.exec_loop, ())
LOG.debug("Thread ID: %d", tid)
contrib/torch/clustering.py
-1
View File
@@ -14,7 +14,6 @@ import torch
from faiss.contrib.clustering import kmeans # noqa: F401 some libraries import kmeans from here
class DatasetAssign:
"""Wrapper for a tensor that offers a function to assign the vectors
to centroids. All other implementations offer the same interface"""
contrib/torch/quantization.py
+6 -5
View File
@@ -11,6 +11,7 @@ import torch
import faiss
import math
from faiss.contrib.torch import clustering
# the kmeans can produce both torch and numpy centroids
@@ -57,7 +58,7 @@ class VectorQuantizer(Quantizer):
class ProductQuantizer(Quantizer):
def __init__(self, d, M, nbits):
""" M: number of subvectors, d%M == 0
"""M: number of subvectors, d%M == 0
nbits: number of bits that each vector is encoded into
"""
assert d % M == 0
@@ -69,20 +70,20 @@ class ProductQuantizer(Quantizer):
self.code_size = code_size
def train(self, x):
nc = 2 ** self.nbits
nc = 2**self.nbits
sd = self.d // self.M
dev = x.device
dtype = x.dtype
self.codebook = torch.zeros((self.M, nc, sd), device=dev, dtype=dtype)
for m in range(self.M):
xsub = x[:, m * self.d // self.M: (m + 1) * self.d // self.M]
xsub = x[:, m * self.d // self.M : (m + 1) * self.d // self.M]
data = clustering.DatasetAssign(xsub.contiguous())
self.codebook[m] = clustering.kmeans(2 ** self.nbits, data)
self.codebook[m] = clustering.kmeans(2**self.nbits, data)
def encode(self, x):
codes = torch.zeros((x.shape[0], self.code_size), dtype=torch.uint8)
for m in range(self.M):
xsub = x[:, m * self.d // self.M:(m + 1) * self.d // self.M]
xsub = x[:, m * self.d // self.M : (m + 1) * self.d // self.M]
_, I = faiss.knn(xsub.contiguous(), self.codebook[m], 1)
codes[:, m] = I.ravel()
return codes
contrib/torch_utils.py
+187 -138
View File
@@ -34,52 +34,59 @@ import numpy as np
def swig_ptr_from_UInt8Tensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.uint8
return faiss.cast_integer_to_uint8_ptr(
x.untyped_storage().data_ptr() + x.storage_offset())
x.untyped_storage().data_ptr() + x.storage_offset()
)
def swig_ptr_from_HalfTensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.float16
# no canonical half type in C/C++
return faiss.cast_integer_to_void_ptr(
x.untyped_storage().data_ptr() + x.storage_offset() * 2)
x.untyped_storage().data_ptr() + x.storage_offset() * 2
)
def swig_ptr_from_FloatTensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.float32
return faiss.cast_integer_to_float_ptr(
x.untyped_storage().data_ptr() + x.storage_offset() * 4)
x.untyped_storage().data_ptr() + x.storage_offset() * 4
)
def swig_ptr_from_BFloat16Tensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.bfloat16
return faiss.cast_integer_to_void_ptr(
x.untyped_storage().data_ptr() + x.storage_offset() * 2)
x.untyped_storage().data_ptr() + x.storage_offset() * 2
)
def swig_ptr_from_IntTensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.int32, 'dtype=%s' % x.dtype
assert x.dtype == torch.int32, "dtype=%s" % x.dtype
return faiss.cast_integer_to_int_ptr(
x.untyped_storage().data_ptr() + x.storage_offset() * 4)
x.untyped_storage().data_ptr() + x.storage_offset() * 4
)
def swig_ptr_from_IndicesTensor(x):
""" gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU) """
"""gets a Faiss SWIG pointer from a pytorch tensor (on CPU or GPU)"""
assert x.is_contiguous()
assert x.dtype == torch.int64, 'dtype=%s' % x.dtype
assert x.dtype == torch.int64, "dtype=%s" % x.dtype
return faiss.cast_integer_to_idx_t_ptr(
x.untyped_storage().data_ptr() + x.storage_offset() * 8)
x.untyped_storage().data_ptr() + x.storage_offset() * 8
)
##################################################################
# utilities
@@ -88,10 +95,10 @@ def swig_ptr_from_IndicesTensor(x):
@contextlib.contextmanager
def using_stream(res, pytorch_stream=None):
""" Creates a scoping object to make Faiss GPU use the same stream
as pytorch, based on torch.cuda.current_stream().
Or, a specific pytorch stream can be passed in as a second
argument, in which case we will use that stream.
"""Creates a scoping object to make Faiss GPU use the same stream
as pytorch, based on torch.cuda.current_stream().
Or, a specific pytorch stream can be passed in as a second
argument, in which case we will use that stream.
"""
if pytorch_stream is None:
@@ -99,7 +106,8 @@ def using_stream(res, pytorch_stream=None):
# This is the cudaStream_t that we wish to use
cuda_stream_s = faiss.cast_integer_to_cudastream_t(
pytorch_stream.cuda_stream)
pytorch_stream.cuda_stream
)
# So we can revert GpuResources stream state upon exit
prior_dev = torch.cuda.current_device()
@@ -114,23 +122,25 @@ def using_stream(res, pytorch_stream=None):
res.setDefaultStream(prior_dev, prior_stream)
def torch_replace_method(the_class, name, replacement,
ignore_missing=False, ignore_no_base=False):
def torch_replace_method(
the_class, name, replacement, ignore_missing=False, ignore_no_base=False
):
try:
orig_method = getattr(the_class, name)
except AttributeError:
if ignore_missing:
return
raise
if orig_method.__name__ == 'torch_replacement_' + name:
if orig_method.__name__ == "torch_replacement_" + name:
# replacement was done in parent class
return
# We should already have the numpy replacement methods patched
assert ignore_no_base or (orig_method.__name__ == 'replacement_' + name)
setattr(the_class, name + '_numpy', orig_method)
assert ignore_no_base or (orig_method.__name__ == "replacement_" + name)
setattr(the_class, name + "_numpy", orig_method)
setattr(the_class, name, replacement)
##################################################################
# Setup wrappers
##################################################################
@@ -155,8 +165,9 @@ def handle_torch_Index(the_class):
)
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()):
@@ -165,8 +176,9 @@ 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)
@@ -184,12 +196,13 @@ def handle_torch_Index(the_class):
)
assert type(ids) is torch.Tensor
assert ids.shape == (n, ), 'not same number of vectors as ids'
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()):
@@ -216,8 +229,9 @@ 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()):
@@ -246,8 +260,9 @@ def handle_torch_Index(the_class):
)
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()):
@@ -283,8 +298,9 @@ 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)
@@ -296,8 +312,9 @@ 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()):
@@ -308,8 +325,9 @@ 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)
@@ -328,21 +346,22 @@ 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)
@@ -365,24 +384,28 @@ 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):
@@ -395,21 +418,22 @@ def handle_torch_Index(the_class):
# If the index is a CPU index, the default device is CPU, otherwise we
# produce a GPU tensor
device = torch.device('cpu')
if hasattr(self, 'getDevice'):
device = torch.device("cpu")
if hasattr(self, "getDevice"):
# same device as the index
device = torch.device('cuda', self.getDevice())
device = torch.device("cuda", self.getDevice())
if x is None:
x = torch.empty(self.d, device=device, dtype=torch.float32)
else:
assert type(x) is torch.Tensor
assert x.shape == (self.d, )
assert x.shape == (self.d,)
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()):
@@ -433,10 +457,10 @@ def handle_torch_Index(the_class):
# If the index is a CPU index, the default device is CPU, otherwise we
# produce a GPU tensor
device = torch.device('cpu')
if hasattr(self, 'getDevice'):
device = torch.device("cpu")
if hasattr(self, "getDevice"):
# same device as the index
device = torch.device('cuda', self.getDevice())
device = torch.device("cuda", self.getDevice())
if x is None:
x = torch.empty(ni, self.d, device=device, dtype=torch.float32)
@@ -446,8 +470,9 @@ 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()):
@@ -464,7 +489,7 @@ def handle_torch_Index(the_class):
return self.update_vectors_numpy(keys, x)
assert type(keys) is torch.Tensor
(n, ) = keys.shape
(n,) = keys.shape
keys_ptr = swig_ptr_from_IndicesTensor(keys)
assert type(x) is torch.Tensor
@@ -472,8 +497,9 @@ 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()):
@@ -494,10 +520,12 @@ 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)
@@ -507,7 +535,8 @@ def handle_torch_Index(the_class):
# 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'))
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())
@@ -531,8 +560,9 @@ 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()):
@@ -561,8 +591,9 @@ 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()):
@@ -573,34 +604,46 @@ 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, '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,
ignore_missing=True)
torch_replace_method(the_class, 'search_and_reconstruct',
torch_replacement_search_and_reconstruct,
ignore_missing=True)
torch_replace_method(the_class, 'search_preassigned',
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)
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, "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,
ignore_missing=True,
)
torch_replace_method(
the_class,
"search_and_reconstruct",
torch_replacement_search_and_reconstruct,
ignore_missing=True,
)
torch_replace_method(
the_class,
"search_preassigned",
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']
faiss_module = sys.modules["faiss"]
# Re-patch anything that inherits from faiss.Index to add the torch bindings
for symbol in dir(faiss_module):
@@ -615,8 +658,7 @@ 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()
@@ -637,31 +679,32 @@ def torch_replacement_knn(xq, xb, k, metric=faiss.METRIC_L2, metric_arg=0):
xq_ptr = swig_ptr_from_FloatTensor(xq)
if metric == faiss.METRIC_L2:
faiss.knn_L2sqr(
xq_ptr, xb_ptr,
d, nq, nb, k, D_ptr, I_ptr
)
faiss.knn_L2sqr(xq_ptr, xb_ptr, d, nq, nb, k, D_ptr, I_ptr)
elif metric == faiss.METRIC_INNER_PRODUCT:
faiss.knn_inner_product(
xq_ptr, xb_ptr,
d, nq, nb, k, D_ptr, I_ptr
)
faiss.knn_inner_product(xq_ptr, xb_ptr, d, nq, nb, k, D_ptr, I_ptr)
else:
faiss.knn_extra_metrics(
xq_ptr, xb_ptr,
d, nq, nb, metric, metric_arg, k, D_ptr, I_ptr
xq_ptr, xb_ptr, d, nq, nb, metric, metric_arg, k, D_ptr, I_ptr
)
return D, I
torch_replace_method(faiss_module, 'knn', torch_replacement_knn, True, True)
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)
@@ -673,7 +716,7 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
xb = xb.t()
xb_row_major = False
else:
raise TypeError('matrix should be row or column-major')
raise TypeError("matrix should be row or column-major")
if xb.dtype == torch.float32:
xb_type = faiss.DistanceDataType_F32
@@ -685,7 +728,7 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
xb_type = faiss.DistanceDataType_BF16
xb_ptr = swig_ptr_from_BFloat16Tensor(xb)
else:
raise TypeError('xq must be float32, float16 or bfloat16')
raise TypeError("xq must be float32, float16 or bfloat16")
nq, d2 = xq.size()
assert d2 == d
@@ -695,7 +738,7 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
xq = xq.t()
xq_row_major = False
else:
raise TypeError('matrix should be row or column-major')
raise TypeError("matrix should be row or column-major")
if xq.dtype == torch.float32:
xq_type = faiss.DistanceDataType_F32
@@ -707,14 +750,14 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
xq_type = faiss.DistanceDataType_BF16
xq_ptr = swig_ptr_from_BFloat16Tensor(xq)
else:
raise TypeError('xq must be float32, float16 or bfloat16')
raise TypeError("xq must be float32, float16 or bfloat16")
if D is None:
D = torch.empty(nq, k, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, k)
# interface takes void*, we need to check this
assert (D.dtype == torch.float32)
assert D.dtype == torch.float32
if I is None:
I = torch.empty(nq, k, device=xb.device, dtype=torch.int64)
@@ -728,7 +771,7 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
I_type = faiss.IndicesDataType_I32
I_ptr = swig_ptr_from_IntTensor(I)
else:
raise TypeError('I must be i64 or i32')
raise TypeError("I must be i64 or i32")
D_ptr = swig_ptr_from_FloatTensor(D)
@@ -756,14 +799,15 @@ def torch_replacement_knn_gpu(res, xq, xb, k, D=None, I=None,
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)
@@ -775,7 +819,7 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
xb = xb.t()
xb_row_major = False
else:
raise TypeError('xb matrix should be row or column-major')
raise TypeError("xb matrix should be row or column-major")
if xb.dtype == torch.float32:
xb_type = faiss.DistanceDataType_F32
@@ -784,7 +828,7 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
xb_type = faiss.DistanceDataType_F16
xb_ptr = swig_ptr_from_HalfTensor(xb)
else:
raise TypeError('xb must be float32 or float16')
raise TypeError("xb must be float32 or float16")
nq, d2 = xq.size()
assert d2 == d
@@ -794,7 +838,7 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
xq = xq.t()
xq_row_major = False
else:
raise TypeError('xq matrix should be row or column-major')
raise TypeError("xq matrix should be row or column-major")
if xq.dtype == torch.float32:
xq_type = faiss.DistanceDataType_F32
@@ -803,14 +847,14 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
xq_type = faiss.DistanceDataType_F16
xq_ptr = swig_ptr_from_HalfTensor(xq)
else:
raise TypeError('xq must be float32 or float16')
raise TypeError("xq must be float32 or float16")
if D is None:
D = torch.empty(nq, nb, device=xb.device, dtype=torch.float32)
else:
assert D.shape == (nq, nb)
# interface takes void*, we need to check this
assert (D.dtype == torch.float32)
assert D.dtype == torch.float32
D_ptr = swig_ptr_from_FloatTensor(D)
@@ -835,5 +879,10 @@ def torch_replacement_pairwise_distance_gpu(res, xq, xb, D=None,
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,
)
contrib/vecs_io.py
+29 -23
View File
@@ -14,51 +14,52 @@ definition of the formats here: http://corpus-texmex.irisa.fr/
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
if sys.byteorder == 'big':
a = np.fromfile(fname, dtype="int32")
if sys.byteorder == "big":
a.byteswap(inplace=True)
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
def ivecs_mmap(fname):
assert sys.byteorder != 'big'
a = np.memmap(fname, dtype='int32', mode='r')
assert sys.byteorder != "big"
a = np.memmap(fname, dtype="int32", mode="r")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:]
def fvecs_mmap(fname):
return ivecs_mmap(fname).view('float32')
return ivecs_mmap(fname).view("float32")
def bvecs_mmap(fname):
x = np.memmap(fname, dtype='uint8', mode='r')
if sys.byteorder == 'big':
x = np.memmap(fname, dtype="uint8", mode="r")
if sys.byteorder == "big":
da = x[:4][::-1].copy()
d = da.view('int32')[0]
d = da.view("int32")[0]
else:
d = x[:4].view('int32')[0]
d = x[:4].view("int32")[0]
return x.reshape(-1, d + 4)[:, 4:]
def ivecs_write(fname, m):
n, d = m.shape
m1 = np.empty((n, d + 1), dtype='int32')
m1 = np.empty((n, d + 1), dtype="int32")
m1[:, 0] = d
m1[:, 1:] = m
if sys.byteorder == 'big':
if sys.byteorder == "big":
m1.byteswap(inplace=True)
m1.tofile(fname)
def fvecs_write(fname, m):
m = m.astype('float32')
ivecs_write(fname, m.view('int32'))
m = m.astype("float32")
ivecs_write(fname, m.view("int32"))
def bvecs_iter(filepath, batch_size=100_000):
"""
@@ -66,19 +67,20 @@ def bvecs_iter(filepath, batch_size=100_000):
"""
file_size = os.path.getsize(filepath)
with open(filepath, 'rb') as f:
dim = np.frombuffer(f.read(4), dtype='<i4')[0]
with open(filepath, "rb") as f:
dim = np.frombuffer(f.read(4), dtype="<i4")[0]
bytes_per_vec = 4 + dim
n_vectors = file_size // bytes_per_vec
mm = np.memmap(filepath, mode='r', dtype=np.uint8)
mm = np.memmap(filepath, mode="r", dtype=np.uint8)
records = mm.reshape(n_vectors, bytes_per_vec)
for start in range(0, n_vectors, batch_size):
end = np.min([start + batch_size, n_vectors])
yield records[start:end, 4:]
def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
"""
Memory-mapped iterator over chunked .bvecs files.
@@ -99,7 +101,11 @@ def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
# Find all chunk files and sort them
chunk_files = []
for entry in os.scandir(chunk_folder):
if entry.is_file() and entry.name.startswith("chunk_") and entry.name.endswith(".bvecs"):
if (
entry.is_file()
and entry.name.startswith("chunk_")
and entry.name.endswith(".bvecs")
):
chunk_files.append(entry.path)
chunk_files.sort()
@@ -111,7 +117,7 @@ def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
for path in chunk_files:
basename = os.path.basename(path)
try:
num_str = basename.split('_')[1].split('.')[0]
num_str = basename.split("_")[1].split(".")[0]
chunk_numbers.append(int(num_str))
except (IndexError, ValueError):
raise ValueError(f"Invalid chunk filename format: {basename}")
@@ -127,8 +133,8 @@ def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
)
# Get dimension from first chunk
with open(chunk_files[0], 'rb') as f:
dim = np.frombuffer(f.read(4), dtype='<i4')[0]
with open(chunk_files[0], "rb") as f:
dim = np.frombuffer(f.read(4), dtype="<i4")[0]
bytes_per_vec = 4 + dim
@@ -142,7 +148,7 @@ def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
n_vectors = file_size // bytes_per_vec
# Memory-map this chunk
mm = np.memmap(chunk_path, mode='r', dtype=np.uint8)
mm = np.memmap(chunk_path, mode="r", dtype=np.uint8)
records = mm.reshape(n_vectors, bytes_per_vec)
vectors = records[:, 4:] # Skip dimension prefix
@@ -166,7 +172,7 @@ def bvecs_iter_chunked(chunk_folder, batch_size=100_000):
# Now process complete batches from this chunk
while start + batch_size <= n_vectors:
yield vectors[start:start + batch_size]
yield vectors[start : start + batch_size]
start += batch_size
remainder = n_vectors - start
demos/demo_auto_tune.py
+44 -26
View File
@@ -10,8 +10,10 @@ import numpy as np
try:
import matplotlib
matplotlib.use('Agg')
matplotlib.use("Agg")
from matplotlib import pyplot
graphical_output = True
except ImportError:
graphical_output = False
@@ -30,15 +32,17 @@ def ivecs_read(fname):
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
def plot_OperatingPoints(ops, nq, **kwargs):
ops = ops.optimal_pts
n = ops.size() * 2 - 1
pyplot.plot([ops.at( i // 2).perf for i in range(n)],
[ops.at((i + 1) // 2).t / nq * 1000 for i in range(n)],
**kwargs)
pyplot.plot(
[ops.at(i // 2).perf for i in range(n)],
[ops.at((i + 1) // 2).t / nq * 1000 for i in range(n)],
**kwargs
)
#################################################################
@@ -46,7 +50,6 @@ def plot_OperatingPoints(ops, nq, **kwargs):
#################################################################
t0 = time.time()
print("load data")
@@ -60,7 +63,7 @@ d = xt.shape[1]
print("load GT")
gt = ivecs_read("sift1M/sift_groundtruth.ivecs")
gt = gt.astype('int64')
gt = gt.astype("int64")
k = gt.shape[1]
print("prepare criterion")
@@ -72,28 +75,41 @@ crit.nnn = k
# indexes that are useful when there is no limitation on memory usage
unlimited_mem_keys = [
"IMI2x10,Flat", "IMI2x11,Flat",
"IVF4096,Flat", "IVF16384,Flat",
"PCA64,IMI2x10,Flat"]
"IMI2x10,Flat",
"IMI2x11,Flat",
"IVF4096,Flat",
"IVF16384,Flat",
"PCA64,IMI2x10,Flat",
]
# memory limited to 16 bytes / vector
keys_mem_16 = [
'IMI2x10,PQ16', 'IVF4096,PQ16',
'IMI2x10,PQ8+8', 'OPQ16_64,IMI2x10,PQ16'
]
"IMI2x10,PQ16",
"IVF4096,PQ16",
"IMI2x10,PQ8+8",
"OPQ16_64,IMI2x10,PQ16",
]
# limited to 32 bytes / vector
keys_mem_32 = [
'IMI2x10,PQ32', 'IVF4096,PQ32', 'IVF16384,PQ32',
'IMI2x10,PQ16+16',
'OPQ32,IVF4096,PQ32', 'IVF4096,PQ16+16', 'OPQ16,IMI2x10,PQ16+16'
]
"IMI2x10,PQ32",
"IVF4096,PQ32",
"IVF16384,PQ32",
"IMI2x10,PQ16+16",
"OPQ32,IVF4096,PQ32",
"IVF4096,PQ16+16",
"OPQ16,IMI2x10,PQ16+16",
]
# indexes that can run on the GPU
keys_gpu = [
"PCA64,IVF4096,Flat",
"PCA64,Flat", "Flat", "IVF4096,Flat", "IVF16384,Flat",
"IVF4096,PQ32"]
"PCA64,Flat",
"Flat",
"IVF4096,Flat",
"IVF16384,Flat",
"IVF4096,PQ32",
]
keys_to_test = unlimited_mem_keys
@@ -102,8 +118,9 @@ use_gpu = False
if use_gpu:
# if this fails, it means that the GPU version was not compiled
assert faiss.StandardGpuResources, \
"Faiss was not compiled with GPU support, or loading _swigfaiss_gpu.so failed"
assert (
faiss.StandardGpuResources
), "Faiss was not compiled with GPU support, or loading _swigfaiss_gpu.so failed"
res = faiss.StandardGpuResources()
dev_no = 0
@@ -122,7 +139,6 @@ for index_key in keys_to_test:
# make the index described by the key
index = faiss.index_factory(d, index_key)
if use_gpu:
# transfer to GPU (may be partial)
index = faiss.index_cpu_to_gpu(res, dev_no, index)
@@ -155,14 +171,16 @@ for index_key in keys_to_test:
fig = pyplot.figure(figsize=(12, 9))
pyplot.xlabel("1-recall at 1")
pyplot.ylabel("search time (ms/query, %d threads)" % faiss.omp_get_max_threads())
pyplot.gca().set_yscale('log')
pyplot.ylabel(
"search time (ms/query, %d threads)" % faiss.omp_get_max_threads()
)
pyplot.gca().set_yscale("log")
pyplot.grid()
for i2, opi2 in op_per_key:
plot_OperatingPoints(opi2, crit.nq, label = i2, marker = 'o')
plot_OperatingPoints(opi2, crit.nq, label=i2, marker="o")
# plot_OperatingPoints(op, crit.nq, label = 'best', marker = 'o', color = 'r')
pyplot.legend(loc=2)
fig.savefig('tmp/demo_auto_tune.png')
fig.savefig("tmp/demo_auto_tune.png")
print("[%.3f s] final result:" % (time.time() - t0))
demos/demo_client_server_ivf.py
+8 -8
View File
@@ -17,13 +17,13 @@ from faiss.contrib.client_server import run_index_server, ClientIndex
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
a = np.fromfile(fname, dtype="int32")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
#################################################################
@@ -32,7 +32,7 @@ def fvecs_read(fname):
stage = int(sys.argv[1])
tmpdir = '/tmp/'
tmpdir = "/tmp/"
if stage == 0:
# train the index
@@ -57,10 +57,10 @@ if 1 <= stage <= 4:
machine_ports = [
('localhost', 12010),
('localhost', 12011),
('localhost', 12012),
('localhost', 12013),
("localhost", 12010),
("localhost", 12011),
("localhost", 12012),
("localhost", 12013),
]
v6 = False
@@ -78,7 +78,7 @@ if 5 <= stage <= 8:
if stage == 9:
client_index = ClientIndex(machine_ports)
print('index size:', client_index.ntotal)
print("index size:", client_index.ntotal)
client_index.set_nprobe(16)
# load query vectors and ground-truth
demos/demo_distributed_kmeans_torch.py
+18 -11
View File
@@ -33,18 +33,19 @@ class DatasetAssignDistributedGPU(clustering.DatasetAssign):
sizes = torch.zeros(nproc, device=self.device, dtype=torch.int64)
sizes[rank] = n
torch.distributed.all_gather(
[sizes[i:i + 1] for i in range(nproc)], sizes[rank:rank + 1])
[sizes[i : i + 1] for i in range(nproc)], sizes[rank : rank + 1]
)
self.sizes = sizes.cpu().numpy()
# begin & end of each shard
self.cs = np.zeros(nproc + 1, dtype='int64')
self.cs = np.zeros(nproc + 1, dtype="int64")
self.cs[1:] = np.cumsum(self.sizes)
def count(self):
return int(self.sizes.sum())
def int_to_slaves(self, i):
" broadcast an int to all workers "
"broadcast an int to all workers"
rank = self.rank
tab = torch.zeros(1, device=self.device, dtype=torch.int64)
if rank == 0:
@@ -64,7 +65,8 @@ class DatasetAssignDistributedGPU(clustering.DatasetAssign):
indices = torch.from_numpy(indices).to(self.device)
else:
indices = torch.zeros(
len_indices, dtype=torch.int64, device=self.device)
len_indices, dtype=torch.int64, device=self.device
)
torch.distributed.broadcast(indices, 0)
# select subset of indices
@@ -73,8 +75,8 @@ class DatasetAssignDistributedGPU(clustering.DatasetAssign):
mask = torch.logical_and(indices < i1, indices >= i0)
output = torch.zeros(
len_indices, self.x.shape[1],
dtype=self.x.dtype, device=self.device)
len_indices, self.x.shape[1], dtype=self.x.dtype, device=self.device
)
output[mask] = self.x[indices[mask] - i0]
torch.distributed.reduce(output, 0) # sum
if rank == 0:
@@ -94,12 +96,14 @@ class DatasetAssignDistributedGPU(clustering.DatasetAssign):
if rank != 0:
centroids = torch.zeros(
nc, self.x.shape[1], dtype=self.x.dtype, device=self.device)
nc, self.x.shape[1], dtype=self.x.dtype, device=self.device
)
torch.distributed.broadcast(centroids, 0)
# perform search
D, I = faiss.knn_gpu(
self.res, self.x, centroids, 1, device=self.device.index)
self.res, self.x, centroids, 1, device=self.device.index
)
I = I.ravel()
D = D.ravel()
@@ -120,11 +124,13 @@ class DatasetAssignDistributedGPU(clustering.DatasetAssign):
all_I = torch.zeros(self.count(), dtype=I.dtype, device=device)
all_D = torch.zeros(self.count(), dtype=D.dtype, device=device)
torch.distributed.gather(
I, [all_I[self.cs[r]:self.cs[r + 1]] for r in range(nproc)],
I,
[all_I[self.cs[r] : self.cs[r + 1]] for r in range(nproc)],
dst=0,
)
torch.distributed.gather(
D, [all_D[self.cs[r]:self.cs[r + 1]] for r in range(nproc)],
D,
[all_D[self.cs[r] : self.cs[r + 1]] for r in range(nproc)],
dst=0,
)
return all_I.cpu().numpy(), all_D, sum_per_centroid
@@ -160,7 +166,8 @@ if __name__ == "__main__":
if rank == 0:
print(f"sizes = {da.sizes}")
centroids, iteration_stats = clustering.kmeans(
k, da, niter=niter, return_stats=True)
k, da, niter=niter, return_stats=True
)
print("clusters:", centroids.cpu().numpy())
else:
# make sure the iterations are aligned with master
demos/demo_ondisk_ivf.py
+5 -8
View File
@@ -15,13 +15,13 @@ from faiss.contrib.ondisk import merge_ondisk
def ivecs_read(fname):
a = np.fromfile(fname, dtype='int32')
a = np.fromfile(fname, dtype="int32")
d = a[0]
return a.reshape(-1, d + 1)[:, 1:].copy()
def fvecs_read(fname):
return ivecs_read(fname).view('float32')
return ivecs_read(fname).view("float32")
#################################################################
@@ -30,7 +30,7 @@ def fvecs_read(fname):
stage = int(sys.argv[1])
tmpdir = '/tmp/'
tmpdir = "/tmp/"
if stage == 0:
# train the index
@@ -55,14 +55,11 @@ if 1 <= stage <= 4:
if stage == 5:
print('loading trained index')
print("loading trained index")
# construct the output index
index = faiss.read_index(tmpdir + "trained.index")
block_fnames = [
tmpdir + "block_%d.index" % bno
for bno in range(4)
]
block_fnames = [tmpdir + "block_%d.index" % bno for bno in range(4)]
merge_ondisk(index, block_fnames, tmpdir + "merged_index.ivfdata")
demos/demo_qinco.py
+2 -2
View File
@@ -5,7 +5,7 @@
"""
This demonstrates how to reproduce the QINCo paper results using the Faiss
QINCo implementation. The code loads the reference model because training
QINCo implementation. The code loads the reference model because training
is not implemented in Faiss.
Prepare the data with
@@ -41,7 +41,7 @@ with torch.no_grad():
torch.set_num_threads(1)
faiss.omp_set_num_threads(1)
x_base = bvecs_mmap("/tmp/bigann1M.bvecs")[:1000].astype('float32')
x_base = bvecs_mmap("/tmp/bigann1M.bvecs")[:1000].astype("float32")
x_scaled = torch.from_numpy(x_base) / qinco.db_scale
t0 = time.time()
demos/offline_ivf/create_sharded_ssnpp_files.py
+1 -1
View File
@@ -29,7 +29,7 @@ def main(args: argparse.Namespace):
), "num of embeddings per file should divide total num of embeddings"
for i in range(num_batches):
xb_batch = ssnpp_data[
i * args.data_batch:(i + 1) * args.data_batch, :
i * args.data_batch : (i + 1) * args.data_batch, :
]
filename = args.output_dir + f"/ssnpp_{(i):010}.npy"
np.save(filename, xb_batch)
demos/offline_ivf/dataset.py
+1 -1
View File
@@ -111,7 +111,7 @@ class MultiFileVectorDataset:
xb = xb[start:]
start = 0
req = min(batch_size - rem, xb.shape[0])
buffer[rem:rem + req] = xb[:req]
buffer[rem : rem + req] = xb[:req]
rem += req
if rem == batch_size:
if self.normalize:
demos/offline_ivf/offline_ivf.py
+18 -13
View File
@@ -179,7 +179,9 @@ class OfflineIVF:
idxs.append(np.empty((0,), dtype=np.uint32))
bs = 1_000_000
i = 0
for buffer in tqdm(self._iterate_transformed(self.xb_ds, 0, bs, np.float32)):
for buffer in tqdm(
self._iterate_transformed(self.xb_ds, 0, bs, np.float32)
):
for j in range(len(codecs)):
codec, codeset, idx = codecs[j], codesets[j], idxs[j]
uniq = codeset.insert(codec.sa_encode(buffer))
@@ -272,13 +274,17 @@ class OfflineIVF:
assert xb_j.shape[1] == index.chain.at(0).d_out
index_ivf.add_with_ids(
xb_j,
np.arange(start + jj, start + jj + xb_j.shape[0]),
np.arange(
start + jj, start + jj + xb_j.shape[0]
),
)
else:
assert xb_j.shape[1] == index.d
index.add_with_ids(
xb_j,
np.arange(start + jj, start + jj + xb_j.shape[0]),
np.arange(
start + jj, start + jj + xb_j.shape[0]
),
)
jj += xb_j.shape[0]
logging.info(jj)
@@ -413,7 +419,9 @@ class OfflineIVF:
D_a_ann_file = (
f"{self.eval_dir}/D_a_ann_{index_factory_fn}_np{nprobe}.npy"
)
D_a_ann_refined_file = f"{self.eval_dir}/D_a_ann_refined_{index_factory_fn}_np{nprobe}.npy"
D_a_ann_refined_file = (
f"{self.eval_dir}/D_a_ann_refined_{index_factory_fn}_np{nprobe}.npy"
)
D_b_gt_file = f"{self.eval_dir}/D_b_gt.npy"
D_b_ann_file = (
f"{self.eval_dir}/D_b_ann_{index_factory_fn}_np{nprobe}.npy"
@@ -684,26 +692,22 @@ class OfflineIVF:
else:
d = self.input_d
with open(Ifn, "xb") as f, open(Dfn, "xb") as g:
xq_i = np.empty(
shape=(self.xq_bs, d), dtype=np.float16
)
xq_i = np.empty(shape=(self.xq_bs, d), dtype=np.float16)
q_assign = np.empty(
(self.xq_bs, self.nprobe), dtype=np.int32
)
j = 0
quantizer = faiss.index_cpu_to_all_gpus(
index_ivf.quantizer
)
quantizer = faiss.index_cpu_to_all_gpus(index_ivf.quantizer)
for xq_i_j in tqdm(
self._iterate_transformed(
self.xq_ds, i, min(100_000, self.xq_bs), np.float16
),
file=sys.stdout,
):
xq_i[j:j + xq_i_j.shape[0]] = xq_i_j
xq_i[j : j + xq_i_j.shape[0]] = xq_i_j
(
_,
q_assign[j:j + xq_i_j.shape[0]],
q_assign[j : j + xq_i_j.shape[0]],
) = quantizer.search(xq_i_j, self.nprobe)
j += xq_i_j.shape[0]
assert j <= xq_i.shape[0]
@@ -847,7 +851,8 @@ class OfflineIVF:
for j in range(SMALL_DATA_SAMPLE):
assert np.where(I[j] == j + r)[0].size > 0, (
f"I[j]: {I[j]}, j: {j}, i: {i}, shard_size:"
f" {self.shard_size}")
f" {self.shard_size}"
)
logging.info("search results...")
index_ivf.nprobe = self.nprobe
demos/offline_ivf/tests/test_iterate_input.py
+6 -6
View File
@@ -38,9 +38,9 @@ class TestUtilsMethods(unittest.TestCase):
data_creator.create_test_data()
args = data_creator.setup_cli()
cfg = load_config(args.config)
db_iterator = create_dataset_from_oivf_config(
cfg, args.xb
).iterate(0, TEST_BATCH_SIZE, np.float32)
db_iterator = create_dataset_from_oivf_config(cfg, args.xb).iterate(
0, TEST_BATCH_SIZE, np.float32
)
for i in range(len(SMALL_FILE_SIZES) - 1):
vecs = next(db_iterator)
@@ -65,9 +65,9 @@ class TestUtilsMethods(unittest.TestCase):
data_creator.create_test_data()
args = data_creator.setup_cli()
cfg = load_config(args.config)
db_iterator = create_dataset_from_oivf_config(
cfg, args.xb
).iterate(0, TEST_BATCH_SIZE, np.float32)
db_iterator = create_dataset_from_oivf_config(cfg, args.xb).iterate(
0, TEST_BATCH_SIZE, np.float32
)
for i in range(len(LARGE_FILE_SIZES) - 1):
vecs = next(db_iterator)
faiss/gpu/perf/slow.py
+8 -7
View File
@@ -11,15 +11,16 @@ import numpy as np
def test_slow():
d = 256
index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(),
0, faiss.IndexFlatL2(d))
x = np.random.rand(10 ** 6, d).astype('float32')
print('add')
index = faiss.index_cpu_to_gpu(
faiss.StandardGpuResources(), 0, faiss.IndexFlatL2(d)
)
x = np.random.rand(10**6, d).astype("float32")
print("add")
index.add(x)
print('search')
print("search")
index.search(x, 10)
print('done')
print("done")
if __name__ == '__main__':
if __name__ == "__main__":
test_slow()
faiss/gpu/test/test_binary_cagra.py
+20 -13
View File
@@ -28,8 +28,8 @@ from faiss.contrib import evaluation
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
"CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
)
class TestInterop(unittest.TestCase):
def do_interop(self):
@@ -40,10 +40,12 @@ class TestInterop(unittest.TestCase):
index = faiss.GpuIndexBinaryCagra(res, d)
xb = np.random.randint(
low=0, high=256, size=(1000000, d // 8), dtype=np.uint8)
low=0, high=256, size=(1000000, d // 8), dtype=np.uint8
)
index.train(xb)
queries = np.random.randint(
low=0, high=256, size=(1000, d // 8), dtype=np.uint8)
low=0, high=256, size=(1000, d // 8), dtype=np.uint8
)
Dnew, Inew = index.search(queries, k)
cpu_index = faiss.index_binary_gpu_to_cpu(index)
@@ -52,7 +54,8 @@ class TestInterop(unittest.TestCase):
evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k)
deserialized_index = faiss.deserialize_index_binary(
faiss.serialize_index_binary(cpu_index))
faiss.serialize_index_binary(cpu_index)
)
gpu_index = faiss.index_binary_cpu_to_gpu(res, 0, deserialized_index)
Dnew2, Inew2 = gpu_index.search(queries, k)
@@ -64,8 +67,8 @@ class TestInterop(unittest.TestCase):
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
"CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
)
class TestComputeGT(unittest.TestCase):
def do_compute_GT(self, build_algo=None):
@@ -73,10 +76,12 @@ class TestComputeGT(unittest.TestCase):
k = 12
flat_index = faiss.IndexBinaryFlat(d)
xb = np.random.randint(
low=0, high=256, size=(1000000, d // 8), dtype=np.uint8)
low=0, high=256, size=(1000000, d // 8), dtype=np.uint8
)
flat_index.add(xb)
queries = np.random.randint(
low=0, high=256, size=(1000, d // 8), dtype=np.uint8)
low=0, high=256, size=(1000, d // 8), dtype=np.uint8
)
Dref, Iref = flat_index.search(queries, k)
res = faiss.StandardGpuResources()
@@ -102,8 +107,8 @@ class TestComputeGT(unittest.TestCase):
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
"CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
)
class TestIndexBinaryIDMap(unittest.TestCase):
"""Test IndexBinaryIDMap functionality with GpuIndexBinaryCagra"""
@@ -123,14 +128,16 @@ class TestIndexBinaryIDMap(unittest.TestCase):
index_idmap = faiss.IndexBinaryIDMap(index_gpu)
xb = np.random.randint(
low=0, high=256, size=(n, d // 8), dtype=np.uint8)
low=0, high=256, size=(n, d // 8), dtype=np.uint8
)
ids = np.arange(1000000, 1000000 + n).astype(np.int64)
index_idmap.add_with_ids(xb, ids)
nq = 1000
xq = np.random.randint(
low=0, high=256, size=(nq, d // 8), dtype=np.uint8)
low=0, high=256, size=(nq, d // 8), dtype=np.uint8
)
D, I = index_idmap.search(xq, k)
self.assertTrue(np.all(I >= 1000000))
faiss/gpu/test/test_cagra.py
+30 -10
View File
@@ -21,8 +21,12 @@ 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:
@@ -32,7 +36,9 @@ class TestComputeGT(unittest.TestCase):
# 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
@@ -95,8 +101,12 @@ class TestInterop(unittest.TestCase):
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:
@@ -122,10 +132,14 @@ class TestInterop(unittest.TestCase):
evaluation.check_ref_knn_with_draws(Dref, Iref, Dnew, Inew, k)
deserialized_index = faiss.deserialize_index(faiss.serialize_index(cpu_index))
deserialized_index = faiss.deserialize_index(
faiss.serialize_index(cpu_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)
Dnew2, Inew2 = gpu_index.search(
data_query_nt, k, numeric_type=numeric_type
)
evaluation.check_ref_knn_with_draws(Dnew2, Inew2, Dnew, Inew, k)
@@ -178,8 +192,12 @@ class TestIDMapCagra(unittest.TestCase):
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:
@@ -202,7 +220,9 @@ class TestIDMapCagra(unittest.TestCase):
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)
Dnew, Inew = idMapIndex.search(data_query_nt, k, 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_contrib_gpu.py
+42 -29
View File
@@ -11,9 +11,13 @@ import numpy as np
from common_faiss_tests import get_dataset_2
from faiss.contrib import datasets, evaluation, big_batch_search
from faiss.contrib.exhaustive_search import knn_ground_truth, \
range_ground_truth, range_search_gpu, \
range_search_max_results, exponential_query_iterator
from faiss.contrib.exhaustive_search import (
knn_ground_truth,
range_ground_truth,
range_search_gpu,
range_search_max_results,
exponential_query_iterator,
)
class TestComputeGT(unittest.TestCase):
@@ -30,7 +34,7 @@ class TestComputeGT(unittest.TestCase):
def matrix_iterator(xb, bs):
for i0 in range(0, xb.shape[0], bs):
yield xb[i0:i0 + bs]
yield xb[i0 : i0 + bs]
Dnew, Inew = knn_ground_truth(xq, matrix_iterator(xb, 1000), 10)
@@ -49,12 +53,11 @@ class TestComputeGT(unittest.TestCase):
ref_lims, ref_D, ref_I = index.range_search(xq, threshold)
new_lims, new_D, new_I = range_ground_truth(
xq, ds.database_iterator(bs=100), threshold,
metric_type=metric)
xq, ds.database_iterator(bs=100), threshold, metric_type=metric
)
evaluation.check_ref_range_results(
ref_lims, ref_D, ref_I,
new_lims, new_D, new_I
ref_lims, ref_D, ref_I, new_lims, new_D, new_I
)
def test_range_L2(self):
@@ -84,14 +87,16 @@ class TestComputeGT(unittest.TestCase):
query_iterator = exponential_query_iterator(xq)
radius1, lims_new, Dnew, Inew = range_search_max_results(
index, query_iterator, ds.d // 2,
min_results=Dref.size, clip_to_min=True,
ngpu=1
index,
query_iterator,
ds.d // 2,
min_results=Dref.size,
clip_to_min=True,
ngpu=1,
)
evaluation.check_ref_range_results(
lims_ref, Dref, Iref,
lims_new, Dnew, Inew
lims_ref, Dref, Iref, lims_new, Dnew, Inew
)
@unittest.skipIf(faiss.get_num_gpus() < 2, "multiple GPU only test")
@@ -113,17 +118,20 @@ class TestBigBatchSearch(unittest.TestCase):
def pairwise_distances(xq, xb, metric=faiss.METRIC_L2):
return faiss.pairwise_distance_gpu(
res, xq, xb, metric=faiss.METRIC_L2)
res, xq, xb, metric=faiss.METRIC_L2
)
def knn_function(xq, xb, k, metric=faiss.METRIC_L2):
return faiss.knn_gpu(res, xq, xb, k, metric=faiss.METRIC_L2)
for method in "pairwise_distances", "knn_function":
Dnew, Inew = big_batch_search.big_batch_search(
index, ds.get_queries(),
k, method=method,
index,
ds.get_queries(),
k,
method=method,
pairwise_distances=pairwise_distances,
knn=knn_function
knn=knn_function,
)
self.assertLess((Inew != Iref).sum() / Iref.size, 1e-4)
np.testing.assert_almost_equal(Dnew, Dref, decimal=4)
@@ -151,16 +159,22 @@ class TestBigBatchSearchMultiGPU(unittest.TestCase):
def knn_function(xq, xb, k, metric=faiss.METRIC_L2, thread_id=None):
return faiss.knn_gpu(
res[thread_id], xq, xb, k,
metric=faiss.METRIC_L2, device=thread_id
res[thread_id],
xq,
xb,
k,
metric=faiss.METRIC_L2,
device=thread_id,
)
Dnew, Inew = big_batch_search.big_batch_search(
index, ds.get_queries(),
k, method="knn_function",
index,
ds.get_queries(),
k,
method="knn_function",
knn=knn_function,
threaded=8,
computation_threads=ngpu
computation_threads=ngpu,
)
self.assertLess((Inew != Iref).sum() / Iref.size, 1e-4)
np.testing.assert_almost_equal(Dnew, Dref, decimal=4)
@@ -192,18 +206,17 @@ class TestRangeSearchGpu(unittest.TestCase):
# mixed GPU / CPU run
Lnew, Dnew, Inew = range_search_gpu(
ds.get_queries(), threshold, index_gpu, index_cpu, gpu_k=4)
evaluation.check_ref_range_results(
Lref, Dref, Iref,
Lnew, Dnew, Inew
ds.get_queries(), threshold, index_gpu, index_cpu, gpu_k=4
)
evaluation.check_ref_range_results(Lref, Dref, Iref, Lnew, Dnew, Inew)
# also test the version without CPU search
Lnew2, Dnew2, Inew2 = range_search_gpu(
ds.get_queries(), threshold, index_gpu, None, gpu_k=4)
ds.get_queries(), threshold, index_gpu, None, gpu_k=4
)
for q in range(ds.nq):
ref = Iref[Lref[q]:Lref[q+1]]
new = Inew2[Lnew2[q]:Lnew2[q+1]]
ref = Iref[Lref[q] : Lref[q + 1]]
new = Inew2[Lnew2[q] : Lnew2[q + 1]]
if nres_per_query[q] <= 4:
self.assertEqual(set(ref), set(new))
else:
faiss/gpu/test/test_cuvs.py
+6 -5
View File
@@ -11,8 +11,8 @@ from faiss.contrib.datasets import SyntheticDataset
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if CUVS is compiled in")
"CUVS" not in faiss.get_compile_options(), "only if CUVS is compiled in"
)
class TestBfKnn(unittest.TestCase):
def test_large_k_search(self):
@@ -29,7 +29,6 @@ class TestBfKnn(unittest.TestCase):
_, I = index_gpu.search(ds.get_queries(), k)
np.testing.assert_equal(I.shape, (ds.nq, k))
def test_bfKnn(self):
ds = SyntheticDataset(32, 0, 4321, 1234)
@@ -40,13 +39,15 @@ class TestBfKnn(unittest.TestCase):
# Faiss internal implementation
Dnew, Inew = faiss.knn_gpu(
res, ds.get_queries(), ds.get_database(), 12, use_cuvs=False)
res, ds.get_queries(), ds.get_database(), 12, use_cuvs=False
)
np.testing.assert_allclose(Dref, Dnew, atol=1e-4)
np.testing.assert_array_equal(Iref, Inew)
# cuVS version
Dnew, Inew = faiss.knn_gpu(
res, ds.get_queries(), ds.get_database(), 12, use_cuvs=True)
res, ds.get_queries(), ds.get_database(), 12, use_cuvs=True
)
np.testing.assert_allclose(Dref, Dnew, atol=1e-4)
np.testing.assert_array_equal(Iref, Inew)
faiss/gpu/test/test_gpu_basics.py
+72 -56
View File
@@ -3,7 +3,12 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import absolute_import, division, print_function, unicode_literals
from __future__ import (
absolute_import,
division,
print_function,
unicode_literals,
)
import unittest
import numpy as np
@@ -15,12 +20,12 @@ from common_faiss_tests import get_dataset_2
class ReferencedObject(unittest.TestCase):
d = 16
xb = np.random.rand(256, d).astype('float32')
xb = np.random.rand(256, d).astype("float32")
nlist = 128
d_bin = 256
xb_bin = np.random.randint(256, size=(10000, d_bin // 8)).astype('uint8')
xq_bin = np.random.randint(256, size=(1000, d_bin // 8)).astype('uint8')
xb_bin = np.random.randint(256, size=(10000, d_bin // 8)).astype("uint8")
xq_bin = np.random.randint(256, size=(1000, d_bin // 8)).astype("uint8")
def test_proxy(self):
index = faiss.IndexReplicas()
@@ -33,32 +38,33 @@ class ReferencedObject(unittest.TestCase):
def test_resources(self):
# this used to crash!
index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0,
faiss.IndexFlatL2(self.d))
index = faiss.index_cpu_to_gpu(
faiss.StandardGpuResources(), 0, faiss.IndexFlatL2(self.d)
)
index.add(self.xb)
def test_flat(self):
index = faiss.GpuIndexFlat(faiss.StandardGpuResources(),
self.d, faiss.METRIC_L2)
index = faiss.GpuIndexFlat(
faiss.StandardGpuResources(), self.d, faiss.METRIC_L2
)
index.add(self.xb)
def test_ivfflat(self):
index = faiss.GpuIndexIVFFlat(
faiss.StandardGpuResources(),
self.d, self.nlist, faiss.METRIC_L2)
faiss.StandardGpuResources(), self.d, self.nlist, faiss.METRIC_L2
)
index.train(self.xb)
def test_ivfpq(self):
index_cpu = faiss.IndexIVFPQ(
faiss.IndexFlatL2(self.d),
self.d, self.nlist, 2, 8)
faiss.IndexFlatL2(self.d), self.d, self.nlist, 2, 8
)
# speed up test
index_cpu.pq.cp.niter = 2
index_cpu.do_polysemous_training = False
index_cpu.train(self.xb)
index = faiss.GpuIndexIVFPQ(
faiss.StandardGpuResources(), index_cpu)
index = faiss.GpuIndexIVFPQ(faiss.StandardGpuResources(), index_cpu)
index.add(self.xb)
def test_binary_flat(self):
@@ -68,8 +74,9 @@ class ReferencedObject(unittest.TestCase):
index_ref.add(self.xb_bin)
D_ref, I_ref = index_ref.search(self.xq_bin, k)
index = faiss.GpuIndexBinaryFlat(faiss.StandardGpuResources(),
self.d_bin)
index = faiss.GpuIndexBinaryFlat(
faiss.StandardGpuResources(), self.d_bin
)
index.add(self.xb_bin)
D, I = index.search(self.xq_bin, k)
@@ -89,15 +96,17 @@ class ReferencedObject(unittest.TestCase):
def test_stress(self):
# a mixture of the above, from issue #631
target = np.random.rand(50, 16).astype('float32')
target = np.random.rand(50, 16).astype("float32")
index = faiss.IndexReplicas()
size, dim = target.shape
num_gpu = 4
for _i in range(num_gpu):
config = faiss.GpuIndexFlatConfig()
config.device = 0 # simulate on a single GPU
sub_index = faiss.GpuIndexFlatIP(faiss.StandardGpuResources(), dim, config)
config.device = 0 # simulate on a single GPU
sub_index = faiss.GpuIndexFlatIP(
faiss.StandardGpuResources(), dim, config
)
index.addIndex(sub_index)
index = faiss.IndexIDMap(index)
@@ -105,7 +114,6 @@ class ReferencedObject(unittest.TestCase):
index.add_with_ids(target, ids)
class TestGPUKmeans(unittest.TestCase):
def test_kmeans(self):
@@ -113,7 +121,7 @@ class TestGPUKmeans(unittest.TestCase):
nb = 1000
k = 10
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
km1 = faiss.Kmeans(d, k)
obj1 = km1.train(xb)
@@ -153,8 +161,8 @@ class TestAlternativeDistances(unittest.TestCase):
nq = 100
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
index_ref = faiss.IndexFlat(d, metric)
index_ref.metric_arg = metric_arg
@@ -176,8 +184,8 @@ class TestAlternativeDistances(unittest.TestCase):
np.testing.assert_array_equal(Inew, Iref)
@unittest.skipIf(
"CUVS" in faiss.get_compile_options(),
"only if CUVS is not compiled in")
"CUVS" in faiss.get_compile_options(), "only if CUVS is not compiled in"
)
def do_test_gower(self):
"""Special test for Gower distance with mixed numeric/categorical data"""
res = faiss.StandardGpuResources()
@@ -188,16 +196,16 @@ class TestAlternativeDistances(unittest.TestCase):
rs = np.random.RandomState(123)
# Create mixed data: first half numeric [0,1], second half categorical (negative)
xb = np.zeros((nb, d), dtype='float32')
xq = np.zeros((nq, d), dtype='float32')
xb = np.zeros((nb, d), dtype="float32")
xq = np.zeros((nq, d), dtype="float32")
# Numeric features in [0,1] range
xb[:, :d//2] = rs.rand(nb, d//2)
xq[:, :d//2] = rs.rand(nq, d//2)
xb[:, : d // 2] = rs.rand(nb, d // 2)
xq[:, : d // 2] = rs.rand(nq, d // 2)
# Categorical features (negative integers)
xb[:, d//2:] = -rs.randint(1, 5, size=(nb, d//2)).astype('float32')
xq[:, d//2:] = -rs.randint(1, 5, size=(nq, d//2)).astype('float32')
xb[:, d // 2 :] = -rs.randint(1, 5, size=(nb, d // 2)).astype("float32")
xq[:, d // 2 :] = -rs.randint(1, 5, size=(nq, d // 2)).astype("float32")
# CPU reference
index_ref = faiss.IndexFlat(d, faiss.METRIC_GOWER)
@@ -235,7 +243,9 @@ class TestGpuRef(unittest.TestCase):
def test_gpu_ref(self):
# this crashes
dim = 256
training_data = np.random.randint(256, size=(10000, dim // 8)).astype('uint8')
training_data = np.random.randint(256, size=(10000, dim // 8)).astype(
"uint8"
)
centroids = 330
def create_cpu(dim):
@@ -267,7 +277,7 @@ def make_t(num, d, clamp=False, seed=None):
x = rs.rand(num, d).astype(np.float32)
if clamp:
x = (x * 255).astype('uint8').astype('float32')
x = (x * 255).astype("uint8").astype("float32")
return x
@@ -315,10 +325,8 @@ class TestKnn(unittest.TestCase):
if vectorsMemoryLimit > 0 or queriesMemoryLimit > 0:
faiss.bfKnn_tiling(
res,
params,
vectorsMemoryLimit,
queriesMemoryLimit)
res, params, vectorsMemoryLimit, queriesMemoryLimit
)
else:
faiss.bfKnn(res, params)
@@ -326,9 +334,16 @@ class TestKnn(unittest.TestCase):
np.testing.assert_array_equal(out_i, ref_i)
faiss.knn_gpu(
res, qs, xs, k, out_d, out_i, device=gpu_id,
res,
qs,
xs,
k,
out_d,
out_i,
device=gpu_id,
vectorsMemoryLimit=vectorsMemoryLimit,
queriesMemoryLimit=queriesMemoryLimit)
queriesMemoryLimit=queriesMemoryLimit,
)
np.testing.assert_allclose(ref_d, out_d, atol=1e-5)
np.testing.assert_array_equal(out_i, ref_i)
@@ -368,8 +383,10 @@ class TestKnn(unittest.TestCase):
faiss.bfKnn(res, params)
self.assertGreaterEqual((out_i_f16 == ref_i_f16).sum(), ref_i_f16.size - 5)
np.testing.assert_allclose(ref_d_f16, out_d_f16, atol = 2e-3)
self.assertGreaterEqual(
(out_i_f16 == ref_i_f16).sum(), ref_i_f16.size - 5
)
np.testing.assert_allclose(ref_d_f16, out_d_f16, atol=2e-3)
class TestAllPairwiseDistance(unittest.TestCase):
@@ -382,7 +399,7 @@ class TestAllPairwiseDistance(unittest.TestCase):
faiss.METRIC_Canberra,
faiss.METRIC_BrayCurtis,
faiss.METRIC_JensenShannon,
faiss.METRIC_Jaccard
faiss.METRIC_Jaccard,
]
for metric in metrics:
@@ -409,7 +426,7 @@ class TestAllPairwiseDistance(unittest.TestCase):
# Try f32 data/queries
params = faiss.GpuDistanceParams()
params.metric = metric
params.k = -1 # all pairwise
params.k = -1 # all pairwise
params.dims = d
params.vectors = faiss.swig_ptr(xs)
params.numVectors = nb
@@ -429,7 +446,7 @@ class TestAllPairwiseDistance(unittest.TestCase):
if faiss.is_similarity_metric(metric):
ref_d = np.sort(ref_d, axis=1)
print('f32', np.abs(ref_d - out_d).max())
print("f32", np.abs(ref_d - out_d).max())
np.testing.assert_allclose(ref_d, out_d, atol=1e-5)
@@ -462,9 +479,9 @@ class TestAllPairwiseDistance(unittest.TestCase):
if faiss.is_similarity_metric(metric):
ref_d_f16 = np.sort(ref_d_f16, axis=1)
print('f16', np.abs(ref_d_f16 - out_d_f16).max())
print("f16", np.abs(ref_d_f16 - out_d_f16).max())
np.testing.assert_allclose(ref_d_f16, out_d_f16, atol = 4e-3)
np.testing.assert_allclose(ref_d_f16, out_d_f16, atol=4e-3)
def test_gower_pairwise(self):
"""Test Gower distance with GPU pairwise distance computation"""
@@ -474,18 +491,18 @@ class TestAllPairwiseDistance(unittest.TestCase):
nq = 10
rs = np.random.RandomState(123)
# Create mixed data: first half numeric [0,1], second half categorical (negative)
xs = np.zeros((nb, d), dtype=np.float32)
qs = np.zeros((nq, d), dtype=np.float32)
# Numeric features in [0,1] range
xs[:, :d//2] = rs.rand(nb, d//2)
qs[:, :d//2] = rs.rand(nq, d//2)
xs[:, : d // 2] = rs.rand(nb, d // 2)
qs[:, : d // 2] = rs.rand(nq, d // 2)
# Categorical features (negative integers)
xs[:, d//2:] = -rs.randint(1, 4, size=(nb, d//2)).astype('float32')
qs[:, d//2:] = -rs.randint(1, 4, size=(nq, d//2)).astype('float32')
xs[:, d // 2 :] = -rs.randint(1, 4, size=(nb, d // 2)).astype("float32")
qs[:, d // 2 :] = -rs.randint(1, 4, size=(nq, d // 2)).astype("float32")
res = faiss.StandardGpuResources()
@@ -513,11 +530,10 @@ class TestAllPairwiseDistance(unittest.TestCase):
# IndexFlat will sort the results, so we need to do the same on our end
out_d = np.sort(out_d, axis=1)
print('Gower max diff:', np.abs(ref_d - out_d).max())
print("Gower max diff:", np.abs(ref_d - out_d).max())
np.testing.assert_allclose(ref_d, out_d, atol=1e-5)
def eval_codec(q, xb):
codes = q.compute_codes(xb)
decoded = q.decode(codes)
@@ -536,7 +552,7 @@ class TestResidualQuantizer(unittest.TestCase):
# know there is a problem with it.
@unittest.skip("Skipped due to ncall memory corruption.")
def test_with_gpu(self):
""" check that we get the same results with a GPU quantizer and a CPU quantizer """
"""check that we get the same results with a GPU quantizer and a CPU quantizer"""
d = 32
nt = 3000
nb = 1000
faiss/gpu/test/test_gpu_index.py
+96 -62
View File
@@ -3,7 +3,12 @@
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import absolute_import, division, print_function, unicode_literals
from __future__ import (
absolute_import,
division,
print_function,
unicode_literals,
)
import math
import unittest
@@ -30,8 +35,8 @@ class TestIVFSearchPreassigned(unittest.TestCase):
idx_gpu.nprobe = nprobe
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -40,7 +45,8 @@ class TestIVFSearchPreassigned(unittest.TestCase):
q_d, q_i = idx_gpu.quantizer.search(xq, nprobe)
preassigned_d, preassigned_i = ivf_tools.search_preassigned(
idx_gpu, xq, k, q_i, q_d)
idx_gpu, xq, k, q_i, q_d
)
# Search using the standard API
d, i = idx_gpu.search(xq, k)
@@ -60,12 +66,14 @@ class TestIVFSearchPreassigned(unittest.TestCase):
config = faiss.GpuIndexIVFPQConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, 4, 8, faiss.METRIC_L2, config)
idx_gpu = faiss.GpuIndexIVFPQ(
res, d, nlist, 4, 8, faiss.METRIC_L2, config
)
idx_gpu.nprobe = nprobe
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -74,7 +82,8 @@ class TestIVFSearchPreassigned(unittest.TestCase):
q_d, q_i = idx_gpu.quantizer.search(xq, nprobe)
preassigned_d, preassigned_i = ivf_tools.search_preassigned(
idx_gpu, xq, k, q_i, q_d)
idx_gpu, xq, k, q_i, q_d
)
# Search using the standard API
d, i = idx_gpu.search(xq, k)
@@ -93,14 +102,17 @@ class TestIVFSearchPreassigned(unittest.TestCase):
k = 50
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist,
res,
d,
nlist,
faiss.ScalarQuantizer.QT_6bit,
faiss.METRIC_INNER_PRODUCT)
faiss.METRIC_INNER_PRODUCT,
)
idx_gpu.nprobe = nprobe
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -109,7 +121,8 @@ class TestIVFSearchPreassigned(unittest.TestCase):
q_d, q_i = idx_gpu.quantizer.search(xq, nprobe)
preassigned_d, preassigned_i = ivf_tools.search_preassigned(
idx_gpu, xq, k, q_i, q_d)
idx_gpu, xq, k, q_i, q_d
)
# Search using the standard API
d, i = idx_gpu.search(xq, k)
@@ -132,8 +145,8 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
idx_cpu = faiss.IndexIVFFlat(q, d, nlist)
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
idx_cpu.train(xb)
idx_cpu.add(xb)
@@ -141,8 +154,10 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
# construct a GPU index using the same trained coarse quantizer
# from the CPU index
config = faiss.GpuIndexIVFFlatConfig()
idx_gpu = faiss.GpuIndexIVFFlat(res, q, d, nlist, faiss.METRIC_L2, config)
assert(idx_gpu.is_trained)
idx_gpu = faiss.GpuIndexIVFFlat(
res, q, d, nlist, faiss.METRIC_L2, config
)
assert idx_gpu.is_trained
idx_gpu.add(xb)
k = 20
@@ -155,7 +170,6 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
self.assertGreaterEqual((i_g == i_c).sum(), i_g.size * 0.9)
self.assertTrue(np.allclose(d_g, d_c, rtol=5e-5, atol=5e-5))
def test_ivfsq_pu_coarse(self):
res = faiss.StandardGpuResources()
d = 128
@@ -168,11 +182,12 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
q = faiss.IndexFlatL2(d)
idx_cpu = faiss.IndexIVFScalarQuantizer(
q, d, nlist, qtype, faiss.METRIC_L2, use_residual)
q, d, nlist, qtype, faiss.METRIC_L2, use_residual
)
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
idx_cpu.train(xb)
idx_cpu.add(xb)
@@ -180,8 +195,9 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
# construct a GPU index using the same trained coarse quantizer
# from the CPU index
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, q, d, nlist, qtype, faiss.METRIC_L2, use_residual)
assert(not idx_gpu.is_trained)
res, q, d, nlist, qtype, faiss.METRIC_L2, use_residual
)
assert not idx_gpu.is_trained
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -208,11 +224,12 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
nprobe_lvl_2 = 10
rs = np.random.RandomState(567)
coarse_centroids = rs.rand(nlist_lvl_2, d).astype('float32')
coarse_centroids = rs.rand(nlist_lvl_2, d).astype("float32")
# Construct an IVFFlat index for usage as a coarse quantizer
idx_coarse_cpu = faiss.IndexIVFFlat(
faiss.IndexFlatL2(d), d, nlist_lvl_1)
faiss.IndexFlatL2(d), d, nlist_lvl_1
)
idx_coarse_cpu.set_direct_map_type(faiss.DirectMap.Hashtable)
idx_coarse_cpu.nprobe = nprobe_lvl_1
@@ -220,12 +237,11 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
idx_coarse_cpu.add(coarse_centroids)
idx_coarse_cpu.make_direct_map()
assert(idx_coarse_cpu.ntotal == nlist_lvl_2)
assert idx_coarse_cpu.ntotal == nlist_lvl_2
idx_cpu = faiss.IndexIVFPQ(
idx_coarse_cpu, d, nlist_lvl_2, 4, 8)
idx_cpu = faiss.IndexIVFPQ(idx_coarse_cpu, d, nlist_lvl_2, 4, 8)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
idx_cpu.train(xb)
idx_cpu.add(xb)
idx_cpu.nprobe = nprobe_lvl_2
@@ -235,8 +251,9 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
config = faiss.GpuIndexIVFPQConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFPQ(
res, idx_coarse_cpu, d, nlist_lvl_2, 4, 8, faiss.METRIC_L2, config)
assert(not idx_gpu.is_trained)
res, idx_coarse_cpu, d, nlist_lvl_2, 4, 8, faiss.METRIC_L2, config
)
assert not idx_gpu.is_trained
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -248,7 +265,7 @@ class TestIVFPluggableCoarseQuantizer(unittest.TestCase):
for use_precomputed in [False, True]:
idx_gpu.setPrecomputedCodes(use_precomputed)
xq = rs.rand(nq, d).astype('float32')
xq = rs.rand(nq, d).astype("float32")
d_g, i_g = idx_gpu.search(xq, k)
d_c, i_c = idx_cpu.search(xq, k)
@@ -265,8 +282,8 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
nq = 20
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
nlist = int(math.sqrt(nb))
sub_q = 16
@@ -274,9 +291,13 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config)
idx_gpu = faiss.GpuIndexIVFPQ(
res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config
)
q = faiss.IndexFlatL2(d)
idx_cpu = faiss.IndexIVFPQ(q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2)
idx_cpu = faiss.IndexIVFPQ(
q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2
)
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -309,8 +330,8 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
nq = 20
rs = np.random.RandomState(234)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
nlist = int(math.sqrt(nb))
sub_q = 16
@@ -319,9 +340,13 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config)
idx_gpu = faiss.GpuIndexIVFPQ(
res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config
)
q = faiss.IndexFlatL2(d)
idx_cpu = faiss.IndexIVFPQ(q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2)
idx_cpu = faiss.IndexIVFPQ(
q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2
)
idx_gpu.train(xb)
idx_gpu.add(xb)
@@ -353,8 +378,8 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
nq = 20
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
nlist = int(math.sqrt(nb))
sub_q = 16
@@ -363,9 +388,13 @@ class TestInterleavedIVFPQLayout(unittest.TestCase):
config = faiss.GpuIndexIVFPQConfig()
config.interleavedLayout = True
idx_gpu = faiss.GpuIndexIVFPQ(res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config)
idx_gpu = faiss.GpuIndexIVFPQ(
res, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2, config
)
q = faiss.IndexFlatL2(d)
idx_cpu = faiss.IndexIVFPQ(q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2)
idx_cpu = faiss.IndexIVFPQ(
q, d, nlist, sub_q, bits_per_code, faiss.METRIC_L2
)
idx_cpu.train(xb)
idx_cpu.add(xb)
@@ -398,11 +427,11 @@ class TestIVFIndices(unittest.TestCase):
nlist = 10
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
xb_indices = (xb_indices_base + 4294967296).astype("int64")
config = faiss.GpuIndexIVFFlatConfig()
idx = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, config)
@@ -432,15 +461,16 @@ class TestIVFIndices(unittest.TestCase):
nbits = 8
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
xb_indices = (xb_indices_base + 4294967296).astype("int64")
config = faiss.GpuIndexIVFPQConfig()
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits,
faiss.METRIC_L2, config)
idx = faiss.GpuIndexIVFPQ(
res, d, nlist, M, nbits, faiss.METRIC_L2, config
)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
@@ -451,8 +481,9 @@ class TestIVFIndices(unittest.TestCase):
config.indicesOptions = faiss.INDICES_32_BIT
# 32-bit indices are not supported with cuVS
config.use_cuvs = False
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits,
faiss.METRIC_L2, config)
idx = faiss.GpuIndexIVFPQ(
res, d, nlist, M, nbits, faiss.METRIC_L2, config
)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
@@ -468,15 +499,16 @@ class TestIVFIndices(unittest.TestCase):
qtype = faiss.ScalarQuantizer.QT_4bit
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
xb_indices = (xb_indices_base + 4294967296).astype("int64")
config = faiss.GpuIndexIVFScalarQuantizerConfig()
idx = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, qtype,
faiss.METRIC_L2, True, config)
idx = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, faiss.METRIC_L2, True, config
)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
@@ -485,8 +517,9 @@ class TestIVFIndices(unittest.TestCase):
# Store values using 32-bit indices instead
config.indicesOptions = faiss.INDICES_32_BIT
idx = faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, qtype,
faiss.METRIC_L2, True, config)
idx = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, faiss.METRIC_L2, True, config
)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
@@ -518,15 +551,16 @@ class TestInvalidParams(unittest.TestCase):
nbits = 8
rs = np.random.RandomState(567)
xb = rs.rand(nb, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xb_indices_base = np.arange(nb, dtype=np.int64)
# Force values to not be representable in int32
xb_indices = (xb_indices_base + 4294967296).astype('int64')
xb_indices = (xb_indices_base + 4294967296).astype("int64")
config = faiss.GpuIndexIVFPQConfig()
idx = faiss.GpuIndexIVFPQ(res, d, nlist, M, nbits,
faiss.METRIC_L2, config)
idx = faiss.GpuIndexIVFPQ(
res, d, nlist, M, nbits, faiss.METRIC_L2, config
)
idx.train(xb)
idx.add_with_ids(xb, xb_indices)
faiss/gpu/test/test_gpu_index_ivfflat.py
+1 -1
View File
@@ -12,7 +12,7 @@ import numpy as np
class TestGpuIndexIvfflat(unittest.TestCase):
def test_reconstruct_n(self):
index = faiss.index_factory(4, "IVF10,Flat")
x = np.random.RandomState(123).rand(10, 4).astype('float32')
x = np.random.RandomState(123).rand(10, 4).astype("float32")
index.train(x)
index.add(x)
res = faiss.StandardGpuResources()
faiss/gpu/test/test_gpu_index_ivfsq.py
+36 -26
View File
@@ -11,9 +11,9 @@ import faiss
def make_t(num, d, clamp=False):
rs = np.random.RandomState(123)
x = rs.rand(num, d).astype('float32')
x = rs.rand(num, d).astype("float32")
if clamp:
x = (x * 255).astype('uint8').astype('float32')
x = (x * 255).astype("uint8").astype("float32")
return x
@@ -22,7 +22,8 @@ def make_indices_copy_from_cpu(nlist, d, qtype, by_residual, metric, clamp):
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(
quantizer_cp, d, nlist, qtype, metric, by_residual)
quantizer_cp, d, nlist, qtype, metric, by_residual
)
idx_cpu.train(to_train)
idx_cpu.add(to_train)
@@ -44,13 +45,15 @@ def make_indices_copy_from_gpu(nlist, d, qtype, by_residual, metric, clamp):
config = faiss.GpuIndexIVFScalarQuantizerConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, metric, by_residual, config)
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)
quantizer_cp, d, nlist, qtype, metric, by_residual
)
idx_gpu.copyTo(idx_cpu)
return idx_cpu, idx_gpu
@@ -61,7 +64,8 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
quantizer_cp = faiss.IndexFlat(d, metric)
idx_cpu = faiss.IndexIVFScalarQuantizer(
quantizer_cp, d, nlist, qtype, metric, by_residual)
quantizer_cp, d, nlist, qtype, metric, by_residual
)
assert by_residual == idx_cpu.by_residual
idx_cpu.train(to_train)
@@ -72,7 +76,8 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
config = faiss.GpuIndexIVFScalarQuantizerConfig()
config.use_cuvs = False
idx_gpu = faiss.GpuIndexIVFScalarQuantizer(
res, d, nlist, qtype, metric, by_residual, config)
res, d, nlist, qtype, metric, by_residual, config
)
assert by_residual == idx_gpu.by_residual
idx_gpu.train(to_train)
@@ -80,6 +85,7 @@ def make_indices_train(nlist, d, qtype, by_residual, metric, clamp):
return idx_cpu, idx_gpu
#
# Testing functions
#
@@ -116,8 +122,8 @@ def compare_results(d1, i1, d2, i2):
# Invalid results should be the same for both
# (except if we happen to hit different centroids)
for inv1, inv2 in zip(invalid1, invalid2):
if (len(inv1) != len(inv2)):
print('mismatch ', len(inv1), len(inv2), inv2[0])
if len(inv1) != len(inv2):
print("mismatch ", len(inv1), len(inv2), inv2[0])
assert len(inv1) == len(inv2)
idx_invalid += len(inv2)
@@ -150,8 +156,8 @@ def check_diffs(total_num, in_window_thresh, diffs, diff_inf, invalid):
if abs(diff) <= diff_window:
in_window += diffs[diff] / total_num
if (in_window < in_window_thresh):
print('error {} {}'.format(in_window, in_window_thresh))
if in_window < in_window_thresh:
print("error {} {}".format(in_window, in_window_thresh))
assert in_window >= in_window_thresh
@@ -163,27 +169,30 @@ def do_test_with_index(ci, gi, nprobe, k, clamp, in_window_thresh):
gi.nprobe = gi.nprobe
total_num = num_query * k
check_diffs(total_num, in_window_thresh,
*compare_results(*ci.search(to_query, k),
*gi.search(to_query, k)))
check_diffs(
total_num,
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,
by_residual, metric, clamp)
clamp = qtype == faiss.ScalarQuantizer.QT_8bit_direct
ci, gi = make_indices_copy_from_cpu(
nlist, d, qtype, by_residual, metric, clamp
)
# A direct copy should be much more closely in agreement
# (except for fp accumulation order differences)
do_test_with_index(ci, gi, nprobe, k, clamp, 0.99)
ci, gi = make_indices_copy_from_gpu(nlist, d, qtype,
by_residual, metric, clamp)
ci, gi = make_indices_copy_from_gpu(
nlist, d, qtype, by_residual, metric, clamp
)
# A direct copy should be much more closely in agreement
# (except for fp accumulation order differences)
do_test_with_index(ci, gi, nprobe, k, clamp, 0.99)
ci, gi = make_indices_train(nlist, d, qtype,
by_residual, metric, clamp)
ci, gi = make_indices_train(nlist, d, qtype, by_residual, metric, clamp)
# Separate training can produce a slightly different coarse quantizer
# and residuals
do_test_with_index(ci, gi, nprobe, k, clamp, 0.8)
@@ -195,15 +204,16 @@ def do_multi_test(qtype):
k = 50
for d in [11, 64, 77]:
if (qtype != faiss.ScalarQuantizer.QT_8bit_direct):
if qtype != faiss.ScalarQuantizer.QT_8bit_direct:
# residual doesn't make sense here
do_test(nlist, d, qtype, True,
faiss.METRIC_L2, nprobe, k)
do_test(nlist, d, qtype, True,
faiss.METRIC_INNER_PRODUCT, nprobe, k)
do_test(nlist, d, qtype, True, faiss.METRIC_L2, nprobe, k)
do_test(
nlist, d, qtype, True, faiss.METRIC_INNER_PRODUCT, nprobe, k
)
do_test(nlist, d, qtype, False, faiss.METRIC_L2, nprobe, k)
do_test(nlist, d, qtype, False, faiss.METRIC_INNER_PRODUCT, nprobe, k)
#
# Test
#
faiss/gpu/test/test_gpu_index_refs.py
+3 -3
View File
@@ -80,8 +80,8 @@ class TestRefs(unittest.TestCase):
def test_GpuIndexIVFPQ(self):
index_to_delete = faiss.IndexIVFPQ(
faiss.IndexFlatL2(self.d),
self.d, self.nlist, 2, 8)
faiss.IndexFlatL2(self.d), self.d, self.nlist, 2, 8
)
idx = faiss.GpuIndexIVFPQ(self.res, index_to_delete)
do_multi_test(idx, index_to_delete, self.db)
@@ -92,7 +92,7 @@ class TestRefs(unittest.TestCase):
self.nlist,
faiss.ScalarQuantizer.QT_8bit_direct,
faiss.METRIC_L2,
False
False,
)
idx = faiss.GpuIndexIVFScalarQuantizer(self.res, index_to_delete)
do_multi_test(idx, index_to_delete, self.db)
faiss/gpu/test/test_gpu_index_serialize.py
+19 -5
View File
@@ -10,7 +10,7 @@ import faiss
def make_t(num, d):
rs = np.random.RandomState(123)
return rs.rand(num, d).astype('float32')
return rs.rand(num, d).astype("float32")
class TestGpuSerialize(unittest.TestCase):
@@ -38,10 +38,22 @@ class TestGpuSerialize(unittest.TestCase):
# IVFSQ
config = faiss.GpuIndexIVFScalarQuantizerConfig()
config.use_cuvs = False
indexes.append(faiss.GpuIndexIVFScalarQuantizer(res, d, nlist, faiss.ScalarQuantizer.QT_fp16, faiss.METRIC_L2, True, config))
indexes.append(
faiss.GpuIndexIVFScalarQuantizer(
res,
d,
nlist,
faiss.ScalarQuantizer.QT_fp16,
faiss.METRIC_L2,
True,
config,
)
)
# IVFPQ
indexes.append(faiss.GpuIndexIVFPQ(res, d, nlist, 4, 8, faiss.METRIC_L2))
indexes.append(
faiss.GpuIndexIVFPQ(res, d, nlist, 4, 8, faiss.METRIC_L2)
)
for index in indexes:
index.train(train)
@@ -51,11 +63,13 @@ class TestGpuSerialize(unittest.TestCase):
ser = faiss.serialize_index(faiss.index_gpu_to_cpu(index))
cpu_index = faiss.deserialize_index(ser)
gpu_cloner_options = faiss.GpuClonerOptions()
if isinstance(index, faiss.GpuIndexIVFScalarQuantizer):
gpu_cloner_options.use_cuvs = False
gpu_index_restore = faiss.index_cpu_to_gpu(res, 0, cpu_index, gpu_cloner_options)
gpu_index_restore = faiss.index_cpu_to_gpu(
res, 0, cpu_index, gpu_cloner_options
)
restore_d, restore_i = gpu_index_restore.search(query, k)
faiss/gpu/test/test_index_cpu_to_gpu.py
+27 -20
View File
@@ -17,7 +17,7 @@ class TestMoveToGpu(unittest.TestCase):
def create_index(self, factory_string):
dimension = 128
n = 2500
db_vectors = np.random.random((n, dimension)).astype('float32')
db_vectors = np.random.random((n, dimension)).astype("float32")
index = faiss.index_factory(dimension, factory_string)
index.train(db_vectors)
if factory_string.startswith("IDMap"):
@@ -26,9 +26,9 @@ class TestMoveToGpu(unittest.TestCase):
index.add(db_vectors)
return index
def create_and_clone(self, factory_string,
allowCpuCoarseQuantizer=None,
use_cuvs=None):
def create_and_clone(
self, factory_string, allowCpuCoarseQuantizer=None, use_cuvs=None
):
idx = self.create_index(factory_string)
config = faiss.GpuClonerOptions()
if allowCpuCoarseQuantizer is not None:
@@ -42,20 +42,25 @@ class TestMoveToGpu(unittest.TestCase):
self.create_and_clone(factory_string)
except Exception as e:
if "not implemented" not in str(e):
self.fail("Expected an exception but no exception was "
"thrown for factory_string: %s." % factory_string)
self.fail(
"Expected an exception but no exception was "
"thrown for factory_string: %s." % factory_string
)
def verify_clones_successfully(self, factory_string,
allowCpuCoarseQuantizer=None,
use_cuvs=None):
def verify_clones_successfully(
self, factory_string, allowCpuCoarseQuantizer=None, use_cuvs=None
):
try:
self.create_and_clone(
factory_string,
allowCpuCoarseQuantizer=allowCpuCoarseQuantizer,
use_cuvs=use_cuvs)
use_cuvs=use_cuvs,
)
except Exception as e:
self.fail("Unexpected exception thrown factory_string: "
"%s; error message: %s." % (factory_string, str(e)))
self.fail(
"Unexpected exception thrown factory_string: "
"%s; error message: %s." % (factory_string, str(e))
)
def test_not_implemented_indices(self):
self.verify_throws_not_implemented_exception("PQ16")
@@ -76,19 +81,21 @@ class TestMoveToGpu(unittest.TestCase):
# set use_cuvs to false, these index types are not supported on cuVS
self.verify_clones_successfully("IVF32,SQ8", use_cuvs=False)
self.verify_clones_successfully(
"PCA32,IVF32,SQ8", use_cuvs=False)
self.verify_clones_successfully("PCA32,IVF32,SQ8", use_cuvs=False)
def test_with_flag(self):
self.verify_clones_successfully("IVF32_HNSW,Flat",
allowCpuCoarseQuantizer=True)
self.verify_clones_successfully("IVF256(PQ2x4fs),Flat",
allowCpuCoarseQuantizer=True)
self.verify_clones_successfully(
"IVF32_HNSW,Flat", allowCpuCoarseQuantizer=True
)
self.verify_clones_successfully(
"IVF256(PQ2x4fs),Flat", allowCpuCoarseQuantizer=True
)
def test_with_flag_set_to_false(self):
try:
self.verify_clones_successfully("IVF32_HNSW,Flat",
allowCpuCoarseQuantizer=False)
self.verify_clones_successfully(
"IVF32_HNSW,Flat", allowCpuCoarseQuantizer=False
)
except Exception as e:
if "set the flag to true to allow the CPU fallback" not in str(e):
self.fail("Unexepected error message thrown: %s." % str(e))
faiss/gpu/test/test_lut_dtype_binding.py
+2 -2
View File
@@ -9,8 +9,8 @@ import faiss
@unittest.skipIf(
"CUVS" not in faiss.get_compile_options(),
"only if cuVS is compiled in")
"CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
)
class TestIVFPQSearchCagraConfigDtypes(unittest.TestCase):
"""Regression tests for the SWIG int typemap on cudaDataType_t.
faiss/gpu/test/test_multi_gpu.py
+17 -16
View File
@@ -21,8 +21,8 @@ class TestShardedFlat(unittest.TestCase):
nq = 200
k = 10
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
index_cpu = faiss.IndexFlatL2(d)
@@ -58,10 +58,11 @@ class TestShardedFlat(unittest.TestCase):
def do_test_sharded_ivf(self, index_key):
ds = SyntheticDataset(32, 8000, 10000, 100)
index = faiss.index_factory(ds.d, index_key)
if 'HNSW' in index_key:
if "HNSW" in index_key:
# make a bit more reproducible...
faiss.ParameterSpace().set_index_parameter(
index, 'quantizer_efSearch', 40)
index, "quantizer_efSearch", 40
)
index.train(ds.get_train())
index.add(ds.get_database())
Dref, Iref = index.search(ds.get_queries(), 10)
@@ -83,7 +84,7 @@ class TestShardedFlat(unittest.TestCase):
np.testing.assert_array_almost_equal(Dref, Dnew, decimal=4)
# the nprobe is taken from the sub-indexes
faiss.GpuParameterSpace().set_index_parameter(index, 'nprobe', 8)
faiss.GpuParameterSpace().set_index_parameter(index, "nprobe", 8)
Dnew8, Inew8 = index.search(ds.get_queries(), 10)
np.testing.assert_array_equal(Iref8, Inew8)
np.testing.assert_array_almost_equal(Dref8, Dnew8, decimal=4)
@@ -154,14 +155,13 @@ class EvalIVFPQAccuracy(unittest.TestCase):
qc = q[:d1, :]
def make_mat(n):
return np.dot(
np.random.random(size=(nb, d1)), qc).astype('float32')
return np.dot(np.random.random(size=(nb, d1)), qc).astype("float32")
return (make_mat(nt), make_mat(nb), make_mat(nq))
def test_mm(self):
# trouble with MKL+fbmake that appears only at runtime. Check it here
x = np.random.random(size=(100, 20)).astype('float32')
x = np.random.random(size=(100, 20)).astype("float32")
mat = faiss.PCAMatrix(20, 10)
mat.train(x)
mat.apply_py(x)
@@ -185,7 +185,7 @@ class EvalIVFPQAccuracy(unittest.TestCase):
# adding some ids because there was a bug in this case;
# those need to be cast to idx_t(= int64_t), because
# on windows the numpy int default is int32
ids = (np.arange(nb) * 3 + 12345).astype('int64')
ids = (np.arange(nb) * 3 + 12345).astype("int64")
index.add_with_ids(xb, ids)
ts.append(time.time())
@@ -203,15 +203,15 @@ class EvalIVFPQAccuracy(unittest.TestCase):
mem_info = res.getMemoryInfo()
assert isinstance(mem_info, dict)
assert isinstance(mem_info[0]['FlatData'], tuple)
assert isinstance(mem_info[0]['FlatData'][0], int)
assert isinstance(mem_info[0]['FlatData'][1], int)
assert isinstance(mem_info[0]["FlatData"], tuple)
assert isinstance(mem_info[0]["FlatData"][0], int)
assert isinstance(mem_info[0]["FlatData"][1], int)
gpu_index.nprobe = 4
Dnew, Inew = gpu_index.search(xq, 10)
ts.append(time.time())
print('times:', [t - ts[0] for t in ts])
print("times:", [t - ts[0] for t in ts])
# Give us some margin of error
self.assertGreaterEqual((Iref == Inew).sum(), Iref.size - 50)
@@ -230,7 +230,8 @@ class EvalIVFPQAccuracy(unittest.TestCase):
gpu_index = faiss.index_cpu_to_gpu_multiple_py(res, index, co)
faiss.GpuParameterSpace().set_index_parameter(
gpu_index, 'nprobe', 4)
gpu_index, "nprobe", 4
)
Dnew, Inew = gpu_index.search(xq, 10)
@@ -239,10 +240,10 @@ class EvalIVFPQAccuracy(unittest.TestCase):
self.assertGreaterEqual((Iref == Inew).sum(), Iref.size * 0.99)
def test_cpu_to_gpu_IVFPQ(self):
self.do_cpu_to_gpu('IVF128,PQ4')
self.do_cpu_to_gpu("IVF128,PQ4")
def test_cpu_to_gpu_IVFFlat(self):
self.do_cpu_to_gpu('IVF128,Flat')
self.do_cpu_to_gpu("IVF128,Flat")
def test_set_gpu_param(self):
index = faiss.index_factory(12, "PCAR8,IVF10,PQ4")
faiss/gpu/test/torch_test_contrib_gpu.py
+60 -25
View File
@@ -14,7 +14,7 @@ from faiss.contrib.torch import clustering
def to_column_major_torch(x):
if hasattr(torch, 'contiguous_format'):
if hasattr(torch, "contiguous_format"):
return x.t().clone(memory_format=torch.contiguous_format).t()
else:
# was default setting before memory_format was introduced
@@ -35,7 +35,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
cpu_index.add(xb_torch.numpy())
# Add to CPU index with torch GPU (should fail)
xb_torch_gpu = torch.rand(10000, 128, device=torch.device('cuda', 0), dtype=torch.float32)
xb_torch_gpu = torch.rand(
10000, 128, device=torch.device("cuda", 0), dtype=torch.float32
)
with self.assertRaises(AssertionError):
cpu_index.add(xb_torch_gpu)
@@ -59,8 +61,12 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(i_torch_cpu.cuda(), i_torch_gpu))
# Search with torch GPU using pre-allocated arrays
new_d_torch_gpu = torch.zeros(10, 10, device=torch.device('cuda', 0), dtype=torch.float32)
new_i_torch_gpu = torch.zeros(10, 10, device=torch.device('cuda', 0), dtype=torch.int64)
new_d_torch_gpu = torch.zeros(
10, 10, device=torch.device("cuda", 0), dtype=torch.float32
)
new_i_torch_gpu = torch.zeros(
10, 10, device=torch.device("cuda", 0), dtype=torch.int64
)
gpu_index.search(xq_torch_gpu, 10, new_d_torch_gpu, new_i_torch_gpu)
self.assertTrue(torch.equal(d_torch_cpu.cuda(), new_d_torch_gpu))
@@ -87,10 +93,17 @@ class TestTorchUtilsGPU(unittest.TestCase):
config.use_cuvs = False
index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, config)
xb = torch.rand(1000, d, device=torch.device('cuda', 0), dtype=torch.float32)
xb = torch.rand(
1000, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.train(xb)
ids = torch.arange(1000, 1000 + xb.shape[0], device=torch.device('cuda', 0), dtype=torch.int64)
ids = torch.arange(
1000,
1000 + xb.shape[0],
device=torch.device("cuda", 0),
dtype=torch.int64,
)
# Test add_with_ids with torch gpu
index.add_with_ids(xb, ids)
@@ -124,7 +137,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
res.noTempMemory()
index = faiss.GpuIndexFlatL2(res, d)
xb = torch.rand(100, d, device=torch.device('cuda', 0), dtype=torch.float32)
xb = torch.rand(
100, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.add(xb)
# Test reconstruct with torch gpu (native return)
@@ -133,7 +148,7 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[7], y))
# Test reconstruct with numpy output provided
y = np.empty(d, dtype='float32')
y = np.empty(d, dtype="float32")
index.reconstruct(11, y)
self.assertTrue(np.array_equal(xb.cpu().numpy()[11], y))
@@ -143,7 +158,7 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[12].cpu(), y))
# Test reconstruct with torch gpu output providesd
y = torch.empty(d, device=torch.device('cuda', 0), dtype=torch.float32)
y = torch.empty(d, device=torch.device("cuda", 0), dtype=torch.float32)
index.reconstruct(13, y)
self.assertTrue(torch.equal(xb[13], y))
@@ -153,7 +168,7 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[10:20], y))
# Test reconstruct with numpy output provided
y = np.empty((10, d), dtype='float32')
y = np.empty((10, d), dtype="float32")
index.reconstruct_n(20, 10, y)
self.assertTrue(np.array_equal(xb.cpu().numpy()[20:30], y))
@@ -163,7 +178,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[40:50].cpu(), y))
# Test reconstruct_n with torch gpu output provided
y = torch.empty(10, d, device=torch.device('cuda', 0), dtype=torch.float32)
y = torch.empty(
10, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.reconstruct_n(50, 10, y)
self.assertTrue(torch.equal(xb[50:60], y))
@@ -177,7 +194,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, config)
xb = torch.rand(100, d, device=torch.device('cuda', 0), dtype=torch.float32)
xb = torch.rand(
100, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.train(xb)
index.add(xb)
@@ -187,7 +206,7 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[10:20], y))
# Test reconstruct with numpy output provided
y = np.empty((10, d), dtype='float32')
y = np.empty((10, d), dtype="float32")
index.reconstruct_n(20, 10, y)
self.assertTrue(np.array_equal(xb.cpu().numpy()[20:30], y))
@@ -197,7 +216,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(torch.equal(xb[40:50].cpu(), y))
# Test reconstruct_n with torch gpu output provided
y = torch.empty(10, d, device=torch.device('cuda', 0), dtype=torch.float32)
y = torch.empty(
10, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.reconstruct_n(50, 10, y)
self.assertTrue(torch.equal(xb[50:60], y))
@@ -208,7 +229,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
res.noTempMemory()
index = faiss.GpuIndexFlatL2(res, d)
xb = torch.rand(10000, d, device=torch.device('cuda', 0), dtype=torch.float32)
xb = torch.rand(
10000, d, device=torch.device("cuda", 0), dtype=torch.float32
)
index.add(xb)
index_cpu = faiss.IndexFlatL2(d)
@@ -216,7 +239,9 @@ class TestTorchUtilsGPU(unittest.TestCase):
# Test assign with native gpu output
# both input as gpu torch and input as cpu torch
xq = torch.rand(10, d, device=torch.device('cuda', 0), dtype=torch.float32)
xq = torch.rand(
10, d, device=torch.device("cuda", 0), dtype=torch.float32
)
labels = index.assign(xq, 5)
labels_cpu = index_cpu.assign(xq.cpu(), 5)
@@ -228,7 +253,7 @@ class TestTorchUtilsGPU(unittest.TestCase):
self.assertTrue(np.array_equal(labels, labels_cpu))
# Test assign with numpy output provided
labels = np.empty((xq.shape[0], 5), dtype='int64')
labels = np.empty((xq.shape[0], 5), dtype="int64")
index.assign(xq.cpu().numpy(), 5, labels)
self.assertTrue(np.array_equal(labels, labels_cpu))
@@ -300,7 +325,9 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
D, I = faiss.knn_gpu(res, xq_c, xb_c, k, use_cuvs=use_cuvs)
self.assertTrue(torch.equal(torch.from_numpy(I), gt_I))
self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4)
self.assertLess(
(torch.from_numpy(D) - gt_D).abs().max(), 1e-4
)
# test torch (cpu, gpu) inputs
for is_cuda in True, False:
@@ -323,7 +350,9 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
xb_c = to_column_major_torch(xb)
assert not xb_c.is_contiguous()
D, I = faiss.knn_gpu(res, xq_c, xb_c, k, use_cuvs=use_cuvs)
D, I = faiss.knn_gpu(
res, xq_c, xb_c, k, use_cuvs=use_cuvs
)
self.assertTrue(torch.equal(I.cpu(), gt_I))
self.assertLess((D.cpu() - gt_D).abs().max(), 1e-4)
@@ -331,7 +360,9 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
# test on subset
try:
# This internally uses the current pytorch stream
D, I = faiss.knn_gpu(res, xq_c[6:8], xb_c, k, use_cuvs=use_cuvs)
D, I = faiss.knn_gpu(
res, xq_c[6:8], xb_c, k, use_cuvs=use_cuvs
)
except TypeError:
if not xq_row_major:
# then it is expected
@@ -343,8 +374,8 @@ class TestTorchUtilsKnnGpu(unittest.TestCase):
self.assertLess((D.cpu() - gt_D[6:8]).abs().max(), 1e-4)
@unittest.skipUnless(
"CUVS" in faiss.get_compile_options(),
"only if CUVS is compiled in")
"CUVS" in faiss.get_compile_options(), "only if CUVS is compiled in"
)
def test_knn_gpu_cuvs(self):
self.test_knn_gpu(use_cuvs=True)
@@ -434,7 +465,9 @@ class TestTorchUtilsPairwiseDistanceGpu(unittest.TestCase):
# do the same on our end
D = np.sort(D, axis=1)
self.assertLess((torch.from_numpy(D) - gt_D).abs().max(), 1e-4)
self.assertLess(
(torch.from_numpy(D) - gt_D).abs().max(), 1e-4
)
# test torch (cpu, gpu) inputs
for is_cuda in True, False:
@@ -468,7 +501,9 @@ class TestTorchUtilsPairwiseDistanceGpu(unittest.TestCase):
# test on subset
try:
# This internally uses the current pytorch stream
D = faiss.pairwise_distance_gpu(res, xq_c[4:8], xb_c)
D = faiss.pairwise_distance_gpu(
res, xq_c[4:8], xb_c
)
except TypeError:
if not xq_row_major:
# then it is expected
@@ -487,7 +522,7 @@ class TestTorchUtilsPairwiseDistanceGpu(unittest.TestCase):
class TestClustering(unittest.TestCase):
def test_python_kmeans(self):
""" Test the python implementation of kmeans """
"""Test the python implementation of kmeans"""
ds = datasets.SyntheticDataset(32, 10000, 0, 0)
x = ds.get_train()
faiss/gpu_metal/test/test_metal_python.py
+2 -1
View File
@@ -123,7 +123,8 @@ class TestMetalIVFFlat(unittest.TestCase):
quantizer = faiss.IndexFlatIP(d)
cpu_index = faiss.IndexIVFFlat(
quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT)
quantizer, d, nlist, faiss.METRIC_INNER_PRODUCT
)
cpu_index.nprobe = nprobe
cpu_index.train(xb)
cpu_index.add(xb)
faiss/python/__init__.py
+3 -1
View File
@@ -390,7 +390,9 @@ def search_with_parameters(index, x, k, params=None, output_stats=False):
range_search_with_parameters_c = range_search_with_parameters
def range_search_with_parameters(index, x, radius, params=None, output_stats=False):
def range_search_with_parameters(
index, x, radius, params=None, output_stats=False
):
x = np.ascontiguousarray(x, dtype="float32")
n, d = x.shape
assert d == index.d
faiss/python/array_conversions.py
+52 -43
View File
@@ -19,6 +19,7 @@ from faiss.loader import *
# classes from the SWIG interface
###########################################
def _make_deprecated_swig_class(deprecated_name, base_name):
"""
Dynamically construct deprecated classes as wrappers around renamed ones
@@ -62,58 +63,59 @@ def _make_deprecated_swig_class(deprecated_name, base_name):
# numpy array / std::vector conversions
###########################################
sizeof_long = array.array('l').itemsize
sizeof_long = array.array("l").itemsize
deprecated_name_map = {
# deprecated: replacement
'Float': 'Float32',
'Double': 'Float64',
'Char': 'Int8',
'Int': 'Int32',
'Long': 'Int32' if sizeof_long == 4 else 'Int64',
'LongLong': 'Int64',
'Byte': 'UInt8',
"Float": "Float32",
"Double": "Float64",
"Char": "Int8",
"Int": "Int32",
"Long": "Int32" if sizeof_long == 4 else "Int64",
"LongLong": "Int64",
"Byte": "UInt8",
# previously misspelled variant
'Uint64': 'UInt64',
"Uint64": "UInt64",
}
for depr_prefix, base_prefix in deprecated_name_map.items():
_make_deprecated_swig_class(depr_prefix + "Vector", base_prefix + "Vector")
# same for the three legacy *VectorVector classes
if depr_prefix in ['Float', 'Long', 'Byte']:
_make_deprecated_swig_class(depr_prefix + "VectorVector",
base_prefix + "VectorVector")
if depr_prefix in ["Float", "Long", "Byte"]:
_make_deprecated_swig_class(
depr_prefix + "VectorVector", base_prefix + "VectorVector"
)
# mapping from vector names in swigfaiss.swig and the numpy dtype names
# TODO: once deprecated classes are removed, remove the dict and just use .lower() below
vector_name_map = {
'Float32': 'float32',
'Float64': 'float64',
'Int8': 'int8',
'Int16': 'int16',
'Int32': 'int32',
'Int64': 'int64',
'UInt8': 'uint8',
'UInt16': 'uint16',
'UInt32': 'uint32',
'UInt64': 'uint64',
**{k: v.lower() for k, v in deprecated_name_map.items()}
"Float32": "float32",
"Float64": "float64",
"Int8": "int8",
"Int16": "int16",
"Int32": "int32",
"Int64": "int64",
"UInt8": "uint8",
"UInt16": "uint16",
"UInt32": "uint32",
"UInt64": "uint64",
**{k: v.lower() for k, v in deprecated_name_map.items()},
}
def vector_to_array(v):
""" convert a C++ vector to a numpy array """
"""convert a C++ vector to a numpy array"""
classname = v.__class__.__name__
if classname.startswith('AlignedTable'):
if classname.startswith("AlignedTable"):
return AlignedTable_to_array(v)
if classname.startswith('MaybeOwnedVector'):
if classname.startswith("MaybeOwnedVector"):
dtype = np.dtype(vector_name_map[classname[16:]])
a = np.empty(v.size(), dtype=dtype)
if v.size() > 0:
memcpy(swig_ptr(a), v.data(), a.nbytes)
return a
assert classname.endswith('Vector')
assert classname.endswith("Vector")
dtype = np.dtype(vector_name_map[classname[:-6]])
a = np.empty(v.size(), dtype=dtype)
if v.size() > 0:
@@ -126,34 +128,41 @@ def vector_float_to_array(v):
def copy_array_to_vector(a, v):
""" copy a numpy array to a vector """
n, = a.shape
"""copy a numpy array to a vector"""
(n,) = a.shape
classname = v.__class__.__name__
if classname.startswith('MaybeOwnedVector'):
assert v.is_owned, 'cannot copy to an non-owned MaybeOwnedVector'
if classname.startswith("MaybeOwnedVector"):
assert v.is_owned, "cannot copy to an non-owned MaybeOwnedVector"
dtype = np.dtype(vector_name_map[classname[16:]])
assert dtype == a.dtype, (
'cannot copy a %s array to a %s (should be %s)' % (
a.dtype, classname, dtype))
assert (
dtype == a.dtype
), "cannot copy a %s array to a %s (should be %s)" % (
a.dtype,
classname,
dtype,
)
v.resize(n)
if n > 0:
memcpy(v.data(), swig_ptr(a), a.nbytes)
return
assert classname.endswith('Vector')
assert classname.endswith("Vector")
dtype = np.dtype(vector_name_map[classname[:-6]])
assert dtype == a.dtype, (
'cannot copy a %s array to a %s (should be %s)' % (
a.dtype, classname, dtype))
assert dtype == a.dtype, "cannot copy a %s array to a %s (should be %s)" % (
a.dtype,
classname,
dtype,
)
v.resize(n)
if n > 0:
memcpy(v.data(), swig_ptr(a), a.nbytes)
# same for AlignedTable
def copy_array_to_AlignedTable(a, v):
n, = a.shape
(n,) = a.shape
# TODO check class name
assert v.itemsize() == a.itemsize
v.resize(n)
@@ -162,9 +171,9 @@ def copy_array_to_AlignedTable(a, v):
def array_to_AlignedTable(a):
if a.dtype == 'uint16':
if a.dtype == "uint16":
v = AlignedTableUint16(a.size)
elif a.dtype == 'uint8':
elif a.dtype == "uint8":
v = AlignedTableUint8(a.size)
else:
assert False
@@ -173,9 +182,9 @@ def array_to_AlignedTable(a):
def AlignedTable_to_array(v):
""" convert an AlignedTable to a numpy array """
"""convert an AlignedTable to a numpy array"""
classname = v.__class__.__name__
assert classname.startswith('AlignedTable')
assert classname.startswith("AlignedTable")
dtype = classname[12:].lower()
a = np.empty(v.size(), dtype=dtype)
if a.size > 0:
faiss/python/class_wrappers.py
+293 -205
View File
File diff suppressed because it is too large Load Diff
faiss/python/extra_wrappers.py
+160 -114
View File
@@ -25,10 +25,10 @@ import collections.abc
def kmin(array, k):
"""return k smallest values (and their indices) of the lines of a
float32 array"""
array = np.ascontiguousarray(array, dtype='float32')
array = np.ascontiguousarray(array, dtype="float32")
m, n = array.shape
I = np.zeros((m, k), dtype='int64')
D = np.zeros((m, k), dtype='float32')
I = np.zeros((m, k), dtype="int64")
D = np.zeros((m, k), dtype="float32")
ha = faiss.float_maxheap_array_t()
ha.ids = swig_ptr(I)
ha.val = swig_ptr(D)
@@ -43,10 +43,10 @@ def kmin(array, k):
def kmax(array, k):
"""return k largest values (and their indices) of the lines of a
float32 array"""
array = np.ascontiguousarray(array, dtype='float32')
array = np.ascontiguousarray(array, dtype="float32")
m, n = array.shape
I = np.zeros((m, k), dtype='int64')
D = np.zeros((m, k), dtype='float32')
I = np.zeros((m, k), dtype="int64")
D = np.zeros((m, k), dtype="float32")
ha = faiss.float_minheap_array_t()
ha.ids = swig_ptr(I)
ha.val = swig_ptr(D)
@@ -61,36 +61,38 @@ def kmax(array, k):
def pairwise_distances(xq, xb, metric=METRIC_L2, metric_arg=0):
"""compute the whole pairwise distance matrix between two sets of
vectors"""
xq = np.ascontiguousarray(xq, dtype='float32')
xb = np.ascontiguousarray(xb, dtype='float32')
xq = np.ascontiguousarray(xq, dtype="float32")
xb = np.ascontiguousarray(xb, dtype="float32")
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
dis = np.empty((nq, nb), dtype='float32')
dis = np.empty((nq, nb), dtype="float32")
if metric == METRIC_L2:
pairwise_L2sqr(
d, nq, swig_ptr(xq),
nb, swig_ptr(xb),
swig_ptr(dis))
pairwise_L2sqr(d, nq, swig_ptr(xq), nb, swig_ptr(xb), swig_ptr(dis))
elif metric == METRIC_INNER_PRODUCT:
dis[:] = xq @ xb.T
else:
pairwise_extra_distances(
d, nq, swig_ptr(xq),
nb, swig_ptr(xb),
metric, metric_arg,
swig_ptr(dis))
d,
nq,
swig_ptr(xq),
nb,
swig_ptr(xb),
metric,
metric_arg,
swig_ptr(dis),
)
return dis
def rand(n, seed=12345):
res = np.empty(n, dtype='float32')
res = np.empty(n, dtype="float32")
float_rand(swig_ptr(res), res.size, seed)
return res
def randint(n, seed=12345, vmax=None):
res = np.empty(n, dtype='int64')
res = np.empty(n, dtype="int64")
if vmax is None:
int64_rand(swig_ptr(res), res.size, seed)
else:
@@ -102,48 +104,53 @@ lrand = randint
def randn(n, seed=12345):
res = np.empty(n, dtype='float32')
res = np.empty(n, dtype="float32")
float_randn(swig_ptr(res), res.size, seed)
return res
def checksum(a):
""" compute a checksum for quick-and-dirty comparisons of arrays """
a = a.view('uint8')
"""compute a checksum for quick-and-dirty comparisons of arrays"""
a = a.view("uint8")
if a.ndim == 1:
return bvec_checksum(a.size, swig_ptr(a))
n, d = a.shape
cs = np.zeros(n, dtype='uint64')
cs = np.zeros(n, dtype="uint64")
bvecs_checksum(n, d, swig_ptr(a), swig_ptr(cs))
return cs
rand_smooth_vectors_c = rand_smooth_vectors
def rand_smooth_vectors(n, d, seed=1234):
res = np.empty((n, d), dtype='float32')
res = np.empty((n, d), dtype="float32")
rand_smooth_vectors_c(n, d, swig_ptr(res), seed)
return res
def eval_intersection(I1, I2):
""" size of intersection between each line of two result tables"""
I1 = np.ascontiguousarray(I1, dtype='int64')
I2 = np.ascontiguousarray(I2, dtype='int64')
"""size of intersection between each line of two result tables"""
I1 = np.ascontiguousarray(I1, dtype="int64")
I2 = np.ascontiguousarray(I2, dtype="int64")
n = I1.shape[0]
assert I2.shape[0] == n
k1, k2 = I1.shape[1], I2.shape[1]
ninter = 0
for i in range(n):
ninter += ranklist_intersection_size(
k1, swig_ptr(I1[i]), k2, swig_ptr(I2[i]))
k1, swig_ptr(I1[i]), k2, swig_ptr(I2[i])
)
return ninter
def normalize_L2(x):
fvec_renorm_L2(x.shape[1], x.shape[0], swig_ptr(x))
bucket_sort_c = bucket_sort
def bucket_sort(tab, nbucket=None, nt=0):
"""Perform a bucket sort on a table of integers.
@@ -166,17 +173,22 @@ def bucket_sort(tab, nbucket=None, nt=0):
tab = np.ascontiguousarray(tab, dtype="int64")
if nbucket is None:
nbucket = int(tab.max() + 1)
lims = np.empty(nbucket + 1, dtype='int64')
perm = np.empty(tab.size, dtype='int64')
lims = np.empty(nbucket + 1, dtype="int64")
perm = np.empty(tab.size, dtype="int64")
bucket_sort_c(
tab.size, faiss.swig_ptr(tab.view('uint64')),
nbucket, faiss.swig_ptr(lims), faiss.swig_ptr(perm),
nt
tab.size,
faiss.swig_ptr(tab.view("uint64")),
nbucket,
faiss.swig_ptr(lims),
faiss.swig_ptr(perm),
nt,
)
return lims, perm
matrix_bucket_sort_inplace_c = matrix_bucket_sort_inplace
def matrix_bucket_sort_inplace(tab, nbucket=None, nt=0):
"""Perform a bucket sort on a matrix, recording the original
row of each element.
@@ -199,15 +211,13 @@ def matrix_bucket_sort_inplace(tab, nbucket=None, nt=0):
lims : array_like
cumulative sum of bucket sizes (size vmax + 1)
"""
assert tab.dtype == 'int32' or tab.dtype == 'int64'
assert tab.dtype == "int32" or tab.dtype == "int64"
nrow, ncol = tab.shape
if nbucket is None:
nbucket = int(tab.max() + 1)
lims = np.empty(nbucket + 1, dtype='int64')
lims = np.empty(nbucket + 1, dtype="int64")
matrix_bucket_sort_inplace_c(
nrow, ncol, faiss.swig_ptr(tab),
nbucket, faiss.swig_ptr(lims),
nt
nrow, ncol, faiss.swig_ptr(tab), nbucket, faiss.swig_ptr(lims), nt
)
return lims
@@ -216,6 +226,7 @@ def matrix_bucket_sort_inplace(tab, nbucket=None, nt=0):
# ResultHeap
###########################################
class ResultHeap:
"""Accumulate query results from a sliced dataset. The final result will
be in self.D, self.I."""
@@ -226,8 +237,8 @@ class ResultHeap:
k: number of results per query
keep_max: keep the top-k maximum values instead of the minima
"""
self.I = np.zeros((nq, k), dtype='int64')
self.D = np.zeros((nq, k), dtype='float32')
self.I = np.zeros((nq, k), dtype="int64")
self.D = np.zeros((nq, k), dtype="float32")
self.nq, self.k = nq, k
if keep_max:
heaps = float_minheap_array_t()
@@ -247,13 +258,11 @@ class ResultHeap:
D, I do not need to be in a particular order (heap or sorted)
"""
nq, kd = D.shape
D = np.ascontiguousarray(D, dtype='float32')
I = np.ascontiguousarray(I, dtype='int64')
D = np.ascontiguousarray(D, dtype="float32")
I = np.ascontiguousarray(I, dtype="int64")
assert I.shape == (nq, kd)
assert nq == self.nq
self.heaps.addn_with_ids(
kd, swig_ptr(D),
swig_ptr(I), kd)
self.heaps.addn_with_ids(kd, swig_ptr(D), swig_ptr(I), kd)
def add_result_subset(self, subset, D, I):
"""
@@ -264,16 +273,17 @@ class ResultHeap:
nsubset, kd = D.shape
assert nsubset == len(subset)
assert (
I.ndim == 2 and D.shape == I.shape or
I.ndim == 1 and I.shape == (kd, )
I.ndim == 2
and D.shape == I.shape
or I.ndim == 1
and I.shape == (kd,)
)
D = np.ascontiguousarray(D, dtype='float32')
I = np.ascontiguousarray(I, dtype='int64')
subset = np.ascontiguousarray(subset, dtype='int64')
D = np.ascontiguousarray(D, dtype="float32")
I = np.ascontiguousarray(I, dtype="int64")
subset = np.ascontiguousarray(subset, dtype="int64")
id_stride = 0 if I.ndim == 1 else kd
self.heaps.addn_query_subset_with_ids(
nsubset, swig_ptr(subset),
kd, swig_ptr(D), swig_ptr(I), id_stride
nsubset, swig_ptr(subset), kd, swig_ptr(D), swig_ptr(I), id_stride
)
def finalize(self):
@@ -292,44 +302,62 @@ def merge_knn_results(Dall, Iall, keep_max=False):
Inew = np.empty((n, k), dtype=Iall.dtype)
func = merge_knn_results_CMax if keep_max else merge_knn_results_CMin
func(
n, k, nshard,
swig_ptr(Dall), swig_ptr(Iall),
swig_ptr(Dnew), swig_ptr(Inew)
n,
k,
nshard,
swig_ptr(Dall),
swig_ptr(Iall),
swig_ptr(Dnew),
swig_ptr(Inew),
)
return Dnew, Inew
######################################################
# Efficient ID to ID map
######################################################
class MapInt64ToInt64:
def __init__(self, capacity):
self.log2_capacity = int(np.log2(capacity))
assert capacity == 2 ** self.log2_capacity, "need power of 2 capacity"
assert capacity == 2**self.log2_capacity, "need power of 2 capacity"
self.capacity = capacity
self.tab = np.empty((capacity, 2), dtype='int64')
faiss.hashtable_int64_to_int64_init(self.log2_capacity, swig_ptr(self.tab))
self.tab = np.empty((capacity, 2), dtype="int64")
faiss.hashtable_int64_to_int64_init(
self.log2_capacity, swig_ptr(self.tab)
)
def add(self, keys, vals):
n, = keys.shape
(n,) = keys.shape
assert vals.shape == (n,)
faiss.hashtable_int64_to_int64_add(
self.log2_capacity, swig_ptr(self.tab),
n, swig_ptr(keys), swig_ptr(vals))
self.log2_capacity,
swig_ptr(self.tab),
n,
swig_ptr(keys),
swig_ptr(vals),
)
def lookup(self, keys):
n, = keys.shape
vals = np.empty((n,), dtype='int64')
(n,) = keys.shape
vals = np.empty((n,), dtype="int64")
faiss.hashtable_int64_to_int64_lookup(
self.log2_capacity, swig_ptr(self.tab),
n, swig_ptr(keys), swig_ptr(vals))
self.log2_capacity,
swig_ptr(self.tab),
n,
swig_ptr(keys),
swig_ptr(vals),
)
return vals
######################################################
# KNN function
######################################################
def knn(xq, xb, k, metric=METRIC_L2, metric_arg=0.0):
"""
Compute the k nearest neighbors of a vector without constructing an index
@@ -355,30 +383,35 @@ def knn(xq, xb, k, metric=METRIC_L2, metric_arg=0.0):
I : array_like
Labels of the nearest neighbors, shape (nq, k)
"""
xq = np.ascontiguousarray(xq, dtype='float32')
xb = np.ascontiguousarray(xb, dtype='float32')
xq = np.ascontiguousarray(xq, dtype="float32")
xb = np.ascontiguousarray(xb, dtype="float32")
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
I = np.empty((nq, k), dtype='int64')
D = np.empty((nq, k), dtype='float32')
I = np.empty((nq, k), dtype="int64")
D = np.empty((nq, k), dtype="float32")
if metric == METRIC_L2:
knn_L2sqr(
swig_ptr(xq), swig_ptr(xb),
d, nq, nb, k, swig_ptr(D), swig_ptr(I)
swig_ptr(xq), swig_ptr(xb), d, nq, nb, k, swig_ptr(D), swig_ptr(I)
)
elif metric == METRIC_INNER_PRODUCT:
knn_inner_product(
swig_ptr(xq), swig_ptr(xb),
d, nq, nb, k, swig_ptr(D), swig_ptr(I)
swig_ptr(xq), swig_ptr(xb), d, nq, nb, k, swig_ptr(D), swig_ptr(I)
)
else:
else:
knn_extra_metrics(
swig_ptr(xq), swig_ptr(xb),
d, nq, nb, metric, metric_arg, k,
swig_ptr(D), swig_ptr(I)
swig_ptr(xq),
swig_ptr(xb),
d,
nq,
nb,
metric,
metric_arg,
k,
swig_ptr(D),
swig_ptr(I),
)
return D, I
@@ -412,8 +445,8 @@ def knn_hamming(xq, xb, k, variant="hc"):
nq, d = xq.shape
nb, d2 = xb.shape
assert d == d2
D = np.empty((nq, k), dtype='int32')
I = np.empty((nq, k), dtype='int64')
D = np.empty((nq, k), dtype="int32")
I = np.empty((nq, k), dtype="int64")
if variant == "hc":
heap = faiss.int_maxheap_array_t()
@@ -422,13 +455,18 @@ def knn_hamming(xq, xb, k, variant="hc"):
heap.ids = faiss.swig_ptr(I)
heap.val = faiss.swig_ptr(D)
faiss.hammings_knn_hc(
heap, faiss.swig_ptr(xq), faiss.swig_ptr(xb), nb,
d, 1
heap, faiss.swig_ptr(xq), faiss.swig_ptr(xb), nb, d, 1
)
elif variant == "mc":
faiss.hammings_knn_mc(
faiss.swig_ptr(xq), faiss.swig_ptr(xb), nq, nb, k, d,
faiss.swig_ptr(D), faiss.swig_ptr(I)
faiss.swig_ptr(xq),
faiss.swig_ptr(xb),
nq,
nb,
k,
d,
faiss.swig_ptr(D),
faiss.swig_ptr(I),
)
else:
raise NotImplementedError
@@ -487,8 +525,8 @@ class Kmeans:
def __init__(self, d, k, **kwargs):
"""d: input dimension, k: nb of centroids. Additional
parameters are passed on the ClusteringParameters object,
including niter=25, verbose=False, spherical = False
parameters are passed on the ClusteringParameters object,
including niter=25, verbose=False, spherical = False
"""
self.d = d
self.reset(k)
@@ -498,7 +536,7 @@ class Kmeans:
else:
self.cp = ClusteringParameters()
for k, v in kwargs.items():
if k == 'gpu':
if k == "gpu":
if v == True or v == -1:
v = get_num_gpus()
self.gpu = v
@@ -516,7 +554,9 @@ class Kmeans:
else:
self.index = IndexFlatL2(d)
if self.gpu:
self.index = faiss.index_cpu_to_all_gpus(self.index, ngpu=self.gpu)
self.index = faiss.index_cpu_to_all_gpus(
self.index, ngpu=self.gpu
)
else:
if self.gpu:
fac = GpuProgressiveDimIndexFactory(ngpu=self.gpu)
@@ -525,8 +565,8 @@ class Kmeans:
self.fac = fac
def reset(self, k=None):
""" prepare k-means object to perform a new clustering, possibly
with another number of centroids """
"""prepare k-means object to perform a new clustering, possibly
with another number of centroids"""
if k is not None:
self.k = int(k)
self.centroids = None
@@ -534,7 +574,7 @@ class Kmeans:
self.iteration_stats = None
def train(self, x, weights=None, init_centroids=None):
""" Perform k-means clustering.
"""Perform k-means clustering.
On output of the function call:
- the centroids are in the centroids field of size (`k`, `d`).
@@ -559,7 +599,7 @@ class Kmeans:
final optimization objective
"""
x = np.ascontiguousarray(x, dtype='float32')
x = np.ascontiguousarray(x, dtype="float32")
n, d = x.shape
assert d == self.d
@@ -569,7 +609,9 @@ class Kmeans:
if init_centroids is not None:
nc, d2 = init_centroids.shape
assert d2 == d
faiss.copy_array_to_vector(init_centroids.ravel(), clus.centroids)
faiss.copy_array_to_vector(
init_centroids.ravel(), clus.centroids
)
clus.train(x, self.index, weights)
else:
# not supported for progressive dim
@@ -586,15 +628,14 @@ class Kmeans:
stats = [stats.at(i) for i in range(stats.size())]
self.obj = np.array([st.obj for st in stats])
# copy all the iteration_stats objects to a python array
stat_fields = 'obj time time_search imbalance_factor nsplit'.split()
stat_fields = "obj time time_search imbalance_factor nsplit".split()
self.iteration_stats = [
{field: getattr(st, field) for field in stat_fields}
for st in stats
{field: getattr(st, field) for field in stat_fields} for st in stats
]
return self.obj[-1] if self.obj.size > 0 else 0.0
def assign(self, x):
x = np.ascontiguousarray(x, dtype='float32')
x = np.ascontiguousarray(x, dtype="float32")
assert self.centroids is not None, "should train before assigning"
self.index.reset()
self.index.add(self.centroids)
@@ -629,9 +670,10 @@ class SuperKmeans(Kmeans):
def train(self, x, weights=None, init_centroids=None):
assert weights is None, "SuperKmeans does not support weights"
assert init_centroids is None, \
"SuperKmeans does not support init_centroids"
x = np.ascontiguousarray(x, dtype='float32')
assert (
init_centroids is None
), "SuperKmeans does not support init_centroids"
x = np.ascontiguousarray(x, dtype="float32")
n, d = x.shape
assert d == self.d
@@ -643,10 +685,9 @@ class SuperKmeans(Kmeans):
stats = sc.iteration_stats
stats = [stats.at(i) for i in range(stats.size())]
self.obj = np.array([st.obj for st in stats])
stat_fields = 'obj time time_search imbalance_factor nsplit'.split()
stat_fields = "obj time time_search imbalance_factor nsplit".split()
self.iteration_stats = [
{field: getattr(st, field) for field in stat_fields}
for st in stats
{field: getattr(st, field) for field in stat_fields} for st in stats
]
self.gemm_pruning_rates = faiss.vector_to_array(sc.gemm_pruning_rates)
return self.obj[-1] if self.obj.size > 0 else 0.0
@@ -656,11 +697,14 @@ class SuperKmeans(Kmeans):
# Packing and unpacking bitstrings
###########################################
def is_sequence(x):
return isinstance(x, collections.abc.Sequence)
pack_bitstrings_c = pack_bitstrings
def pack_bitstrings(a, nbit):
"""
Pack a set integers (i, j) where i=0:n and j=0:M into
@@ -672,22 +716,25 @@ def pack_bitstrings(a, nbit):
If nbit is an array: entry (i, j) takes nbit[j] bits.
"""
n, M = a.shape
a = np.ascontiguousarray(a, dtype='int32')
a = np.ascontiguousarray(a, dtype="int32")
if is_sequence(nbit):
nbit = np.ascontiguousarray(nbit, dtype='int32')
nbit = np.ascontiguousarray(nbit, dtype="int32")
assert nbit.shape == (M,)
code_size = int((nbit.sum() + 7) // 8)
b = np.empty((n, code_size), dtype='uint8')
b = np.empty((n, code_size), dtype="uint8")
pack_bitstrings_c(
n, M, swig_ptr(nbit), swig_ptr(a), swig_ptr(b), code_size)
n, M, swig_ptr(nbit), swig_ptr(a), swig_ptr(b), code_size
)
else:
code_size = (M * nbit + 7) // 8
b = np.empty((n, code_size), dtype='uint8')
b = np.empty((n, code_size), dtype="uint8")
pack_bitstrings_c(n, M, nbit, swig_ptr(a), swig_ptr(b), code_size)
return b
unpack_bitstrings_c = unpack_bitstrings
def unpack_bitstrings(b, M_or_nbits, nbit=None):
"""
Unpack a set integers (i, j) where i=0:n and j=0:M from
@@ -703,19 +750,18 @@ def unpack_bitstrings(b, M_or_nbits, nbit=None):
"""
n, code_size = b.shape
if nbit is None:
nbit = np.ascontiguousarray(M_or_nbits, dtype='int32')
nbit = np.ascontiguousarray(M_or_nbits, dtype="int32")
M = len(nbit)
min_code_size = int((nbit.sum() + 7) // 8)
assert code_size >= min_code_size
a = np.empty((n, M), dtype='int32')
a = np.empty((n, M), dtype="int32")
unpack_bitstrings_c(
n, M, swig_ptr(nbit),
swig_ptr(b), code_size, swig_ptr(a))
n, M, swig_ptr(nbit), swig_ptr(b), code_size, swig_ptr(a)
)
else:
M = M_or_nbits
min_code_size = (M * nbit + 7) // 8
assert code_size >= min_code_size
a = np.empty((n, M), dtype='int32')
unpack_bitstrings_c(
n, M, nbit, swig_ptr(b), code_size, swig_ptr(a))
a = np.empty((n, M), dtype="int32")
unpack_bitstrings_c(n, M, nbit, swig_ptr(b), code_size, swig_ptr(a))
return a
faiss/python/gpu_wrappers.py
+27 -14
View File
@@ -19,9 +19,9 @@ from faiss.loader import *
def index_cpu_to_gpu_multiple_py(resources, index, co=None, gpus=None):
""" builds the C++ vectors for the GPU indices and the
"""builds the C++ vectors for the GPU indices and the
resources. Handles the case where the resources are assigned to
the list of GPUs """
the list of GPUs"""
if gpus is None:
gpus = range(len(resources))
vres = GpuResourcesVector()
@@ -41,7 +41,7 @@ def index_cpu_to_all_gpus(index, co=None, ngpu=-1):
def index_cpu_to_gpus_list(index, co=None, gpus=None, ngpu=-1):
""" Here we can pass list of GPU ids as a parameter or ngpu to
"""Here we can pass list of GPU ids as a parameter or ngpu to
use first n GPU's. gpus mut be a list or None.
co is a GpuMultipleClonerOptions
"""
@@ -53,10 +53,23 @@ def index_cpu_to_gpus_list(index, co=None, gpus=None, ngpu=-1):
index_gpu = index_cpu_to_gpu_multiple_py(res, index, co, gpus)
return index_gpu
# allows numpy ndarray usage with bfKnn
def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuvs=False, vectorsMemoryLimit=0, queriesMemoryLimit=0):
def knn_gpu(
res,
xq,
xb,
k,
D=None,
I=None,
metric=METRIC_L2,
device=-1,
use_cuvs=False,
vectorsMemoryLimit=0,
queriesMemoryLimit=0,
):
"""
Compute the k nearest neighbors of a vector on one GPU without constructing an index
@@ -107,7 +120,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
xq = xq.T
xq_row_major = False
else:
xq = np.ascontiguousarray(xq, dtype='float32')
xq = np.ascontiguousarray(xq, dtype="float32")
xq_row_major = True
xq_ptr = swig_ptr(xq)
@@ -117,7 +130,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
elif xq.dtype == np.float16:
xq_type = DistanceDataType_F16
else:
raise TypeError('xq must be f32 or f16')
raise TypeError("xq must be f32 or f16")
nb, d2 = xb.shape
assert d2 == d
@@ -127,7 +140,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
xb = xb.T
xb_row_major = False
else:
xb = np.ascontiguousarray(xb, dtype='float32')
xb = np.ascontiguousarray(xb, dtype="float32")
xb_row_major = True
xb_ptr = swig_ptr(xb)
@@ -137,7 +150,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
elif xb.dtype == np.float16:
xb_type = DistanceDataType_F16
else:
raise TypeError('xb must be float32 or float16')
raise TypeError("xb must be float32 or float16")
if D is None:
D = np.empty((nq, k), dtype=np.float32)
@@ -160,7 +173,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
elif I.dtype == I.dtype == np.int32:
I_type = IndicesDataType_I32
else:
raise TypeError('I must be i64 or i32')
raise TypeError("I must be i64 or i32")
args = GpuDistanceParams()
args.metric = metric
@@ -189,6 +202,7 @@ def knn_gpu(res, xq, xb, k, D=None, I=None, metric=METRIC_L2, device=-1, use_cuv
return D, I
# allows numpy ndarray usage with bfKnn for all pairwise distances
@@ -229,8 +243,7 @@ def pairwise_distance_gpu(res, xq, xb, D=None, metric=METRIC_L2, device=-1):
xq = xq.T
xq_row_major = False
else:
raise TypeError(
'xq matrix should be row (C) or column-major (Fortran)')
raise TypeError("xq matrix should be row (C) or column-major (Fortran)")
xq_ptr = swig_ptr(xq)
@@ -239,7 +252,7 @@ def pairwise_distance_gpu(res, xq, xb, D=None, metric=METRIC_L2, device=-1):
elif xq.dtype == np.float16:
xq_type = DistanceDataType_F16
else:
xq = np.ascontiguousarray(xb, dtype='float32')
xq = np.ascontiguousarray(xb, dtype="float32")
xq_row_major = True
nb, d2 = xb.shape
@@ -250,7 +263,7 @@ def pairwise_distance_gpu(res, xq, xb, D=None, metric=METRIC_L2, device=-1):
xb = xb.T
xb_row_major = False
else:
xb = np.ascontiguousarray(xb, dtype='float32')
xb = np.ascontiguousarray(xb, dtype="float32")
xb_row_major = True
xb_ptr = swig_ptr(xb)
@@ -260,7 +273,7 @@ def pairwise_distance_gpu(res, xq, xb, D=None, metric=METRIC_L2, device=-1):
elif xb.dtype == np.float16:
xb_type = DistanceDataType_F16
else:
raise TypeError('xb must be float32 or float16')
raise TypeError("xb must be float32 or float16")
if D is None:
D = np.empty((nq, nb), dtype=np.float32)
faiss/python/loader.py
+41 -11
View File
@@ -37,23 +37,34 @@ def supported_instruction_sets():
return False
# Numpy 2.0 supports SVE detection by __cpu_features__, so just skip
import numpy
if Version(numpy.__version__) >= Version("2.0"):
return False
# platform-dependent legacy fallback using numpy.distutils.cpuinfo
try:
import numpy.distutils.cpuinfo
return "sve" in numpy.distutils.cpuinfo.cpu.info[0].get('Features', "").split()
return (
"sve"
in numpy.distutils.cpuinfo.cpu.info[0]
.get("Features", "")
.split()
)
except ImportError:
# check if SVE is supported by checking the auxval
# using values defined as:
# #define AT_HWCAP 16
# #define HWCAP_SVE (1 << 22)
return bool(__import__('ctypes').CDLL(None).getauxval(16) & (1<<22))
return bool(
__import__("ctypes").CDLL(None).getauxval(16) & (1 << 22)
)
import numpy
if Version(numpy.__version__) >= Version("1.19"):
# use private API as next-best thing until numpy/numpy#18058 is solved
from numpy._core._multiarray_umath import __cpu_features__
# __cpu_features__ is a dictionary with CPU features
# as keys, and True / False as values
supported = {k for k, v in __cpu_features__.items() if v}
@@ -65,16 +76,24 @@ def supported_instruction_sets():
# platform-dependent legacy fallback before numpy 1.19, no windows
if platform.system() == "Darwin":
if subprocess.check_output(["/usr/sbin/sysctl", "hw.optional.avx2_0"])[-1] == '1':
if (
subprocess.check_output(["/usr/sbin/sysctl", "hw.optional.avx2_0"])[
-1
]
== "1"
):
return {"AVX2"}
elif platform.system() == "Linux":
import numpy.distutils.cpuinfo
result = set()
if "avx2" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
if "avx2" in numpy.distutils.cpuinfo.cpu.info[0].get("flags", ""):
result.add("AVX2")
if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
if "avx512" in numpy.distutils.cpuinfo.cpu.info[0].get("flags", ""):
result.add("AVX512")
if "avx512_fp16" in numpy.distutils.cpuinfo.cpu.info[0].get('flags', ""):
if "avx512_fp16" in numpy.distutils.cpuinfo.cpu.info[0].get(
"flags", ""
):
# avx512_fp16 is supported starting SPR
result.add("AVX512_SPR")
if is_sve_supported():
@@ -84,6 +103,7 @@ def supported_instruction_sets():
return result
return set()
logger = logging.getLogger(__name__)
instruction_sets = None
@@ -93,8 +113,10 @@ opt_env_variable_name = "FAISS_OPT_LEVEL"
opt_level = os.environ.get(opt_env_variable_name, None)
if opt_level is None:
logger.debug(f"Environment variable {opt_env_variable_name} is not set, " \
"so let's pick the instruction set according to the current CPU")
logger.debug(
f"Environment variable {opt_env_variable_name} is not set, "
"so let's pick the instruction set according to the current CPU"
)
instruction_sets = supported_instruction_sets()
else:
logger.debug(f"Using {opt_level} as an instruction set.")
@@ -107,10 +129,13 @@ if has_AVX512_SPR:
try:
logger.info("Loading faiss with AVX512-SPR support.")
from .swigfaiss_avx512_spr import *
logger.info("Successfully loaded faiss with AVX512-SPR support.")
loaded = True
except ImportError as e:
logger.info(f"Could not load library with AVX512-SPR support due to:\n{e!r}")
logger.info(
f"Could not load library with AVX512-SPR support due to:\n{e!r}"
)
# reset so that we load without AVX512 below
loaded = False
@@ -119,10 +144,13 @@ if has_AVX512 and not loaded:
try:
logger.info("Loading faiss with AVX512 support.")
from .swigfaiss_avx512 import *
logger.info("Successfully loaded faiss with AVX512 support.")
loaded = True
except ImportError as e:
logger.info(f"Could not load library with AVX512 support due to:\n{e!r}")
logger.info(
f"Could not load library with AVX512 support due to:\n{e!r}"
)
# reset so that we load without AVX512 below
loaded = False
@@ -131,6 +159,7 @@ if has_AVX2 and not loaded:
try:
logger.info("Loading faiss with AVX2 support.")
from .swigfaiss_avx2 import *
logger.info("Successfully loaded faiss with AVX2 support.")
loaded = True
except ImportError as e:
@@ -143,6 +172,7 @@ if has_SVE and not loaded:
try:
logger.info("Loading faiss with SVE support.")
from .swigfaiss_sve import *
logger.info("Successfully loaded faiss with SVE support.")
loaded = True
except ImportError as e:
@@ -155,6 +185,7 @@ if not loaded:
# we import * so that the symbol X can be accessed as faiss.X
logger.info("Loading faiss.")
from .swigfaiss import *
logger.info("Successfully loaded faiss.")
except ModuleNotFoundError:
formatted_ins_sets = ", ".join(supported_instruction_sets())
@@ -165,7 +196,6 @@ if not loaded:
f"A) Set the correct FAISS_OPT_LEVEL value when executing "
f"'cmake'.\n"
f"B) Build the correct SWIG wrapper.\n\n"
f"These are the supported instruction sets on your system:\n"
f"{formatted_ins_sets}\n"
f"- If 'AVX512_SPR' (case insensitive) is supported on your "
faiss/python/setup.py
+5 -2
View File
@@ -83,7 +83,9 @@ if found_swigfaiss_avx512:
if found_swigfaiss_avx512_spr:
print(f"Copying {swigfaiss_avx512_spr_lib}")
shutil.copyfile("swigfaiss_avx512_spr.py", "faiss/swigfaiss_avx512_spr.py")
shutil.copyfile(swigfaiss_avx512_spr_lib, f"faiss/_swigfaiss_avx512_spr{ext}")
shutil.copyfile(
swigfaiss_avx512_spr_lib, f"faiss/_swigfaiss_avx512_spr{ext}"
)
if found_callbacks:
print(f"Copying {callbacks_lib}")
@@ -97,7 +99,8 @@ if found_swigfaiss_sve:
if found_faiss_example_external_module_lib:
print(f"Copying {faiss_example_external_module_lib}")
shutil.copyfile(
"faiss_example_external_module.py", "faiss/faiss_example_external_module.py"
"faiss_example_external_module.py",
"faiss/faiss_example_external_module.py",
)
shutil.copyfile(
faiss_example_external_module_lib,
perf_tests/bench_hnsw.py
+3 -5
View File
@@ -45,9 +45,7 @@ def is_perf_counter(key: str) -> bool:
def accumulate_perf_counter(
phase: str,
t: PerfCounters,
counters: Dict[str, int]
phase: str, t: PerfCounters, counters: Dict[str, int]
):
counters[f"{phase}_wall_time_us"] = int(t.wall_time_s * US_IN_S)
counters[f"{phase}_user_time_us"] = int(t.user_time_s * US_IN_S)
@@ -198,8 +196,8 @@ def main():
for counter, values in result.items():
if is_perf_counter(counter):
print(
"%s t=%.3f us (± %.4f)" %
(counter, np.mean(values), np.std(values))
"%s t=%.3f us (± %.4f)"
% (counter, np.mean(values), np.std(values))
)
tests/common_faiss_tests.py
+25 -15
View File
@@ -5,7 +5,12 @@
# a few common functions for the tests
from __future__ import absolute_import, division, print_function, unicode_literals
from __future__ import (
absolute_import,
division,
print_function,
unicode_literals,
)
import sys
import unittest
@@ -50,8 +55,9 @@ class Randu10k:
D, I = DI
e = {}
for rank in 1, 10, 100:
e[rank] = ((I[:, :rank] == self.gt.reshape(-1, 1)).sum() /
float(self.nq))
e[rank] = (I[:, :rank] == self.gt.reshape(-1, 1)).sum() / float(
self.nq
)
return e
@@ -77,9 +83,9 @@ class Randu10kUnbalanced(Randu10k):
def get_dataset(d, nb, nt, nq):
rs = np.random.RandomState(123)
xb = rs.rand(nb, d).astype('float32')
xt = rs.rand(nt, d).astype('float32')
xq = rs.rand(nq, d).astype('float32')
xb = rs.rand(nb, d).astype("float32")
xt = rs.rand(nt, d).astype("float32")
xq = rs.rand(nq, d).astype("float32")
return (xt, xb, xq)
@@ -88,7 +94,7 @@ def get_dataset_2(d, nt, nb, nq):
"""A dataset that is not completely random but still challenging to
index
"""
d1 = 10 # intrinsic dimension (more or less)
d1 = 10 # intrinsic dimension (more or less)
n = nb + nt + nq
rs = np.random.RandomState(1338)
x = rs.normal(size=(n, d1))
@@ -97,14 +103,14 @@ def get_dataset_2(d, nt, nb, nq):
# higher factor (>4) -> higher frequency -> less linear
x = x * (rs.rand(d) * 4 + 0.1)
x = np.sin(x)
x = x.astype('float32')
return x[:nt], x[nt:nt + nb], x[nt + nb:]
x = x.astype("float32")
return x[:nt], x[nt : nt + nb], x[nt + nb :]
def make_binary_dataset(d, nt, nb, nq):
assert d % 8 == 0
rs = np.random.RandomState(123)
x = rs.randint(256, size=(nb + nq + nt, int(d / 8))).astype('uint8')
x = rs.randint(256, size=(nb + nq + nt, int(d / 8))).astype("uint8")
return x[:nt], x[nt:-nq], x[-nq:]
@@ -115,19 +121,22 @@ def compare_binary_result_lists(D1, I1, D2, I2):
assert D1.shape == I1.shape == D2.shape == I2.shape
n, k = D1.shape
ndiff = (D1 != D2).sum()
assert ndiff == 0, '%d differences in distance matrix %s' % (
ndiff, D1.shape)
assert ndiff == 0, "%d differences in distance matrix %s" % (
ndiff,
D1.shape,
)
def normalize_DI(D, I):
norm = I.max() + 1.0
Dr = D.astype('float64') + I / norm
Dr = D.astype("float64") + I / norm
# ignore -1s and elements on last column
Dr[I1 == -1] = 1e20
Dr[D == D[:, -1:]] = 1e20
Dr.sort(axis=1)
return Dr
ndiff = (normalize_DI(D1, I1) != normalize_DI(D2, I2)).sum()
assert ndiff == 0, '%d differences in normalized D matrix' % ndiff
assert ndiff == 0, "%d differences in normalized D matrix" % ndiff
# Map SIMDLevel enum values to their string names.
@@ -218,7 +227,8 @@ def for_all_simd_levels(cls):
def tearDown(self):
super().tearDown()
faiss.SIMDConfig.set_level(
faiss.SIMDConfig.get_dispatched_level())
faiss.SIMDConfig.get_dispatched_level()
)
Parameterized.__name__ = f"{cls.__name__}_{ln}"
Parameterized.__qualname__ = f"{cls.__qualname__}_{ln}"
tests/index_io_backward_compatibility/cmake_conda_io_utils.py
+10 -27
View File
@@ -121,9 +121,7 @@ INDEX_BINARY_TYPES = [
def sanitize_index_key(index_key: str) -> str:
return (
index_key.replace(",", "__").replace("(", "--").replace(")", "--")
)
return index_key.replace(",", "__").replace("(", "--").replace(")", "--")
def generate_vectors(d: int, nb: int, seed: int) -> np.ndarray:
@@ -133,17 +131,13 @@ def generate_vectors(d: int, nb: int, seed: int) -> np.ndarray:
return xb
def write_vectors(
output_dir: str, vectors: np.ndarray, filename: str
) -> None:
def write_vectors(output_dir: str, vectors: np.ndarray, filename: str) -> None:
path = os.path.join(output_dir, filename)
print(f"Writing vectors to: {path}")
np.save(path, vectors)
def create_and_populate_index(
index_key: str, xb: np.ndarray
) -> faiss.Index:
def create_and_populate_index(index_key: str, xb: np.ndarray) -> faiss.Index:
d = xb.shape[1]
index = faiss.index_factory(d, index_key)
print(f"Created index type: {type(index).__name__}")
@@ -156,9 +150,7 @@ def create_and_populate_index(
# IDMap indexes require add_with_ids() instead of add()
if "IDMap" in index_key:
ids = np.arange(xb.shape[0], dtype=np.int64)
print(
f"Using add_with_ids() with IDs from 0 to {xb.shape[0] - 1}"
)
print(f"Using add_with_ids() with IDs from 0 to {xb.shape[0] - 1}")
index.add_with_ids(xb, ids)
else:
index.add(xb)
@@ -182,9 +174,7 @@ def create_and_populate_binary_index(
# IDMap indexes require add_with_ids() instead of add()
if "IDMap" in index_key:
ids = np.arange(xb.shape[0], dtype=np.int64)
print(
f"Using add_with_ids() with IDs from 0 to {xb.shape[0] - 1}"
)
print(f"Using add_with_ids() with IDs from 0 to {xb.shape[0] - 1}")
index.add_with_ids(xb, ids)
else:
index.add(xb)
@@ -193,9 +183,7 @@ def create_and_populate_binary_index(
return index
def write_index(
index: faiss.Index, output_dir: str, filename: str
) -> None:
def write_index(index: faiss.Index, output_dir: str, filename: str) -> None:
index_path = os.path.join(output_dir, filename)
# Delete existing file if present
@@ -359,9 +347,7 @@ def test_write_binary_index_type(
}
def write_test_all_files(
writer: str, output_dir: str, seed: int
) -> int:
def write_test_all_files(writer: str, output_dir: str, seed: int) -> int:
print(f"{writer.capitalize()} Writer: Testing index serialization")
os.makedirs(output_dir, exist_ok=True)
@@ -401,13 +387,12 @@ def write_test_all_files(
# Now write binary indexes
print(
f"\n{writer.capitalize()} Writer: "
"Testing binary index serialization"
f"\n{writer.capitalize()} Writer: " "Testing binary index serialization"
)
print(f"\nGenerating binary vectors: {nb} database vectors, dimension {d}")
assert d % 8 == 0
np.random.seed(seed)
xb_binary = np.random.randint(256, size=(nb, int(d / 8))).astype('uint8')
xb_binary = np.random.randint(256, size=(nb, int(d / 8))).astype("uint8")
write_vectors(output_dir, xb_binary, f"vectors_db_binary_{writer}.npy")
binary_results = []
@@ -508,9 +493,7 @@ def read_test_all_files(reader: str, writer: str, input_dir: str) -> int:
# Now read binary indexes
print(f"\n{reader} Reader: Testing binary index deserialization")
binary_metadata = read_metadata(
input_dir, f"{writer}_binary_metadata.json"
)
binary_metadata = read_metadata(input_dir, f"{writer}_binary_metadata.json")
print(f"Binary metadata from {writer} build:")
print_metadata(binary_metadata)

Some files were not shown because too many files have changed in this diff Show More