cuvs.neighbors.mg.ivf_pq

cuvs.neighbors.mg.ivf_pq#

8 min read time

Applies to Linux

Submodules#

cuvs.neighbors.mg.ivf_pq.ivf_pq

Classes#

`Index`	Multi-GPU IVF-PQ index object. Stores the trained multi-GPU IVF-PQ
`IndexParams`	IndexParams(distribution_mode=u'sharded', , *kwargs)
`SearchParams`	SearchParams(n_probes=20, , search_mode=u'load_balancer', merge_mode=u'merge_on_root_rank', n_rows_per_batch=1000, *kwargs)

Functions#

`build`(*args[, resources])	build(IndexParams index_params, dataset, resources=None)
`distribute`(*args[, resources])	distribute(filename, resources=None)
`extend`(*args[, resources])	extend(Index index, new_vectors, new_indices=None, resources=None)
`load`(*args[, resources])	load(filename, resources=None)
`save`(*args[, resources])	save(Index index, filename, resources=None)
`search`(*args[, resources])	search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)

Package Contents#

class cuvs.neighbors.mg.ivf_pq.Index#

Multi-GPU IVF-PQ index object. Stores the trained multi-GPU IVF-PQ index state which can be used to perform nearest neighbors searches across multiple GPUs.

static __reduce__(*args, **kwargs)#: Index.__reduce_cython__(self)

static __repr__(*args, **kwargs)#

static __setstate__(*args, **kwargs)#: Index.__setstate_cython__(self, __pyx_state)

class cuvs.neighbors.mg.ivf_pq.IndexParams#

Bases: cuvs.neighbors.ivf_pq.ivf_pq.IndexParams

IndexParams(distribution_mode=u’sharded’, *, **kwargs)

Parameters to build multi-GPU IVF-PQ index for efficient search. Extends single-GPU IndexParams with multi-GPU specific parameters.

distribution_modestr, default = “sharded”
Distribution mode for multi-GPU setup. Valid values: [“replicated”, “sharded”]

**kwargs : Additional parameters passed to single-GPU IndexParams

static __reduce__(*args, **kwargs)#: IndexParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#: IndexParams.__setstate_cython__(self, __pyx_state)

get_handle()#: IndexParams.get_handle(self)

class cuvs.neighbors.mg.ivf_pq.SearchParams#

Bases: cuvs.neighbors.ivf_pq.ivf_pq.SearchParams

SearchParams(n_probes=20, *, search_mode=u’load_balancer’, merge_mode=u’merge_on_root_rank’, n_rows_per_batch=1000, **kwargs)

Parameters to search multi-GPU IVF-PQ index.

static __reduce__(*args, **kwargs)#: SearchParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#: SearchParams.__setstate_cython__(self, __pyx_state)

get_handle()#: SearchParams.get_handle(self)

cuvs.neighbors.mg.ivf_pq.build(*args, resources=None, **kwargs)#

build(IndexParams index_params, dataset, resources=None)

Build the multi-GPU IVF-PQ index from the dataset for efficient search.

