mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Tia Lim
b37e929947
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5363 `ResultHeap` is now part of the faiss Python package: it is defined in `faiss/python/extra_wrappers.py` and re-exported into the top-level namespace from `faiss/python/__init__.py`, so it is importable as `faiss.ResultHeap`. The `faiss.contrib` modules (`big_batch_search.py`, `exhaustive_search.py`, `client_server.py`) already consume the packaged class. Two benchmark scripts still carried their own older, byte-for-byte-duplicated copy of the class: - `benchs/bench_for_interrupt.py` - `benchs/distributed_ondisk/search_server.py` This removes both local copies and uses `faiss.ResultHeap` instead. The packaged class exposes a slightly different (and richer) API than the old copies: it has `add_result(D, I)` rather than `add_batch_result(D, I, i0)`. The old `add_batch_result` simply applied an id offset (`I += i0`) before pushing the batch onto the heaps, so the call sites are translated as: - `rh.add_batch_result(Di, Ii, i)` -> `Ii += i; rh.add_result(Di, Ii)` - `rh.add_batch_result(Di, Ii, 0)` -> `rh.add_result(Di, Ii)` (offset was 0) This is the code half of task T276876996. The other half is the broken wiki link on the "Brute force search without an index" page, which points at `benchs/bench_all_ivf/datasets.py#L75` on `master`. Now that `ResultHeap` lives in the package, the link is best repointed at the canonical implementation: `faiss/python/extra_wrappers.py`. That edit is on the GitHub wiki (a separate repo) and is not part of this diff. Reviewed By: mnorris11 Differential Revision: D110118635 fbshipit-source-id: b9e0ba5ecbec3e2c22bdf6d26da931be976ed944
139 lines
3.2 KiB
Python
139 lines
3.2 KiB
Python
#! /usr/bin/env python3
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from __future__ import print_function
|
|
import numpy as np
|
|
import faiss
|
|
import time
|
|
import os
|
|
import argparse
|
|
|
|
|
|
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)
|
|
|
|
|
|
args = parser.parse_args()
|
|
print("args:", args)
|
|
|
|
|
|
d = args.dim # dimension
|
|
nb = args.nb # database size
|
|
nq = 1000 # nb of queries
|
|
nt = 100000
|
|
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")
|
|
k = 100
|
|
|
|
if args.no_intcallback:
|
|
faiss.InterruptCallback.clear_instance()
|
|
|
|
if args.nt != -1:
|
|
faiss.omp_set_num_threads(args.nt)
|
|
|
|
nprobe = args.nprobe
|
|
key = args.key
|
|
# key = 'IVF1000,Flat'
|
|
# key = 'IVF1000,PQ64'
|
|
# key = 'IVF100_HNSW32,PQ64'
|
|
|
|
# faiss.omp_set_num_threads(1)
|
|
|
|
pf = "dim%d_" % d
|
|
if d == 64:
|
|
pf = ""
|
|
|
|
basename = "/tmp/base%s%s.index" % (pf, key)
|
|
|
|
if os.path.exists(basename):
|
|
print("load", basename)
|
|
index_1 = faiss.read_index(basename)
|
|
else:
|
|
print("train + write", basename)
|
|
index_1 = faiss.index_factory(d, key)
|
|
index_1.train(xt)
|
|
faiss.write_index(index_1, basename)
|
|
|
|
print("add")
|
|
index_1.add(xb)
|
|
|
|
print("set nprobe=", nprobe)
|
|
faiss.ParameterSpace().set_index_parameter(index_1, "nprobe", nprobe)
|
|
|
|
|
|
stats = faiss.cvar.indexIVF_stats
|
|
stats.reset()
|
|
|
|
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)
|
|
)
|
|
|
|
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])
|
|
indexes[i] = index
|
|
|
|
rh = faiss.ResultHeap(nq, k)
|
|
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])
|
|
else:
|
|
index = indexes[i]
|
|
|
|
stats.reset()
|
|
start = time.time()
|
|
Di, Ii = index.search(xq, k)
|
|
sum_time = sum_time + time.time() - start
|
|
tq += stats.quantization_time
|
|
ts += stats.search_time
|
|
Ii += i
|
|
rh.add_result(Di, Ii)
|
|
|
|
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))
|