cuvs.neighbors.ivf_pq

cuvs.neighbors.ivf_pq#

10 min read time

Applies to Linux

Submodules#

cuvs.neighbors.ivf_pq.ivf_pq

Classes#

`Index`	IvfPq index object. This object stores the trained IvfPq index state
`IndexParams`	IndexParams(n_lists=1024, *, metric=u'sqeuclidean', metric_arg=2.0, kmeans_n_iters=20, kmeans_trainset_fraction=0.5, pq_bits=8, pq_dim=0, codebook_kind=u'subspace', force_random_rotation=False, add_data_on_build=True, conservative_memory_allocation=False, max_train_points_per_pq_code=256)
`SearchParams`	SearchParams(n_probes=20, *, lut_dtype=np.float32, internal_distance_dtype=np.float32)

Functions#

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

Package Contents#

class cuvs.neighbors.ivf_pq.Index#

IvfPq index object. This object stores the trained IvfPq index state which can be used to perform nearest neighbors searches.

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.ivf_pq.IndexParams#

IndexParams(n_lists=1024, *, metric=u’sqeuclidean’, metric_arg=2.0, kmeans_n_iters=20, kmeans_trainset_fraction=0.5, pq_bits=8, pq_dim=0, codebook_kind=u’subspace’, force_random_rotation=False, add_data_on_build=True, conservative_memory_allocation=False, max_train_points_per_pq_code=256)

Parameters to build index for IvfPq nearest neighbor search

n_listsint, default = 1024
The number of clusters used in the coarse quantizer.

metricstr, default=”sqeuclidean”
String denoting the metric type. Valid values for metric: [“sqeuclidean”, “inner_product”, “euclidean”], where:

sqeuclidean is the euclidean distance without the square root operation, i.e.: distance(a,b) = sum_i (a_i - b_i)^2,

euclidean is the euclidean distance

inner product distance is defined as distance(a, b) = sum_i a_i * b_i.

kmeans_n_itersint, default = 20
The number of iterations searching for kmeans centers during index building.

kmeans_trainset_fractionint, default = 0.5
If kmeans_trainset_fraction is less than 1, then the dataset is subsampled, and only n_samples * kmeans_trainset_fraction rows are used for training.

pq_bitsint, default = 8
The bit length of the vector element after quantization.

pq_dimint, default = 0
The dimensionality of a the vector after product quantization. When zero, an optimal value is selected using a heuristic. Note pq_dim * pq_bits must be a multiple of 8. Hint: a smaller ‘pq_dim’ results in a smaller index size and better search performance, but lower recall. If ‘pq_bits’ is 8, ‘pq_dim’ can be set to any number, but multiple of 8 are desirable for good performance. If ‘pq_bits’ is not 8, ‘pq_dim’ should be a multiple of 8. For good performance, it is desirable that ‘pq_dim’ is a multiple of 32. Ideally, ‘pq_dim’ should be also a divisor of the dataset dim.

codebook_kindstring, default = “subspace”
Valid values [“subspace”, “cluster”]

force_random_rotationbool, default = False
Apply a random rotation matrix on the input data and queries even if dim % pq_dim == 0. Note: if dim is not multiple of pq_dim, a random rotation is always applied to the input data and queries to transform the working space from dim to rot_dim, which may be slightly larger than the original space and and is a multiple of pq_dim (rot_dim % pq_dim == 0). However, this transform is not necessary when dim is multiple of pq_dim (dim == rot_dim, hence no need in adding “extra” data columns / features). By default, if dim == rot_dim, the rotation transform is initialized with the identity matrix. When force_random_rotation == True, a random orthogonal transform matrix is generated regardless of the values of dim and pq_dim.

add_data_on_buildbool, default = True
After training the coarse and fine quantizers, we will populate the index with the dataset if add_data_on_build == True, otherwise the index is left empty, and the extend method can be used to add new vectors to the index.

conservative_memory_allocationbool, default = True
By default, the algorithm allocates more space than necessary for individual clusters (list_data). This allows to amortize the cost of memory allocation and reduce the number of data copies during repeated calls to extend (extending the database). To disable this behavior and use as little GPU memory for the database as possible, set this flat to True.