index_params : cuvs.neighbors.ivf_pq.IndexParams dataset : Array interface compliant matrix shape (n_samples, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU IVF-PQ, the dataset MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

index: py:class:cuvs.neighbors.ivf_pq.Index
>>> import numpy as np
>>> from cuvs.neighbors.mg import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU IVF-PQ, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = ivf_pq.IndexParams(metric="sqeuclidean")
>>> index = ivf_pq.build(build_params, dataset)
>>> distances, neighbors = ivf_pq.search(
...     ivf_pq.SearchParams(),
...     index, dataset, k)
>>> # Results are already in host memory (NumPy arrays)

cuvs.neighbors.mg.ivf_pq.distribute(*args, resources=None, **kwargs)#

distribute(filename, resources=None)

Distribute a single-GPU IVF-PQ index across multiple GPUs from a file.

filenamestr
The filename to distribute the index from.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

indexIndex
The distributed index.
>>> from cuvs.neighbors.mg import ivf_pq
>>> index = ivf_pq.distribute("single_gpu_index.bin")

cuvs.neighbors.mg.ivf_pq.extend(*args, resources=None, **kwargs)#

extend(Index index, new_vectors, new_indices=None, resources=None)

Extend the multi-GPU IVF-PQ index with new vectors.

index : cuvs.neighbors.ivf_pq.Index new_vectors : Array interface compliant matrix shape (n_new_vectors, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU IVF-PQ, new_vectors MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

new_indicesArray interface compliant matrix shape (n_new_vectors,)
, optional If provided, these indices will be used for the new vectors. If not provided, indices will be automatically assigned. IMPORTANT: Must be in host memory (CPU) for multi-GPU IVF-PQ.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.
>>> import numpy as np
>>> from cuvs.neighbors.mg import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_new_vectors = 1000
>>> # For multi-GPU IVF-PQ, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> new_vectors = np.random.random_sample(
...     (n_new_vectors, n_features)).astype(np.float32)
>>> new_indices = np.arange(n_samples, n_new_vectors, dtype=np.int64)
>>> build_params = ivf_pq.IndexParams(metric="sqeuclidean")
>>> index = ivf_pq.build(build_params, dataset)
>>> ivf_pq.extend(index, new_vectors, new_indices)

cuvs.neighbors.mg.ivf_pq.load(*args, resources=None, **kwargs)#

load(filename, resources=None)

Deserialize the multi-GPU IVF-PQ index from a file.

filenamestr
The filename to deserialize the index from.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

indexIndex
The deserialized index.
>>> from cuvs.neighbors.mg import ivf_pq
>>> index = ivf_pq.load("index.bin")

cuvs.neighbors.mg.ivf_pq.save(*args, resources=None, **kwargs)#

save(Index index, filename, resources=None)

Serialize the multi-GPU IVF-PQ index to a file.

index : cuvs.neighbors.ivf_pq.Index filename : str

The filename to serialize the index to.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.
>>> import numpy as np
>>> from cuvs.neighbors.mg import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> # For multi-GPU IVF-PQ, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> build_params = ivf_pq.IndexParams(metric="sqeuclidean")
>>> index = ivf_pq.build(build_params, dataset)
>>> ivf_pq.save(index, "index.bin")

cuvs.neighbors.mg.ivf_pq.search(*args, resources=None, **kwargs)#

search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)

Search the multi-GPU IVF-PQ index for the k-nearest neighbors of each query.

search_params : cuvs.neighbors.ivf_pq.SearchParams index : cuvs.neighbors.ivf_pq.Index queries : Array interface compliant matrix shape (n_queries, dim)

Supported dtype [float32, float16, int8, uint8] IMPORTANT: For multi-GPU IVF-PQ, queries MUST be in host memory (CPU). If using CuPy/device arrays, transfer to host with array.get() or cp.asnumpy(array).

kint
The number of neighbors to search for each query.

neighborsArray interface compliant matrix shape (n_queries, k), optional
If provided, this array will be filled with the indices of the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU IVF-PQ.

distancesArray interface compliant matrix shape (n_queries, k), optional
If provided, this array will be filled with the distances to the k-nearest neighbors. If not provided, a new host array will be allocated. IMPORTANT: Must be in host memory (CPU) for multi-GPU IVF-PQ.

resourcesOptional cuVS Multi-GPU Resource handle for reusing CUDA resources.
If Multi-GPU Resources aren’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If resources are supplied, you will need to explicitly synchronize yourself by calling resources.sync() before accessing the output.

distancesnumpy.ndarray
The distances to the k-nearest neighbors for each query (in host memory).

neighborsnumpy.ndarray
The indices of the k-nearest neighbors for each query (in host memory).
>>> import numpy as np
>>> from cuvs.neighbors.mg import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> # For multi-GPU IVF-PQ, use host (NumPy) arrays
>>> dataset = np.random.random_sample((n_samples, n_features)).astype(
...     np.float32)
>>> queries = np.random.random_sample((n_queries, n_features)).astype(
...     np.float32)
>>> build_params = ivf_pq.IndexParams(metric="sqeuclidean")
>>> index = ivf_pq.build(build_params, dataset)
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
...                                       index, queries, k)
>>> # Results are already in host memory (NumPy arrays)