def test_shards(self):
k = 32
ref_index = faiss.IndexFlatL2(d)
print('ref search')
ref_index.add(xb)
_Dref, Iref = ref_index.search(xq, k)
print(Iref[:5, :6])
shard_index = faiss.IndexShards(d)
shard_index_2 = faiss.IndexShards(d, True, False)
ni = 3
for i in range(ni):
i0 = int(i * nb / ni)
i1 = int((i + 1) * nb / ni)
index = faiss.IndexFlatL2(d)
index.add(xb[i0:i1])
shard_index.add_shard(index)
index_2 = faiss.IndexFlatL2(d)
irm = faiss.IndexIDMap(index_2)
shard_index_2.add_shard(irm)
# test parallel add
shard_index_2.verbose = True
shard_index_2.add(xb)
for test_no in range(3):
with_threads = test_no == 1
print('shard search test_no = %d' % test_no)
if with_threads:
remember_nt = faiss.omp_get_max_threads()
faiss.omp_set_num_threads(1)
shard_index.threaded = True
else:
shard_index.threaded = False
if test_no != 2:
_D, I = shard_index.search(xq, k)
else:
_D, I = shard_index_2.search(xq, k)
print(I[:5, :6])
if with_threads:
faiss.omp_set_num_threads(remember_nt)
ndiff = (I != Iref).sum()
print('%d / %d differences' % (ndiff, nq * k))
assert (ndiff < nq * k / 1000.)
def test_shard_flag_propagation(self):
d = 64 # dimension
nb = 1000
rs = np.random.RandomState(1234)
xb = rs.rand(nb, d).astype('float32')
nlist = 10
quantizer1 = faiss.IndexFlatL2(d)
quantizer2 = faiss.IndexFlatL2(d)
index1 = faiss.IndexIVFFlat(quantizer1, d, nlist)
index2 = faiss.IndexIVFFlat(quantizer2, d, nlist)
index = faiss.IndexShards(d, True)
index.add_shard(index1)
index.add_shard(index2)
self.assertFalse(index.is_trained)
index.train(xb)
self.assertTrue(index.is_trained)
self.assertEqual(index.ntotal, 0)
index.add(xb)
self.assertEqual(index.ntotal, nb)
index.remove_shard(index2)
self.assertEqual(index.ntotal, nb / 2)
index.remove_shard(index1)
self.assertEqual(index.ntotal, 0)
def __init__(self, shard_dir, nprobe: int = 4):
"""
For deploying multiple, pre-made IVF indexes as shards.
(intended for on-disk indexes that do not fit in memory)
Note: The index shards must be true partitions with no overlapping ids
:param shard_dir: Dir containing faiss index shards
:param nprobe: Number of clusters to visit during search
(speed accuracy trade-off)
"""
super().__init__()
self.paths_to_shards = self.get_index_paths(shard_dir)
self.nprobe = nprobe
# Load shards
shards = list()
for shard_path in self.paths_to_shards:
shard = self.load_shard(path_to_shard=shard_path,
nprobe=self.nprobe)
shards.append(shard)
# Merge shards
self.index = faiss.IndexShards(512,
threaded=True,
successive_ids=False)
for shard in shards:
self.index.add_shard(shard)
def test_shards(self):
index = faiss.IndexShards(d)
for i in range(3):
sub_index = faiss.IndexFlatL2(d)
sub_index.add(xb)
index.add_shard(sub_index)
gc.collect()
index.search(xb, 10)
def execute(cls, ctx, op):
(y, ), device_id, xp = as_same_device([ctx[op.input.key]],
device=op.device,
ret_extra=True)
indexes = [
_load_index(ctx[index.key], device_id) for index in op.inputs[1:]
]
with device(device_id):
y = xp.ascontiguousarray(y, dtype=np.float32)
if len(indexes) == 1:
index = indexes[0]
else:
index = faiss.IndexShards(indexes[0].d)
[index.add_shard(ind) for ind in indexes]
if op.metric == "cosine":
# faiss does not support cosine distances directly,
# data needs to be normalize before searching,
# refer to:
# https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
faiss.normalize_L2(y)
if op.nprobe is not None:
index.nprobe = op.nprobe
if device_id >= 0: # pragma: no cover
n = y.shape[0]
k = op.n_neighbors
distances = xp.empty((n, k), dtype=xp.float32)
indices = xp.empty((n, k), dtype=xp.int64)
index.search_c(
n,
_swig_ptr_from_cupy_float32_array(y),
k,
_swig_ptr_from_cupy_float32_array(distances),
_swig_ptr_from_cupy_int64_array(indices),
)
else:
distances, indices = index.search(y, op.n_neighbors)
if op.return_distance:
if index.metric_type == faiss.METRIC_L2:
# make it equivalent to `pairwise.euclidean_distances`
distances = xp.sqrt(distances, out=distances)
elif op.metric == "cosine":
# make it equivalent to `pairwise.cosine_distances`
distances = xp.subtract(1, distances, out=distances)
ctx[op.outputs[0].key] = distances
ctx[op.outputs[-1].key] = indices
def index_init_gpu(ngpus, feat_dim):
flat_config = []
for i in range(ngpus):
cfg = faiss.GpuIndexFlatConfig()
cfg.useFloat16 = False
cfg.device = i
flat_config.append(cfg)
res = [faiss.StandardGpuResources() for i in range(ngpus)]
indexes = [faiss.GpuIndexFlatL2(res[i], feat_dim, flat_config[i]) for i in range(ngpus)]
index = faiss.IndexShards(feat_dim)
for sub_index in indexes:
index.add_shard(sub_index)
index.reset()
return index
def load_faiss_index(self):
from .url_utils import use_s3
item_index_input_dir = use_s3('%sfaiss/item_index/' %
self.model_in_path)
self.faiss_index = faiss.IndexShards(self.item_embedding_size, True,
False)
partition_count = self.get_partition_count()
for rank in range(partition_count):
item_index_input_path = '%spart_%d_%d.dat' % (
item_index_input_dir, partition_count, rank)
item_index_stream = _mindalpha.InputStream(item_index_input_path)
item_index_reader = faiss.PyCallbackIOReader(
item_index_stream.read)
index = faiss.read_index(item_index_reader)
self.faiss_index.add_shard(index)
print('faiss index ntotal: %d' % self.faiss_index.ntotal)
def test_manual_build_faiss_index(setup):
d = 8
n = 50
n_test = 10
x = np.random.RandomState(0).rand(n, d).astype(np.float32)
y = np.random.RandomState(0).rand(n_test, d).astype(np.float32)
nn = NearestNeighbors(algorithm='kd_tree')
nn.fit(x)
_, expected_indices = nn.kneighbors(y, 5)
# test brute-force search
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X,
'Flat',
None,
random_state=0,
same_distribution=True)
faiss_index = index.execute().fetch()
index_shards = faiss.IndexShards(d)
for ind in faiss_index:
shard = _load_index(ind, -1)
index_shards.add_shard(shard)
faiss_index = index_shards
faiss_index.nprob = 10
_, indices = faiss_index.search(y, k=5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
# test one chunk, brute force
X = mt.tensor(x, chunk_size=50)
index = build_faiss_index(X,
'Flat',
None,
random_state=0,
same_distribution=True)
faiss_index = _load_index(index.execute().fetch(), -1)
faiss_index.nprob = 10
_, indices = faiss_index.search(y, k=5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
# test train, same distribution
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X,
'IVF30,Flat',
30,
random_state=0,
same_distribution=True)
faiss_index = _load_index(index.execute().fetch(), -1)
assert isinstance(faiss_index, faiss.IndexIVFFlat)
assert faiss_index.ntotal == n
assert len(tile(index).chunks) == 1
# test train, distributions are variant
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X,
'IVF10,Flat',
None,
random_state=0,
same_distribution=False)
faiss_index = index.execute().fetch()
assert len(faiss_index) == 5
for ind in faiss_index:
ind = _load_index(ind, -1)
assert isinstance(ind, faiss.IndexIVFFlat)
assert ind.ntotal == 10
# test more index type
index = build_faiss_index(X, 'PCAR6,IVF8_HNSW32,SQ8', 10, random_state=0)
faiss_index = index.execute().fetch()
assert len(faiss_index) == 5
for ind in faiss_index:
ind = _load_index(ind, -1)
assert isinstance(ind, faiss.IndexPreTransform)
assert ind.ntotal == 10
# test one chunk, train
X = mt.tensor(x, chunk_size=50)
index = build_faiss_index(X,
'IVF30,Flat',
30,
random_state=0,
same_distribution=True)
faiss_index = _load_index(index.execute().fetch(), -1)
assert isinstance(faiss_index, faiss.IndexIVFFlat)
assert faiss_index.ntotal == n
# test wrong index
with pytest.raises(ValueError):
build_faiss_index(X, 'unknown_index', None)
# test unknown metric
with pytest.raises(ValueError):
build_faiss_index(X, 'Flat', None, metric='unknown_metric')
def testManualBuildFaissIndex(self):
d = 8
n = 50
n_test = 10
x = np.random.RandomState(0).rand(n, d).astype(np.float32)
y = np.random.RandomState(0).rand(n_test, d).astype(np.float32)
nn = NearestNeighbors(algorithm='kd_tree')
nn.fit(x)
_, expected_indices = nn.kneighbors(y, 5)
for index_type in ['object', 'filename', 'bytes']:
# test brute-force search
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X, 'Flat', None, random_state=0,
same_distribution=True, return_index_type=index_type)
faiss_index = self.executor.execute_tileable(index)
index_shards = faiss.IndexShards(d)
for ind in faiss_index:
shard = _load_index(None, index.op, ind, -1)
index_shards.add_shard(shard)
faiss_index = index_shards
faiss_index.nprob = 10
_, indices = faiss_index.search(y, k=5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
# test one chunk, brute force
X = mt.tensor(x, chunk_size=50)
index = build_faiss_index(X, 'Flat', None, random_state=0,
same_distribution=True, return_index_type='object')
faiss_index = self.executor.execute_tileable(index)[0]
faiss_index.nprob = 10
_, indices = faiss_index.search(y, k=5)
np.testing.assert_array_equal(indices, expected_indices.fetch())
# test train, same distribution
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X, 'IVF30,Flat', 30, random_state=0,
same_distribution=True, return_index_type='object')
faiss_index = self.executor.execute_tileable(index)[0]
self.assertIsInstance(faiss_index, faiss.IndexIVFFlat)
self.assertEqual(faiss_index.ntotal, n)
self.assertEqual(len(get_tiled(index).chunks), 1)
# test train, distributions are variant
X = mt.tensor(x, chunk_size=10)
index = build_faiss_index(X, 'IVF10,Flat', None, random_state=0,
same_distribution=False, return_index_type='object')
faiss_index = self.executor.execute_tileable(index)
self.assertEqual(len(faiss_index), 5)
for ind in faiss_index:
self.assertIsInstance(ind, faiss.IndexIVFFlat)
self.assertEqual(ind.ntotal, 10)
# test one chunk, train
X = mt.tensor(x, chunk_size=50)
index = build_faiss_index(X, 'IVF30,Flat', 30, random_state=0,
same_distribution=True, return_index_type='object')
faiss_index = self.executor.execute_tileable(index)[0]
self.assertIsInstance(faiss_index, faiss.IndexIVFFlat)
self.assertEqual(faiss_index.ntotal, n)
# test wrong index
with self.assertRaises(ValueError):
build_faiss_index(X, 'unknown_index', None)
# test unknown metric
with self.assertRaises(ValueError):
build_faiss_index(X, 'Flat', None, metric='unknown_metric')
def test_shards(self):
k = 32
ref_index = faiss.IndexFlatL2(d)
print('ref search')
ref_index.add(xb)
_Dref, Iref = ref_index.search(xq, k)
print(Iref[:5, :6])
# there is a OpenMP bug in this configuration, so disable threading
if sys.platform == "darwin" and "Clang 12" in sys.version:
nthreads = faiss.omp_get_max_threads()
faiss.omp_set_num_threads(1)
else:
nthreads = None
shard_index = faiss.IndexShards(d)
shard_index_2 = faiss.IndexShards(d, True, False)
ni = 3
for i in range(ni):
i0 = int(i * nb / ni)
i1 = int((i + 1) * nb / ni)
index = faiss.IndexFlatL2(d)
index.add(xb[i0:i1])
shard_index.add_shard(index)
index_2 = faiss.IndexFlatL2(d)
irm = faiss.IndexIDMap(index_2)
shard_index_2.add_shard(irm)
# test parallel add
shard_index_2.verbose = True
shard_index_2.add(xb)
for test_no in range(3):
with_threads = test_no == 1
print('shard search test_no = %d' % test_no)
if with_threads:
remember_nt = faiss.omp_get_max_threads()
faiss.omp_set_num_threads(1)
shard_index.threaded = True
else:
shard_index.threaded = False
if test_no != 2:
_D, I = shard_index.search(xq, k)
else:
_D, I = shard_index_2.search(xq, k)
print(I[:5, :6])
if with_threads:
faiss.omp_set_num_threads(remember_nt)
ndiff = (I != Iref).sum()
print('%d / %d differences' % (ndiff, nq * k))
assert(ndiff < nq * k / 1000.)
if nthreads is not None:
faiss.omp_set_num_threads(nthreads)
#test_time(cpu_index, xb, xq, K, eD, eI)
#cpu_index.nprobe = 10
#test_time(cpu_index, xb, xq, K, eD, eI)
#cpu_index.nprobe = 100
cpu_index.nprobe = 50
test_time(cpu_index, xb, xq, K, eD, eI)
print(cpu_index_l2)
test_time(cpu_index_l2, xb, xq, K, eD, eI)
#print(index)
#test_time(index, xb, xq, K, eD, eI)
print(index_gpu)
test_time(index_gpu, xb, xq, K, eD, eI)
# test sharding
# JK no sharding
idxs = faiss.IndexShards(d)
for i in range(100):
xb = np.random.random((nb // 2, d)).astype('float32')
xb[:, 0] += np.arange(xb.shape[0]) / 1000.
index = faiss.index_factory(d, "IVF4096,PQ32", faiss.METRIC_L2)
res = faiss.StandardGpuResources()
index_gpu = faiss.index_cpu_to_gpu(res, 3, index)
index_gpu.train(xb)
index_gpu.add(xb)
idxs.add_shard(index_gpu)
xq = xb[:32]
t = time.time()
idxs.search(xq, K)
print(f"{time.time()} - {t}")
