mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Evaluation script for hybrid CPU / GPU search
Summary: Implementation of various combinations of coarse quantization / scaning code on CPU and GPU. Used to generate the results of https://github.com/facebookresearch/faiss/wiki/Hybrid-CPU-GPU-search-and-multiple-GPUs Reviewed By: alexanderguzhva Differential Revision: D43041802 fbshipit-source-id: 12608812ab351d60d4a6dc45be1ca493f76d4375
This commit is contained in:
Matthijs Douze
committed by
Facebook GitHub Bot
Facebook GitHub Bot
parent
1dc992bf26
commit
a80c96c0de
@@ -0,0 +1,606 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import pickle
|
||||
import time
|
||||
from multiprocessing.pool import ThreadPool
|
||||
|
||||
import faiss
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from faiss.contrib.datasets_fb import dataset_from_name
|
||||
except ImportError:
|
||||
from faiss.contrib.datasets import dataset_from_name
|
||||
|
||||
from faiss.contrib.evaluation import OperatingPointsWithRanges
|
||||
from faiss.contrib.ivf_tools import replace_ivf_quantizer
|
||||
|
||||
#################################################################
|
||||
# Preassigned search functions
|
||||
#################################################################
|
||||
|
||||
|
||||
def search_preassigned(xq, k, index, quantizer, batch_size=0):
|
||||
"""
|
||||
Explicitly call the coarse quantizer and the search_preassigned
|
||||
on the index.
|
||||
"""
|
||||
n, d = xq.shape
|
||||
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')
|
||||
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)
|
||||
return D, I
|
||||
|
||||
|
||||
def tiled_search_preassigned(xq, k, index, quantizer, batch_size=32768):
|
||||
"""
|
||||
Explicitly call the coarse quantizer and the search_preassigned
|
||||
on the index. Allow overlapping between coarse quantization and
|
||||
scanning the inverted lists.
|
||||
"""
|
||||
n, d = xq.shape
|
||||
|
||||
# prepare a thread that will run the quantizer
|
||||
qq_pool = ThreadPool(1)
|
||||
nprobe = index.nprobe
|
||||
|
||||
def coarse_quant(i0):
|
||||
if i0 >= n:
|
||||
return None
|
||||
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')
|
||||
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, ))
|
||||
Dq, Iq = qq
|
||||
index.search_preassigned(
|
||||
xq[i0:i1], k, Iq=Iq, Dq=Dq, I=I[i0:i1], D=D[i0:i1])
|
||||
qq = qq_next.get()
|
||||
|
||||
qq_pool.close()
|
||||
return D, I
|
||||
|
||||
|
||||
#################################################################
|
||||
# IVF index objects with a separate coarse quantizer
|
||||
#################################################################
|
||||
|
||||
class SeparateCoarseQuantizationIndex:
|
||||
"""
|
||||
Separately manage the coarse quantizer and the IVF index.
|
||||
"""
|
||||
|
||||
def __init__(self, quantizer, index, bs=-1, seq_tiling=False):
|
||||
self.index = index
|
||||
self.index_ivf = extract_index_ivf(index)
|
||||
if isinstance(self.index_ivf, faiss.IndexIVF):
|
||||
self.index_ivf.parallel_mode
|
||||
self.index_ivf.parallel_mode = 3
|
||||
|
||||
self.quantizer = quantizer
|
||||
assert self.quantizer.d == self.index_ivf.d
|
||||
# populate quantizer if it was not done before
|
||||
if quantizer.ntotal > 0:
|
||||
assert quantizer.ntotal == self.index_ivf.nlist
|
||||
else:
|
||||
centroids = self.index_ivf.quantizer.reconstruct_n()
|
||||
print(f"adding centroids size {centroids.shape} to quantizer")
|
||||
quantizer.train(centroids)
|
||||
quantizer.add(centroids)
|
||||
self.bs = bs
|
||||
self.seq_tiling = seq_tiling
|
||||
|
||||
def search(self, xq, k):
|
||||
# perform coarse quantization
|
||||
if isinstance(self.index, faiss.IndexPreTransform):
|
||||
# print("applying pre-transform")
|
||||
assert self.index.chain.size() == 1
|
||||
xq = self.index.chain.at(0).apply(xq)
|
||||
if self.bs <= 0:
|
||||
# non batched
|
||||
nprobe = self.index_ivf.nprobe
|
||||
Dq, Iq = self.quantizer.search(xq, nprobe)
|
||||
|
||||
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)
|
||||
else:
|
||||
return tiled_search_preassigned(
|
||||
xq, k, self.index_ivf, self.quantizer, self.bs)
|
||||
|
||||
|
||||
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
|
||||
if isinstance(index, faiss.IndexPreTransform):
|
||||
index = faiss.downcast_index(index.index)
|
||||
ngpu = index.count()
|
||||
self.pool = ThreadPool(ngpu)
|
||||
self.bs = bs
|
||||
if bs > 0:
|
||||
self.q_pool = ThreadPool(1)
|
||||
|
||||
def __del__(self):
|
||||
self.pool.close()
|
||||
if self.bs > 0:
|
||||
self.q_pool.close()
|
||||
|
||||
def search(self, xq, k):
|
||||
nq = len(xq)
|
||||
# perform coarse quantization
|
||||
index = self.cpu_index
|
||||
if isinstance(self.cpu_index, faiss.IndexPreTransform):
|
||||
assert index.chain.size() == 1
|
||||
xq = self.cpu_index.chain.at(0).apply(xq)
|
||||
index = faiss.downcast_index(index.index)
|
||||
ngpu = index.count()
|
||||
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')
|
||||
bs = self.bs
|
||||
if bs <= 0:
|
||||
|
||||
Dq, Iq = self.quantizer.search(xq, nprobe)
|
||||
|
||||
def do_search(rank):
|
||||
gpu_index = faiss.downcast_index(index.at(rank))
|
||||
Dall[rank], Iall[rank] = gpu_index.search_preassigned(
|
||||
xq, k, Iq, Dq)
|
||||
list(self.pool.map(do_search, range(ngpu)))
|
||||
else:
|
||||
qq_pool = self.q_pool
|
||||
bs = self.bs
|
||||
|
||||
def coarse_quant(i0):
|
||||
if i0 >= nq:
|
||||
return None
|
||||
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)
|
||||
|
||||
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 = 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 """
|
||||
try:
|
||||
return faiss.extract_index_ivf(index)
|
||||
except RuntimeError:
|
||||
if index.__class__ == faiss.IndexPreTransform:
|
||||
index = faiss.downcast_index(index.index)
|
||||
if isinstance(index, faiss.GpuIndexIVF):
|
||||
return index
|
||||
raise RuntimeError(f"could not extract IVF index from {index}")
|
||||
|
||||
|
||||
def set_index_parameter(index, name, val):
|
||||
"""
|
||||
Index parameter setting that works on the index lookalikes defined above
|
||||
"""
|
||||
if index.__class__ == SeparateCoarseQuantizationIndex:
|
||||
if name == "nprobe":
|
||||
set_index_parameter(index.index_ivf, name, val)
|
||||
elif name.startswith("quantizer_"):
|
||||
set_index_parameter(
|
||||
index.quantizer, name[name.find("_") + 1:], val)
|
||||
else:
|
||||
raise RuntimeError()
|
||||
return
|
||||
|
||||
if index.__class__ == ShardedGPUIndex:
|
||||
if name == "nprobe":
|
||||
set_index_parameter(index.cpu_index, name, val)
|
||||
elif name.startswith("quantizer_"):
|
||||
set_index_parameter(
|
||||
index.quantizer, name[name.find("_") + 1:], val)
|
||||
else:
|
||||
raise RuntimeError()
|
||||
return
|
||||
|
||||
# then it's a Faiss index
|
||||
index = faiss.downcast_index(index)
|
||||
|
||||
if isinstance(index, faiss.IndexPreTransform):
|
||||
set_index_parameter(index.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)
|
||||
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)):
|
||||
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)
|
||||
elif name == "efSearch":
|
||||
index.hnsw.efSearch
|
||||
index.hnsw.efSearch = int(val)
|
||||
elif name == "nprobe":
|
||||
index_ivf = extract_index_ivf(index)
|
||||
index_ivf.nprobe
|
||||
index_ivf.nprobe = int(val)
|
||||
else:
|
||||
raise RuntimeError(f"could not set param {name} on {index}")
|
||||
|
||||
|
||||
#####################################################################
|
||||
# Driver routine
|
||||
#####################################################################
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
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('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",
|
||||
"multi_gpu_sharded1_ivf",
|
||||
"multi_gpu_Csharded1", "multi_gpu_Csharded1_flat",
|
||||
"multi_gpu_Csharded1_ivf",
|
||||
],
|
||||
help="how to search"
|
||||
)
|
||||
aa('--ivf_quant_nlist', type=int, default=1024,
|
||||
help="nb of invlists 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")
|
||||
|
||||
args = parser.parse_args()
|
||||
print("args:", args)
|
||||
|
||||
if not args.quiet:
|
||||
# log some stats about the machine
|
||||
os.system("grep -m1 'model name' < /proc/cpuinfo")
|
||||
os.system("grep -E 'MemTotal|MemFree' /proc/meminfo")
|
||||
os.system("nvidia-smi")
|
||||
|
||||
print("prepare dataset", args.db)
|
||||
ds = dataset_from_name(args.db)
|
||||
print(ds)
|
||||
|
||||
print("Faiss nb GPUs:", faiss.get_num_gpus())
|
||||
|
||||
xq = ds.get_queries()
|
||||
if args.nq > len(xq):
|
||||
xqx = []
|
||||
n = 0
|
||||
while n < args.nq:
|
||||
xqx.append(xq[:args.nq - n])
|
||||
n += len(xqx[-1])
|
||||
print(f"increased nb queries from {len(xq)} to {n}")
|
||||
xq = np.vstack(xqx)
|
||||
|
||||
if args.nt != -1:
|
||||
print("setting nb openmp threads to", args.nt)
|
||||
faiss.omp_set_num_threads(args.nt)
|
||||
|
||||
print("loading index")
|
||||
|
||||
if args.mmap:
|
||||
io_flag = faiss.IO_FLAG_READ_ONLY | faiss.IO_FLAG_MMAP
|
||||
else:
|
||||
io_flag = 0
|
||||
|
||||
print(f"load index {args.indexname} {io_flag=:x}")
|
||||
index = faiss.read_index(args.indexname, io_flag)
|
||||
index_ivf = faiss.extract_index_ivf(index)
|
||||
|
||||
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
|
||||
]
|
||||
)
|
||||
|
||||
# prepare options for GPU clone
|
||||
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.useFloat16 = args.useFloat16
|
||||
co.useFloat16CoarseQuantizer = args.useFloat16CoarseQuantizer
|
||||
co.usePrecomputed = args.usePrecomputed
|
||||
co.shard_type = args.shard_type
|
||||
|
||||
if args.search_type == "cpu":
|
||||
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.restrict_range("nprobe", 2049)
|
||||
elif args.search_type == "gpu_tiled":
|
||||
print("move index to 1 GPU")
|
||||
new_quantizer = faiss.IndexFlatL2(index_ivf.d)
|
||||
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.restrict_range("nprobe", 2049)
|
||||
index = SeparateCoarseQuantizationIndex(
|
||||
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")
|
||||
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)]
|
||||
)
|
||||
op.restrict_range("nprobe", 1025)
|
||||
elif args.search_type == "gpu_flat_quantizer":
|
||||
index_ivf = faiss.extract_index_ivf(index)
|
||||
new_quantizer = faiss.IndexFlatL2(index_ivf.d)
|
||||
replace_ivf_quantizer(index_ivf, new_quantizer)
|
||||
res = faiss.StandardGpuResources()
|
||||
index = faiss.index_cpu_to_gpu(res, 0, index, co)
|
||||
op.restrict_range("nprobe", 2049)
|
||||
elif args.search_type == "cpu_flat_gpu_quantizer":
|
||||
index_ivf = faiss.extract_index_ivf(index)
|
||||
quantizer = faiss.IndexFlatL2(index_ivf.d)
|
||||
res = faiss.StandardGpuResources()
|
||||
quantizer = faiss.index_cpu_to_gpu(res, 0, quantizer, co)
|
||||
index = SeparateCoarseQuantizationIndex(
|
||||
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.restrict_range("nprobe", 2049)
|
||||
elif args.search_type in (
|
||||
"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"):
|
||||
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)
|
||||
co.shard
|
||||
co.shard = True
|
||||
gpus = list(range(faiss.get_num_gpus()))
|
||||
res = [faiss.StandardGpuResources() for _ in gpus]
|
||||
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)]
|
||||
)
|
||||
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")
|
||||
quantizer.train(centroids)
|
||||
quantizer.add(centroids)
|
||||
co.shard = False
|
||||
quantizer = faiss.index_cpu_to_gpu_multiple_py(
|
||||
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.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)
|
||||
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"):
|
||||
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)]
|
||||
)
|
||||
index = faiss.index_cpu_to_all_gpus(index, co)
|
||||
elif args.search_type == "multi_gpu_Csharded1_flat":
|
||||
new_quantizer = faiss.IndexFlatL2(index_ivf.d)
|
||||
quantizer_hnsw = replace_ivf_quantizer(index_ivf, new_quantizer)
|
||||
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 = faiss.index_cpu_to_all_gpus(index, co)
|
||||
else:
|
||||
raise RuntimeError()
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
totex = op.num_experiments()
|
||||
rs = np.random.RandomState(123)
|
||||
if totex < args.n_autotune:
|
||||
experiments = rs.permutation(totex - 2) + 1
|
||||
else:
|
||||
experiments = rs.randint(
|
||||
totex - 2, size=args.n_autotune - 2, replace=False)
|
||||
|
||||
experiments = [0, totex - 1] + list(experiments)
|
||||
print(f"total nb experiments {totex}, running {len(experiments)}")
|
||||
|
||||
print("perform search")
|
||||
gt = ds.get_groundtruth(100)
|
||||
|
||||
# piggyback on operating points so that this gets stored in the stats file
|
||||
op.all_experiments = []
|
||||
op.platform = {
|
||||
"loadavg": open("/proc/loadavg", "r").readlines(),
|
||||
"procesor": [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()
|
||||
}
|
||||
op.args = args
|
||||
if args.stats:
|
||||
print(f"storing stats in {args.stats} after each experiment")
|
||||
|
||||
for cno in experiments:
|
||||
key = op.cno_to_key(cno)
|
||||
parameters = op.get_parameters(key)
|
||||
print(f"{cno=:4d} {str(parameters):50}", end=": ", flush=True)
|
||||
|
||||
(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}", )
|
||||
continue
|
||||
|
||||
for name, val in parameters.items():
|
||||
set_index_parameter(index, name, val)
|
||||
|
||||
if cno == 0:
|
||||
# warmup
|
||||
for _ in range(5):
|
||||
D, I = index.search(xq, 100)
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
D, I = index.search(xq, 100)
|
||||
except RuntimeError as e:
|
||||
print(f"ERROR {e}")
|
||||
continue
|
||||
t1 = time.time()
|
||||
|
||||
recalls = {}
|
||||
for rank in 1, 10, 100:
|
||||
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
|
||||
})
|
||||
|
||||
if args.stats:
|
||||
pickle.dump(op, open(args.stats, "wb"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -310,3 +310,39 @@ class DatasetMusic100(Dataset):
|
||||
assert k <= 100
|
||||
gt = gt[:, :k]
|
||||
return gt
|
||||
|
||||
|
||||
|
||||
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':
|
||||
return DatasetSIFT1M()
|
||||
|
||||
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':
|
||||
dbsize = 1000 * int(szsuf[:-1])
|
||||
else:
|
||||
assert False, "did not recognize suffix " + szsuf
|
||||
return DatasetDeep1B(nb=dbsize)
|
||||
|
||||
elif dataset == "music-100":
|
||||
return DatasetMusic100()
|
||||
|
||||
elif dataset == "glove":
|
||||
return DatasetGlove(download=download)
|
||||
|
||||
else:
|
||||
raise RuntimeError("unknown dataset " + dataset)
|
||||
|
||||
@@ -264,3 +264,125 @@ def test_ref_range_results(lims_ref, Dref, Iref,
|
||||
(Ii_new, Di_new) = sort_by_ids(Ii_new, Di_new)
|
||||
np.testing.assert_array_equal(Ii_ref, Ii_new)
|
||||
np.testing.assert_array_almost_equal(Di_ref, Di_new, decimal=5)
|
||||
|
||||
|
||||
###############################################################
|
||||
# OperatingPoints functions
|
||||
# this is the Python version of the AutoTune object in C++
|
||||
|
||||
class OperatingPoints:
|
||||
"""
|
||||
Manages a set of search parameters with associated performance and time.
|
||||
Keeps the Pareto optimal points.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
# list of (key, perf, t)
|
||||
self.operating_points = [
|
||||
# (self.do_nothing_key(), 0.0, 0.0)
|
||||
]
|
||||
self.suboptimal_points = []
|
||||
|
||||
def compare_keys(self, k1, k2):
|
||||
""" return -1 if k1 > k2, 1 if k2 > k1, 0 otherwise """
|
||||
raise NotImplemented
|
||||
|
||||
def do_nothing_key(self):
|
||||
""" parameters to say we do noting, takes 0 time and has 0 performance"""
|
||||
raise NotImplemented
|
||||
|
||||
def is_pareto_optimal(self, perf_new, t_new):
|
||||
for _, perf, t in self.operating_points:
|
||||
if perf >= perf_new and t <= t_new:
|
||||
return False
|
||||
return True
|
||||
|
||||
def predict_bounds(self, key):
|
||||
""" 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 t > min_time:
|
||||
min_time = t
|
||||
if cmp < 0: # key2 < key
|
||||
if perf < max_perf:
|
||||
max_perf = perf
|
||||
return max_perf, min_time
|
||||
|
||||
def should_run_experiment(self, key):
|
||||
(max_perf, min_time) = self.predict_bounds(key)
|
||||
return self.is_pareto_optimal(max_perf, min_time)
|
||||
|
||||
def add_operating_point(self, key, perf, t):
|
||||
if self.is_pareto_optimal(perf, t):
|
||||
i = 0
|
||||
# maybe it shadows some other operating point completely?
|
||||
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))
|
||||
else:
|
||||
i += 1
|
||||
self.operating_points.append((key, perf, t))
|
||||
return True
|
||||
else:
|
||||
self.suboptimal_points.append((key, perf, t))
|
||||
return False
|
||||
|
||||
|
||||
class OperatingPointsWithRanges(OperatingPoints):
|
||||
"""
|
||||
Set of parameters that are each picked from a discrete range of values.
|
||||
An increase of each parameter is assumed to make the operation slower
|
||||
and more accurate.
|
||||
A key = int array of indices in the ordered set of parameters.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
OperatingPoints.__init__(self)
|
||||
# list of (name, values)
|
||||
self.ranges = []
|
||||
|
||||
def add_range(self, name, values):
|
||||
self.ranges.append((name, values))
|
||||
|
||||
def compare_keys(self, k1, k2):
|
||||
if np.all(k1 >= k2):
|
||||
return 1
|
||||
if np.all(k2 >= k1):
|
||||
return -1
|
||||
return 0
|
||||
|
||||
def do_nothing_key(self):
|
||||
return np.zeros(len(self.ranges), dtype=int)
|
||||
|
||||
def num_experiments(self):
|
||||
return np.prod([len(values) for name, values in self.ranges])
|
||||
|
||||
def cno_to_key(self, cno):
|
||||
"""Convert a sequential experiment number to a key"""
|
||||
k = np.zeros(len(self.ranges), dtype=int)
|
||||
for i, (name, values) in enumerate(self.ranges):
|
||||
k[i] = cno % len(values)
|
||||
cno //= len(values)
|
||||
assert cno == 0
|
||||
return k
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
def restrict_range(self, name, max_val):
|
||||
""" 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]
|
||||
values[:] = val2
|
||||
return
|
||||
raise RuntimeError(f"parameter {name} not found")
|
||||
|
||||
@@ -94,6 +94,30 @@ def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None):
|
||||
return lims, dist, indices
|
||||
|
||||
|
||||
def replace_ivf_quantizer(index_ivf, new_quantizer):
|
||||
""" replace the IVF quantizer with a flat quantizer and return the
|
||||
old quantizer"""
|
||||
if new_quantizer.ntotal == 0:
|
||||
centroids = index_ivf.quantizer.reconstruct_n()
|
||||
new_quantizer.train(centroids)
|
||||
new_quantizer.add(centroids)
|
||||
else:
|
||||
assert new_quantizer.ntotal == index_ivf.nlist
|
||||
|
||||
# cleanly dealloc old quantizer
|
||||
old_own = index_ivf.own_fields
|
||||
index_ivf.own_fields = False
|
||||
old_quantizer = faiss.downcast_index(index_ivf.quantizer)
|
||||
old_quantizer.this.own(old_own)
|
||||
index_ivf.quantizer = new_quantizer
|
||||
|
||||
if hasattr(index_ivf, "referenced_objects"):
|
||||
index_ivf.referenced_objects.append(new_quantizer)
|
||||
else:
|
||||
index_ivf.referenced_objects = [new_quantizer]
|
||||
return old_quantizer
|
||||
|
||||
|
||||
class BigBatchSearcher:
|
||||
"""
|
||||
Object that manages all the data related to the computation
|
||||
|
||||
Reference in New Issue
Block a user