max_train_points_per_pq_codeint, default = 256
The max number of data points to use per PQ code during PQ codebook training. Using more data points per PQ code may increase the quality of PQ codebook but may also increase the build time. The parameter is applied to both PQ codebook generation methods, i.e., PER_SUBSPACE and PER_CLUSTER. In both cases, we will use pq_book_size * max_train_points_per_pq_code training points to train each codebook.

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

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

class cuvs.neighbors.ivf_pq.SearchParams#

SearchParams(n_probes=20, *, lut_dtype=np.float32, internal_distance_dtype=np.float32)

Supplemental parameters to search IVF-Pq index

n_probes: int
The number of clusters to search.

lut_dtype: default = np.float32
Data type of look up table to be created dynamically at search time. The use of low-precision types reduces the amount of shared memory required at search time, so fast shared memory kernels can be used even for datasets with large dimansionality. Note that the recall is slightly degraded when low-precision type is selected. Possible values [np.float32, np.float16, np.uint8]

internal_distance_dtype: default = np.float32
Storage data type for distance/similarity computation. Possible values [np.float32, np.float16]

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

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

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

build(IndexParams index_params, dataset, resources=None)

Build the IvfPq index from the dataset for efficient search.

index_paramscuvs.neighbors.ivf_pq.IndexParams
Parameters on how to build the index

datasetCUDA array interface compliant matrix shape (n_samples, dim)
Supported dtype [float, int8, uint8]

resourcesOptional cuVS Resource handle for reusing CUDA resources.
If 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: cuvs.neighbors.ivf_pq.Index
>>> import cupy as cp
>>> from cuvs.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> k = 10
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.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)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

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

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

Extend an existing index with new vectors.

The input array can be either CUDA array interface compliant matrix or array interface compliant matrix in host memory.

indexivf_pq.Index
Trained ivf_pq object.

new_vectorsarray interface compliant matrix shape (n_samples, dim)
Supported dtype [float, int8, uint8]

new_indicesarray interface compliant vector shape (n_samples)
Supported dtype [int64]

resourcesOptional cuVS Resource handle for reusing CUDA resources.
If 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 cupy as cp
>>> from cuvs.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
>>> n_rows = 100
>>> more_data = cp.random.random_sample((n_rows, n_features),
...                                     dtype=cp.float32)
>>> indices = n_samples + cp.arange(n_rows, dtype=cp.int64)
>>> index = ivf_pq.extend(index, more_data, indices)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
...                                      index, queries,
...                                      k=10)

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

load(filename, resources=None)

Loads index from file.

Saving / loading the index is experimental. The serialization format is subject to change, therefore loading an index saved with a previous version of cuvs is not guaranteed to work.

filenamestring
Name of the file.

resourcesOptional cuVS Resource handle for reusing CUDA resources.
If 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 : Index

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

save(filename, Index index, bool include_dataset=True, resources=None)

Saves the index to a file.

Saving / loading the index is experimental. The serialization format is subject to change.

filenamestring
Name of the file.

indexIndex
Trained IVF-PQ index.

resourcesOptional cuVS Resource handle for reusing CUDA resources.
If 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 cupy as cp
>>> from cuvs.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
>>> # Serialize and deserialize the ivf_pq index built
>>> ivf_pq.save("my_index.bin", index)
>>> index_loaded = ivf_pq.load("my_index.bin")

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

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

Find the k nearest neighbors for each query.

search_paramscuvs.neighbors.ivf_pq.SearchParams
Parameters on how to search the index

indexcuvs.neighbors.ivf_pq.Index
Trained IvfPq index.

queriesCUDA array interface compliant matrix shape (n_samples, dim)
Supported dtype [float, int8, uint8]

kint
The number of neighbors.

neighborsOptional CUDA array interface compliant matrix shape
(n_queries, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None)

distancesOptional CUDA array interface compliant matrix shape
(n_queries, k) If supplied, the distances to the neighbors will be written here in-place. (default None)

resourcesOptional cuVS Resource handle for reusing CUDA resources.
If 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 cupy as cp
>>> from cuvs.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build the index
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
>>>
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> search_params = ivf_pq.SearchParams(n_probes=20)
>>>
>>> distances, neighbors = ivf_pq.search(search_params, index, queries,
...                                     k)