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}")