def test_equiv_rcq_rq(self):
""" make sure that the codes generated by the standalone codec are the same
between an
IndexRefine with ResidualQuantizer
and
IVF with ResidualCoarseQuantizer
both are the centroid id concatenated with the code.
"""
ds = SyntheticDataset(16, 400, 100, 0)
index1 = faiss.index_factory(ds.d, "RQ2x3,Refine(Flat)")
index1.train(ds.get_train())
irq = faiss.downcast_index(index1.base_index)
# because the default beam factor for RCQ is 4
irq.rq.max_beam_size = 4
index2 = faiss.index_factory(ds.d, "IVF64(RCQ2x3),Flat")
index2.train(ds.get_train())
quantizer = faiss.downcast_index(index2.quantizer)
quantizer.rq = irq.rq
index2.is_trained = True
codes1 = index1.sa_encode(ds.get_database())
codes2 = index2.sa_encode(ds.get_database())
np.testing.assert_array_equal(codes1, codes2)
def test_factory_NSG(self):
index = faiss.index_factory(12, "NSG64")
assert isinstance(index, faiss.IndexNSGFlat)
assert index.nsg.R == 64
index = faiss.index_factory(12, "NSG64", faiss.METRIC_INNER_PRODUCT)
assert isinstance(index, faiss.IndexNSGFlat)
assert index.nsg.R == 64
assert index.metric_type == faiss.METRIC_INNER_PRODUCT
index = faiss.index_factory(12, "NSG64,Flat")
assert isinstance(index, faiss.IndexNSGFlat)
assert index.nsg.R == 64
index = faiss.index_factory(12, "IVF65536_NSG64,Flat")
index_nsg = faiss.downcast_index(index.quantizer)
assert isinstance(index, faiss.IndexIVFFlat)
assert isinstance(index_nsg, faiss.IndexNSGFlat)
assert index.nlist == 65536 and index_nsg.nsg.R == 64
index = faiss.index_factory(12, "IVF65536_NSG64,PQ2x8")
index_nsg = faiss.downcast_index(index.quantizer)
assert isinstance(index, faiss.IndexIVFPQ)
assert isinstance(index_nsg, faiss.IndexNSGFlat)
assert index.nlist == 65536 and index_nsg.nsg.R == 64
assert index.pq.M == 2 and index.pq.nbits == 8
def check_index_recon(embeds_path,
index_or_index_path,
embeds_format='labeled_numpy',
sort=True,
**kwargs):
index = faiss.read_index(index_or_index_path) if isinstance(
index_or_index_path, str) else index_or_index_path
faiss.downcast_index(index).make_direct_map()
embeds_list, _ = load_embeds(embeds_path=embeds_path,
format=embeds_format,
sort=sort,
**kwargs)
# tic("Gathering targets ...")
# all_tgt_embeds = []
# for file_path in Tqdm.tqdm(embeds_paths):
# embeds_group, batch_group = pickle_load(file_path)
# for embeds, batch in zip(embeds_group[embeds_key], batch_group):
# all_tgt_embeds.append(embeds[embeds_idx])
#
# toc("Done!")
tic("Checking embedding reconstruction difference ...")
all_embeds = np.concatenate(embeds_list)
all_embeds_recon = index.reconstruct_n(0, len(all_embeds))
embeds_diff = np.linalg.norm(all_embeds - all_embeds_recon)
toc("Passed embedding reconstruction difference check.") \
if embeds_diff == 0 else toc(f"Embedding reconstruction difference: {embeds_diff}.")
def reverse_index_factory(index):
"""
attempts to get the factory string the index was built with
"""
index = faiss.downcast_index(index)
if isinstance(index, faiss.IndexFlat):
return "Flat"
if isinstance(index, faiss.IndexIVF):
quantizer = faiss.downcast_index(index.quantizer)
if isinstance(quantizer, faiss.IndexFlat):
prefix = "IVF%d" % index.nlist
elif isinstance(quantizer, faiss.MultiIndexQuantizer):
prefix = "IMI%dx%d" % (quantizer.pq.M, quantizer.pq.nbit)
elif isinstance(quantizer, faiss.IndexHNSW):
prefix = "IVF%d_HNSW%d" % (index.nlist, quantizer.hnsw.M)
else:
prefix = "IVF%d(%s)" % (index.nlist, reverse_index_factory(quantizer))
if isinstance(index, faiss.IndexIVFFlat):
return prefix + ",Flat"
if isinstance(index, faiss.IndexIVFScalarQuantizer):
return prefix + ",SQ8"
raise NotImplementedError()
def test_equiv_sh(self):
""" make sure that the IVFSpectralHash sa_encode function gives the same
result as the concatenated RQ + LSH index sa_encode """
ds = SyntheticDataset(32, 500, 100, 0)
index1 = faiss.index_factory(ds.d, "RQ1x4,Refine(ITQ16,LSH)")
index1.train(ds.get_train())
# reproduce this in an IndexIVFSpectralHash
coarse_quantizer = faiss.IndexFlat(ds.d)
rq = faiss.downcast_index(index1.base_index).rq
centroids = get_additive_quantizer_codebooks(rq)[0]
coarse_quantizer.add(centroids)
encoder = faiss.downcast_index(index1.refine_index)
# larger than the magnitude of the vectors
# negative because otherwise the bits are flipped
period = -100000.0
index2 = faiss.IndexIVFSpectralHash(coarse_quantizer, ds.d,
coarse_quantizer.ntotal,
encoder.sa_code_size() * 8, period)
# replace with the vt of the encoder. Binarization is performed by
# the IndexIVFSpectralHash itself
index2.replace_vt(encoder)
codes1 = index1.sa_encode(ds.get_database())
codes2 = index2.sa_encode(ds.get_database())
np.testing.assert_array_equal(codes1, codes2)
def test_nprobe_4(self):
index = faiss.index_factory(32, "PCAR32,IVF32,SQ8,RFlat")
ps = faiss.ParameterSpace()
ps.set_index_parameter(index, "nprobe", 5)
index2 = faiss.downcast_index(index.base_index)
index2 = faiss.downcast_index(index2.index)
self.assertEqual(index2.nprobe, 5)
def unwind_index_ivf(index):
if isinstance(index, faiss.IndexPreTransform):
assert index.chain.size() == 1
vt = index.chain.at(0)
index_ivf, vt2 = unwind_index_ivf(faiss.downcast_index(index.index))
assert vt2 is None
return index_ivf, vt
if hasattr(faiss, "IndexRefine") and isinstance(index, faiss.IndexRefine):
return unwind_index_ivf(faiss.downcast_index(index.base_index))
if isinstance(index, faiss.IndexIVF):
return index, None
else:
return None, None
def search_single_scan(index, xq, k, bs=128):
"""performs a search so that the inverted lists are accessed
sequentially by blocks of size bs"""
# handle pretransform
if isinstance(index, faiss.IndexPreTransform):
xq = index.apply_py(xq)
index = faiss.downcast_index(index.index)
# coarse assignment
nprobe = min(index.nprobe, index.nlist)
coarse_dis, assign = index.quantizer.search(xq, nprobe)
nlist = index.nlist
assign_buckets = assign // bs
nq = len(xq)
rh = faiss.ResultHeap(nq, k)
index.parallel_mode |= index.PARALLEL_MODE_NO_HEAP_INIT
for l0 in range(0, nlist, bs):
bucket_no = l0 // bs
skip_rows, skip_cols = np.where(assign_buckets != bucket_no)
sub_assign = assign.copy()
sub_assign[skip_rows, skip_cols] = -1
index.search_preassigned(nq, faiss.swig_ptr(xq), k,
faiss.swig_ptr(sub_assign),
faiss.swig_ptr(coarse_dis),
faiss.swig_ptr(rh.D), faiss.swig_ptr(rh.I),
False, None)
rh.finalize()
return rh.D, rh.I
def test_nsg_sq(self):
"""Test IndexNSGSQ"""
d = self.xq.shape[1]
R = 32
index = faiss.index_factory(d, f"NSG{R}_SQ8")
assert isinstance(index, faiss.IndexNSGSQ)
idxsq = faiss.downcast_index(index.storage)
assert index.nsg.R == R
assert idxsq.sq.qtype == faiss.ScalarQuantizer.QT_8bit
flat_index = faiss.IndexFlat(d)
flat_index.add(self.xb)
Dref, Iref = flat_index.search(self.xq, k=1)
index.train(self.xb)
index.add(self.xb)
D, I = index.search(self.xq, k=1)
# test accuracy
recalls = (Iref == I).sum()
print("IndexNSGSQ", recalls)
self.assertGreaterEqual(recalls, 405) # 411
# test I/O
self.subtest_io_and_clone(index, D, I)
def test_nsg_pq(self):
"""Test IndexNSGPQ"""
d = self.xq.shape[1]
R, pq_M = 32, 4
index = faiss.index_factory(d, f"NSG{R}_PQ{pq_M}")
assert isinstance(index, faiss.IndexNSGPQ)
idxpq = faiss.downcast_index(index.storage)
assert index.nsg.R == R and idxpq.pq.M == pq_M
flat_index = faiss.IndexFlat(d)
flat_index.add(self.xb)
Dref, Iref = flat_index.search(self.xq, k=1)
index.GK = 32
index.train(self.xb)
index.add(self.xb)
D, I = index.search(self.xq, k=1)
# test accuracy
recalls = (Iref == I).sum()
print("IndexNSGPQ", recalls)
self.assertGreaterEqual(recalls, 190) # 193
# test I/O
self.subtest_io_and_clone(index, D, I)
def query(self, vecs, topk, param):
#self.index.nprobe = param["nprobe"]
#faiss.omp_set_num_threads(1) # Make sure this is on a single thread mode
for n_gpu in range(self.index.count()):
faiss.downcast_index(self.index.at(n_gpu)).nprobe = param["nprobe"]
_, ids = self.index.search(x=vecs, k=topk)
return ids
def train_ivf_index_with_2level(index, xt, **args):
"""
Applies 2-level clustering to an index_ivf embedded in an index.
"""
# handle PreTransforms
index = faiss.downcast_index(index)
if isinstance(index, faiss.IndexPreTransform):
for i in range(index.chain.size()):
vt = index.chain.at(i)
vt.train(xt)
xt = vt.apply(xt)
train_ivf_index_with_2level(index.index, xt)
index.is_trained = True
return
assert isinstance(index, faiss.IndexIVF)
assert index.metric_type == faiss.METRIC_L2
# now do 2-level clustering
nc1 = int(np.sqrt(index.nlist))
cc = np.arange(nc1 + 1) * index.nlist // nc1
all_nc2 = cc[1:] - cc[:-1]
centroids, _ = two_level_clustering(xt, nc1, all_nc2, **args)
index.quantizer.train(centroids)
index.quantizer.add(centroids)
# finish training
index.train(xt)
def search(img, k=10, nprobe=10):
"""Find k near neighbours in index
Arguments:
img {PIL.Image} -- Face
Keyword Arguments:
k {int} -- Count of nearest neighbours (default: {10})
nprobe {int} -- How many nearest clusters to scan (read in Faiss docs) (default: {10})
Returns:
np.array -- k distances
np.array -- k indices
"""
backbone = ResNet_50([112, 112])
pth = os.path.join(settings.BACKBONE_DIR,
BACKBONE_FILE) # Pretrained backbone for ResNet50
index = faiss.read_index(os.path.join(DATASET_PATH,
DATASET_INDEX)) # load index
index_ivf = faiss.downcast_index(index.index)
index_ivf.nprobe = nprobe # change nprobe
query = np.array(extract_one_embedding(img, backbone,
pth)).astype('float32').reshape(
1, -1)
D, I = index.search(query, k)
return D[0], I[0]
def nearest(self, vector=None, n=12, nprobe=16):
logging.info("Index size {} with {} loaded entries in {}".format(
self.faiss_index.ntotal, len(self.loaded_entries), self.name))
if type(self.faiss_index) == faiss.swigfaiss.IndexPreTransform:
index_ivf = faiss.downcast_index(self.faiss_index.index)
index_ivf.nprobe = nprobe
else:
self.faiss_index.nprobe = nprobe
vector = np.atleast_2d(vector)
if vector.shape[-1] != self.faiss_index.d:
vector = vector.T
results = []
dist, ids = self.faiss_index.search(vector, n)
for i, k in enumerate(ids[0]):
if k >= 0:
index_entry = sorted(self.tree[k])[0]
temp = {
'rank': i + 1,
'algo': self.name,
'dist': float(dist[0, i]),
'indexentries_pk': index_entry.data,
'offset': k - index_entry.begin
}
results.append(temp)
return results
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. """
indexall = prepare_trained_index(preproc)
co = faiss.GpuMultipleClonerOptions()
co.useFloat16 = use_float16
co.useFloat16CoarseQuantizer = False
co.usePrecomputed = use_precomputed_tables
co.indicesOptions = faiss.INDICES_CPU
co.verbose = True
co.reserveVecs = max_add if max_add > 0 else xb.shape[0]
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)
print("add...")
t0 = time.time()
nb = xb.shape[0]
for i0, xs in dataset_iterator(xb, preproc, add_batch_size):
i1 = i0 + xs.shape[0]
gpu_index.add_with_ids(xs, np.arange(i0, i1))
if max_add > 0 and gpu_index.ntotal > max_add:
print("Flush indexes to CPU")
for i in range(ngpu):
index_src_gpu = faiss.downcast_index(gpu_index.at(i))
index_src = faiss.index_gpu_to_cpu(index_src_gpu)
print(" index %d size %d" % (i, index_src.ntotal))
index_src.copy_subset_to(indexall, 0, 0, nb)
index_src_gpu.reset()
index_src_gpu.reserveMemory(max_add)
gpu_index.sync_with_shard_indexes()
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'):
# it is a sharded index
for i in range(ngpu):
index_src = faiss.index_gpu_to_cpu(gpu_index.at(i))
print(" index %d size %d" % (i, index_src.ntotal))
index_src.copy_subset_to(indexall, 0, 0, nb)
else:
# simple index
index_src = faiss.index_gpu_to_cpu(gpu_index)
index_src.copy_subset_to(indexall, 0, 0, nb)
print(" done in %.3f s" % (time.time() - t0))
if max_add > 0:
# it does not contain all the vectors
gpu_index = None
return gpu_index, indexall
def _load_ann_index(index_filename: str, device: int) -> faiss.Index:
"""
Load the ANN index from the given file and move it to the GPU(s).
Parameters
----------
index_filename : str
The ANN index filename.
Returns
-------
faiss.Index
The Faiss `Index`.
"""
# https://github.com/facebookresearch/faiss/blob/2cce2e5f59a5047aa9a1729141e773da9bec6b78/benchs/bench_gpu_1bn.py#L608
# logger.debug('Load the ANN index from file %s', index_filename)
index_cpu = faiss.read_index(index_filename)
res = faiss.StandardGpuResources()
co = faiss.GpuClonerOptions()
co.useFloat16 = True
co.useFloat16CoarseQuantizer = False
co.indicesOptions = faiss.INDICES_CPU
co.reserveVecs = index_cpu.ntotal
index = faiss.index_cpu_to_gpu(res, device, index_cpu, co)
if hasattr(index, 'at'):
for i in range(index.count()):
simple_index = faiss.downcast_index(index.at(i))
simple_index.nprobe = min(math.ceil(simple_index.nlist / 2),
config.num_probe)
else:
index.nprobe = min(math.ceil(index.nlist / 2), config.num_probe)
return index
def test_ivfsq(self):
ds = datasets.SyntheticDataset(32, 3000, 1000, 100)
xt = ds.get_train()
xb = ds.get_database()
gt = ds.get_groundtruth(1)
# RQ 2x5 = 10 bits = 1024 centroids
index = faiss.index_factory(ds.d, "IVF1024(RCQ2x5),SQ8")
quantizer = faiss.downcast_index(index.quantizer)
rq = quantizer.rq
rq.train_type = faiss.ResidualQuantizer.Train_default
index.train(xt)
index.add(xb)
# make sure that increasing the nprobe increases accuracy
index.nprobe = 10
D, I = index.search(ds.get_queries(), 10)
r10 = (I == gt[None, :]).sum() / ds.nq
index.nprobe = 40
D, I = index.search(ds.get_queries(), 10)
r40 = (I == gt[None, :]).sum() / ds.nq
self.assertGreater(r40, r10)
def set_nprobe(self, nprobe) -> int:
"""Set the value of nprobe.
Args:
nprobe: The new value for nprobe
"""
faiss.ParameterSpace().set_index_parameter(self.index, "nprobe",
nprobe)
return faiss.downcast_index(self.index).nprobe
def flushGPUIndex(self):
for i in range(self.ngpu):
if self.ngpu > 1:
index_src_gpu = faiss.downcast_index(self.gpu_index.at(i))
else:
index_src_gpu = faiss.downcast_index(self.gpu_index)
index_src = faiss.index_gpu_to_cpu(index_src_gpu)
# print(" index %d size %d" % (i, index_src.ntotal))
if self.index is None:
self.index = faiss.read_index(self.emptyIndexPath)
index_src.copy_subset_to(
self.index, 0, 0,
self.totalImagesToIndex * self.feats_per_file)
index_src_gpu.reset()
index_src.reset()
index_src_gpu.reserveMemory(self.max_add)
if self.ngpu > 1:
self.gpu_index.sync_with_shard_indexes()
def test_factory_3(self):
index = faiss.index_factory(12, "IVF10,PQ4")
faiss.ParameterSpace().set_index_parameter(index, "nprobe", 3)
assert index.nprobe == 3
index = faiss.index_factory(12, "PCAR8,IVF10,PQ4")
faiss.ParameterSpace().set_index_parameter(index, "nprobe", 3)
assert faiss.downcast_index(index.index).nprobe == 3
def test_factory_3(self):
index = faiss.index_factory(12, "IVF10,PQ4")
faiss.ParameterSpace().set_index_parameter(index, "nprobe", 3)
assert index.nprobe == 3
index = faiss.index_factory(12, "PCAR8,IVF10,PQ4")
faiss.ParameterSpace().set_index_parameter(index, "nprobe", 3)
assert faiss.downcast_index(index.index).nprobe == 3
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. """
indexall = prepare_trained_index(preproc)
co = faiss.GpuMultipleClonerOptions()
co.useFloat16 = use_float16
co.useFloat16CoarseQuantizer = False
co.usePrecomputed = use_precomputed_tables
co.indicesOptions = faiss.INDICES_CPU
co.verbose = 10
co.reserveVecs = max_add if max_add > 0 else xb.shape[0]
co.shard = True
vres, vdev = make_vres_vdev()
gpu_index = faiss.index_cpu_to_gpu_multiple(
vres, vdev, indexall, co)
print "add..."
t0 = time.time()
nb = xb.shape[0]
for i0, xs in dataset_iterator(xb, preproc, add_batch_size):
i1 = i0 + xs.shape[0]
gpu_index.add_with_ids(xs, np.arange(i0, i1))
if max_add > 0 and gpu_index.ntotal > max_add:
print "Flush indexes to CPU"
for i in range(ngpu):
index_src_gpu = faiss.downcast_index(gpu_index.at(i))
index_src = faiss.index_gpu_to_cpu(index_src_gpu)
print " index %d size %d" % (i, index_src.ntotal)
index_src.copy_subset_to(indexall, 0, 0, nb)
index_src_gpu.reset()
index_src_gpu.reserveMemory(max_add)
gpu_index.sync_with_shard_indexes()
print '\r%d/%d (%.3f s) ' % (
i0, nb, time.time() - t0),
sys.stdout.flush()
print "Add time: %.3f s" % (time.time() - t0)
print "Aggregate indexes to CPU"
t0 = time.time()
for i in range(ngpu):
index_src = faiss.index_gpu_to_cpu(gpu_index.at(i))
print " index %d size %d" % (i, index_src.ntotal)
index_src.copy_subset_to(indexall, 0, 0, nb)
print " done in %.3f s" % (time.time() - t0)
if max_add > 0:
# it does not contain all the vectors
gpu_index = None
return gpu_index, indexall
def test_factory_HNSW_newstyle(self):
index = faiss.index_factory(12, "HNSW32,Flat")
assert index.storage.sa_code_size() == 12 * 4
index = faiss.index_factory(12, "HNSW32,SQ8", faiss.METRIC_INNER_PRODUCT)
assert index.storage.sa_code_size() == 12
assert index.metric_type == faiss.METRIC_INNER_PRODUCT
index = faiss.index_factory(12, "HNSW32,PQ4")
assert index.storage.sa_code_size() == 4
index = faiss.index_factory(12, "HNSW32,PQ4np")
indexpq = faiss.downcast_index(index.storage)
assert not indexpq.do_polysemous_training
def test_equiv_rq(self):
"""
make sure it is equivalent to search a RQ and to search an IVF
with RCQ + RQ with the same codebooks.
"""
ds = datasets.SyntheticDataset(32, 3000, 1000, 50)
# make a flat RQ
iflat = faiss.IndexResidualQuantizer(ds.d, 5, 4)
iflat.rq.train_type = faiss.ResidualQuantizer.Train_default
iflat.train(ds.get_train())
iflat.add(ds.get_database())
# ref search result
Dref, Iref = iflat.search(ds.get_queries(), 10)
# get its codebooks + encoded version of the dataset
codebooks = get_additive_quantizer_codebooks(iflat.rq)
codes = faiss.vector_to_array(iflat.codes).reshape(-1, iflat.code_size)
# make an IVF with 2x4 + 3x4 = 5x4 bits
ivf = faiss.index_factory(ds.d, "IVF256(RCQ2x4),RQ3x4")
# initialize the codebooks
rcq = faiss.downcast_index(ivf.quantizer)
faiss.copy_array_to_vector(
np.vstack(codebooks[:rcq.rq.M]).ravel(),
rcq.rq.codebooks
)
rcq.rq.is_trained = True
# translation of AdditiveCoarseQuantizer::train
rcq.ntotal = 1 << rcq.rq.tot_bits
rcq.centroid_norms.resize(rcq.ntotal)
rcq.rq.compute_centroid_norms(rcq.centroid_norms.data())
rcq.is_trained = True
faiss.copy_array_to_vector(
np.vstack(codebooks[rcq.rq.M:]).ravel(),
ivf.rq.codebooks
)
ivf.rq.is_trained = True
ivf.is_trained = True
# add the codes (this works because 2x4 is a multiple of 8 bits)
ivf.add_sa_codes(codes)
# perform exhaustive search
ivf.nprobe = ivf.nlist
Dnew, Inew = ivf.search(ds.get_queries(), 10)
np.testing.assert_array_equal(Iref, Inew)
np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5)
def __init__(self, ds, indexfile):
self.d = ds.d
self.metric = ds.metric
self.nq = ds.nq
self.xq = ds.get_queries()
# get the xb set
src_index = faiss.read_index(indexfile)
src_quant = faiss.downcast_index(src_index.quantizer)
centroids = faiss.vector_to_array(src_quant.xb)
self.xb = centroids.reshape(-1, self.d)
self.nb = self.nt = len(self.xb)
def _flush_to_cpu(self, index, nb, offset):
print("Flush indexes to CPU")
for i in range(self.ngpu):
index_src_gpu = faiss.downcast_index(self.gpu_index if self.ngpu ==
1 else self.gpu_index.at(i))
index_src = faiss.index_gpu_to_cpu(index_src_gpu)
# index_src.copy_subset_to(index, 0, 0, nb) # original
index_src.copy_subset_to(index, 0, offset, offset + nb)
index_src_gpu.reset()
index_src_gpu.reserveMemory(self.max_add)
if self.ngpu > 1:
self.gpu_index.sync_with_shard_indexes()
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)
index = faiss.read_index(fname)
indexes.append(index)
il = faiss.extract_index_ivf(index).invlists
else:
assert False
ilv.push_back(il)
print()
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)
self.masked_index = faiss.read_index(
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)
print('loading empty index', empty_index_fname)
self.index = faiss.read_index(empty_index_fname)
ntotal = self.big_il.compute_ntotal()
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
quantizer = faiss.downcast_index(index_ivf.quantizer)
quantizer.hnsw.efSearch = 1024
def test_rcq_LUT(self):
ds = datasets.SyntheticDataset(32, 3000, 1000, 100)
xt = ds.get_train()
xb = ds.get_database()
# RQ 2x5 = 10 bits = 1024 centroids
index = faiss.index_factory(ds.d, "IVF1024(RCQ2x5),SQ8")
quantizer = faiss.downcast_index(index.quantizer)
rq = quantizer.rq
rq.train_type = faiss.ResidualQuantizer.Train_default
index.train(xt)
index.add(xb)
index.nprobe = 10
# set exact centroids as coarse quantizer
all_centroids = quantizer.reconstruct_n(0, quantizer.ntotal)
q2 = faiss.IndexFlatL2(32)
q2.add(all_centroids)
index.quantizer = q2
Dref, Iref = index.search(ds.get_queries(), 10)
index.quantizer = quantizer
# search with LUT
quantizer.set_beam_factor(-1)
Dnew, Inew = index.search(ds.get_queries(), 10)
np.testing.assert_array_almost_equal(Dref, Dnew, decimal=5)
np.testing.assert_array_equal(Iref, Inew)
# check i/o
CDref, CIref = quantizer.search(ds.get_queries(), 10)
quantizer2 = faiss.deserialize_index(faiss.serialize_index(quantizer))
quantizer2.search(ds.get_queries(), 10)
CDnew, CInew = quantizer2.search(ds.get_queries(), 10)
np.testing.assert_array_almost_equal(CDref, CDnew, decimal=5)
np.testing.assert_array_equal(CIref, CInew)
def subtest_add2col(self, xb, xq, index, qname):
"""Test with 2 additional dimensions to take also the non-SIMD
codepath. We don't retrain anything but add 2 dims to the
queries, the centroids and the trained ScalarQuantizer.
"""
nb, d = xb.shape
d2 = d + 2
xb2 = self.add2columns(xb)
xq2 = self.add2columns(xq)
nlist = index.nlist
quantizer = faiss.downcast_index(index.quantizer)
quantizer2 = faiss.IndexFlat(d2, index.metric_type)
centroids = faiss.vector_to_array(quantizer.xb).reshape(nlist, d)
centroids2 = self.add2columns(centroids)
quantizer2.add(centroids2)
index2 = faiss.IndexIVFScalarQuantizer(
quantizer2, d2, index.nlist, index.sq.qtype,
index.metric_type)
index2.nprobe = 4
if qname in ('8bit', '4bit'):
trained = faiss.vector_to_array(index.sq.trained).reshape(2, -1)
nt = trained.shape[1]
# 2 lines: vmins and vdiffs
new_nt = int(nt * d2 / d)
trained2 = np.hstack((
trained,
np.zeros((2, new_nt - nt), dtype='float32')
))
trained2[1, nt:] = 1.0 # set vdiff to 1 to avoid div by 0
faiss.copy_array_to_vector(trained2.ravel(), index2.sq.trained)
else:
index2.sq.trained = index.sq.trained
index2.is_trained = True
index2.add(xb2)
return index2.search(xq2, 10)
xq.shape, gt.shape)
nq, d = xq.shape
nb, d = xb.shape
######################################################
# Make index
######################################################
if args.indexfile and os.path.exists(args.indexfile):
print "reading", args.indexfile
index = faiss.read_index(args.indexfile)
if isinstance(index, faiss.IndexPreTransform):
index_ivf = faiss.downcast_index(index.index)
else:
index_ivf = index
assert isinstance(index_ivf, faiss.IndexIVF)
vec_transform = lambda x: x
assert isinstance(index_ivf, faiss.IndexIVF)
else:
print "build index, key=", args.indexkey
index = faiss.index_factory(d, args.indexkey)
if isinstance(index, faiss.IndexPreTransform):
index_ivf = faiss.downcast_index(index.index)
vec_transform = index.chain.at(0).apply_py
def test_rcq(self):
index = faiss.index_factory(12, "IVF256(RCQ2x4),RQ3x4")
self.assertEqual(
faiss.downcast_index(index.quantizer).__class__,
faiss.ResidualCoarseQuantizer
)
def test_ivf_parent(self):
index = faiss.index_factory(123, "IVF100(LSHr),Flat")
quantizer = faiss.downcast_index(index.quantizer)
self.assertEqual(quantizer.__class__, faiss.IndexLSH)
