Python faiss.index_cpu_to_gpu() Examples
The following are 11
code examples of faiss.index_cpu_to_gpu().
You can vote up the ones you like or vote down the ones you don't like,
and go to the original project or source file by following the links above each example.
You may also want to check out all available functions/classes of the module
faiss
, or try the search function
.
.
Example #1
| Source File: faiss.py From gntp with MIT License | 7 votes |
|
def _build_approximate_index(self,
data: np.ndarray):
dimensionality = data.shape[1]
nlist = 100 if data.shape[0] > 100 else 2
if self.kernel_name in {'rbf'}:
quantizer = faiss.IndexFlatL2(dimensionality)
cpu_index_flat = faiss.IndexIVFFlat(quantizer, dimensionality, nlist, faiss.METRIC_L2)
else:
quantizer = faiss.IndexFlatIP(dimensionality)
cpu_index_flat = faiss.IndexIVFFlat(quantizer, dimensionality, nlist)
gpu_index_ivf = faiss.index_cpu_to_gpu(self.resource, 0, cpu_index_flat)
gpu_index_ivf.train(data)
gpu_index_ivf.add(data)
self.index = gpu_index_ivf
Example #2
| Source File: losses.py From SO-Net with MIT License | 6 votes |
|
def build_nn_index(self, database):
'''
:param database: numpy array of Nx3
:return: Faiss index, in CPU
'''
# index = faiss.GpuIndexFlatL2(self.res, self.dimension, self.flat_config) # dimension is 3
index_cpu = faiss.IndexFlatL2(self.dimension)
index = faiss.index_cpu_to_gpu(self.res, self.opt.gpu_id, index_cpu)
index.add(database)
return index
Example #3
| Source File: run_index.py From denspi with Apache License 2.0 | 6 votes |
|
def train_coarse_quantizer(data, quantizer_path, num_clusters, hnsw=False, niter=10, cuda=False):
d = data.shape[1]
index_flat = faiss.IndexFlatL2(d)
# make it into a gpu index
if cuda:
res = faiss.StandardGpuResources()
index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)
clus = faiss.Clustering(d, num_clusters)
clus.verbose = True
clus.niter = niter
clus.train(data, index_flat)
centroids = faiss.vector_float_to_array(clus.centroids)
centroids = centroids.reshape(num_clusters, d)
if hnsw:
quantizer = faiss.IndexHNSWFlat(d, 32)
quantizer.hnsw.efSearch = 128
quantizer.train(centroids)
quantizer.add(centroids)
else:
quantizer = faiss.IndexFlatL2(d)
quantizer.add(centroids)
faiss.write_index(quantizer, quantizer_path)
Example #4
| Source File: run_index.py From denspi with Apache License 2.0 | 6 votes |
|
def train_index(data, quantizer_path, trained_index_path, fine_quant='SQ8', cuda=False):
quantizer = faiss.read_index(quantizer_path)
if fine_quant == 'SQ8':
trained_index = faiss.IndexIVFScalarQuantizer(quantizer, quantizer.d, quantizer.ntotal, faiss.METRIC_L2)
elif fine_quant.startswith('PQ'):
m = int(fine_quant[2:])
trained_index = faiss.IndexIVFPQ(quantizer, quantizer.d, quantizer.ntotal, m, 8)
else:
raise ValueError(fine_quant)
if cuda:
if fine_quant.startswith('PQ'):
print('PQ not supported on GPU; keeping CPU.')
else:
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, trained_index)
gpu_index.train(data)
trained_index = faiss.index_gpu_to_cpu(gpu_index)
else:
trained_index.train(data)
faiss.write_index(trained_index, trained_index_path)
Example #5
| Source File: _faiss.py From mars with Apache License 2.0 | 5 votes |
|
def _index_to_gpu(index, device_id): # pragma: no cover
res = faiss.StandardGpuResources()
return faiss.index_cpu_to_gpu(res, device_id, index)
Example #6
| Source File: faiss.py From gntp with MIT License | 5 votes |
|
def _build_exact_index(self,
data: np.ndarray):
dimensionality = data.shape[1]
if self.kernel_name in {'rbf'}:
self.cpu_index_flat = faiss.IndexFlatL2(dimensionality)
else:
self.cpu_index_flat = faiss.IndexFlatIP(dimensionality)
if not self.cpu:
self.index = faiss.index_cpu_to_gpu(self.resource, 0, self.cpu_index_flat)
else:
self.index = self.cpu_index_flat
self.index.add(data)
Example #7
| Source File: lossess.py From RL-GAN-Net with MIT License | 5 votes |
|
def build_nn_index(self, database):
'''
:param database: numpy array of Nx3
:return: Faiss index, in CPU
'''
# index = faiss.GpuIndexFlatL2(self.res, self.dimension, self.flat_config) # dimension is 3
index_cpu = faiss.IndexFlatL2(self.dimension)
index = faiss.index_cpu_to_gpu(self.res, self.opt.gpu_id, index_cpu)
index.add(database)
return index
Example #8
| Source File: lossess.py From RL-GAN-Net with MIT License | 5 votes |
|
def build_nn_index(self, database):
'''
:param database: numpy array of Nx3
:return: Faiss index, in CPU
'''
# index = faiss.GpuIndexFlatL2(self.res, self.dimension, self.flat_config) # dimension is 3
index_cpu = faiss.IndexFlatL2(self.dimension)
index = faiss.index_cpu_to_gpu(self.res, self.opt.gpu_id, index_cpu)
index.add(database)
return index
Example #9
| Source File: faiss_gpu.py From ann-benchmarks with MIT License | 5 votes |
|
def fit(self, X):
X = X.astype(numpy.float32)
self._index = faiss.GpuIndexIVFFlat(self._res, len(X[0]), self._n_bits,
faiss.METRIC_L2)
# self._index = faiss.index_factory(len(X[0]),
# "IVF%d,Flat" % self._n_bits)
# co = faiss.GpuClonerOptions()
# co.useFloat16 = True
# self._index = faiss.index_cpu_to_gpu(self._res, 0,
# self._index, co)
self._index.train(X)
self._index.add(X)
self._index.setNumProbes(self._n_probes)
Example #10
| Source File: embedding_based_indexer.py From forte with Apache License 2.0 | 5 votes |
|
def load(self, path: str, device: Optional[str] = None) -> None:
r"""Load the index and meta data from ``path`` directory.
Args:
path (str): A path to the directory to load the index from.
device (optional str): Device to load the index into. If None,
value will be picked from hyperparameters.
"""
if not os.path.exists(path):
raise ValueError(f"Failed to load the index. {path} "
f"does not exist.")
cpu_index = faiss.read_index(f"{path}/index.faiss")
if device is None:
device = self._config.device
if device.lower().startswith("gpu"):
gpu_resource = faiss.StandardGpuResources()
gpu_id = int(device[3:])
if faiss.get_num_gpus() < gpu_id:
gpu_id = 0
logging.warning("Cannot create the index on device %s. "
"Total number of GPUs on this machine is "
"%s. Using the gpu0 for the index.",
device, faiss.get_num_gpus())
self._index = faiss.index_cpu_to_gpu(
gpu_resource, gpu_id, cpu_index)
else:
self._index = cpu_index
with open(f"{path}/index.meta_data", "rb") as f:
self._meta_data = pickle.load(f)
Example #11
| Source File: faiss_gpu.py From learn-to-cluster with MIT License | 4 votes |
|
def __init__(self,
target,
nprobe=128,
index_factory_str=None,
verbose=False,
mode='proxy',
using_gpu=True):
self._res_list = []
num_gpu = faiss.get_num_gpus()
print('[faiss gpu] #GPU: {}'.format(num_gpu))
size, dim = target.shape
assert size > 0, "size: {}".format(size)
index_factory_str = "IVF{},PQ{}".format(
min(8192, 16 * round(np.sqrt(size))),
32) if index_factory_str is None else index_factory_str
cpu_index = faiss.index_factory(dim, index_factory_str)
cpu_index.nprobe = nprobe
if mode == 'proxy':
co = faiss.GpuClonerOptions()
co.useFloat16 = True
co.usePrecomputed = False
index = faiss.IndexProxy()
for i in range(num_gpu):
res = faiss.StandardGpuResources()
self._res_list.append(res)
sub_index = faiss.index_cpu_to_gpu(
res, i, cpu_index, co) if using_gpu else cpu_index
index.addIndex(sub_index)
elif mode == 'shard':
co = faiss.GpuMultipleClonerOptions()
co.useFloat16 = True
co.usePrecomputed = False
co.shard = True
index = faiss.index_cpu_to_all_gpus(cpu_index,
co,
ngpu=num_gpu)
else:
raise KeyError("Unknown index mode")
index = faiss.IndexIDMap(index)
index.verbose = verbose
# get nlist to decide how many samples used for training
nlist = int([
item for item in index_factory_str.split(",") if 'IVF' in item
][0].replace("IVF", ""))
# training
if not index.is_trained:
indexes_sample_for_train = np.random.randint(
0, size, nlist * 256)
index.train(target[indexes_sample_for_train])
# add with ids
target_ids = np.arange(0, size)
index.add_with_ids(target, target_ids)
self.index = index
