cuvs.neighbors.cagra

cuvs.neighbors.cagra#

12 min read time

Applies to Linux

Submodules#

cuvs.neighbors.cagra.cagra

Classes#

`CompressionParams`	CompressionParams(pq_bits=8, *, pq_dim=0, vq_n_centers=0, kmeans_n_iters=25, vq_kmeans_trainset_fraction=0.0, pq_kmeans_trainset_fraction=0.0)
`ExtendParams`	ExtendParams(max_chunk_size=None, *)
`Index`	CAGRA index object. This object stores the trained CAGRA index state
`IndexParams`	IndexParams(metric=u'sqeuclidean', *, intermediate_graph_degree=128, graph_degree=64, build_algo=u'ivf_pq', nn_descent_niter=20, compression=None, ivf_pq_build_params: ivf_pq.IndexParams = None, ivf_pq_search_params: ivf_pq.SearchParams = None, double refinement_rate: float = 1.0)
`SearchParams`	SearchParams(max_queries=0, *, itopk_size=64, max_iterations=0, algo=u'auto', team_size=0, search_width=1, min_iterations=0, thread_block_size=0, hashmap_mode=u'auto', hashmap_min_bitlen=0, hashmap_max_fill_rate=0.5, num_random_samplings=1, rand_xor_mask=0x128394, persistent=False, persistent_lifetime=None, persistent_device_usage=None)

Functions#

`build`(*args[, resources])	build(IndexParams index_params, dataset, resources=None)
`extend`(*args[, resources])	extend(ExtendParams params, Index index, additional_dataset, resources=None)
`from_graph`(*args[, resources])	from_graph(graph, dataset, metric=u'sqeuclidean', 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, filter=None)

Package Contents#

class cuvs.neighbors.cagra.CompressionParams#

CompressionParams(pq_bits=8, *, pq_dim=0, vq_n_centers=0, kmeans_n_iters=25, vq_kmeans_trainset_fraction=0.0, pq_kmeans_trainset_fraction=0.0)

Parameters for VPQ Compression

pq_bits: int
The bit length of the vector element after compression by PQ. Possible values: [4, 5, 6, 7, 8]. The smaller the ‘pq_bits’, the smaller the index size and the better the search performance, but the lower the recall.

pq_dim: int
The dimensionality of the vector after compression by PQ. When zero, an optimal value is selected using a heuristic.

vq_n_centers: int
Vector Quantization (VQ) codebook size - number of “coarse cluster centers”. When zero, an optimal value is selected using a heuristic.

kmeans_n_iters: int
The number of iterations searching for kmeans centers (both VQ & PQ phases).

vq_kmeans_trainset_fraction: float
The fraction of data to use during iterative kmeans building (VQ phase). When zero, an optimal value is selected using a heuristic.

vq_kmeans_trainset_fraction: float
The fraction of data to use during iterative kmeans building (PQ phase). When zero, an optimal value is selected using a heuristic.

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

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

get_handle()#: CompressionParams.get_handle(self)

class cuvs.neighbors.cagra.ExtendParams#

ExtendParams(max_chunk_size=None, *)

Supplemental parameters to extend CAGRA Index

max_chunk_sizeint
The additional dataset is divided into chunks and added to the graph. This is the knob to adjust the tradeoff between the recall and operation throughput. Large chunk sizes can result in high throughput, but use more working memory (O(max_chunk_size*degree^2)). This can also degrade recall because no edges are added between the nodes in the same chunk. Auto select when 0.

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

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

class cuvs.neighbors.cagra.Index#

CAGRA index object. This object stores the trained CAGRA 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)

__len__()#: Return len(self).

class cuvs.neighbors.cagra.IndexParams#

IndexParams(metric=u’sqeuclidean’, *, intermediate_graph_degree=128, graph_degree=64, build_algo=u’ivf_pq’, nn_descent_niter=20, compression=None, ivf_pq_build_params: ivf_pq.IndexParams = None, ivf_pq_search_params: ivf_pq.SearchParams = None, double refinement_rate: float = 1.0)

Parameters to build index for CAGRA nearest neighbor search

metricstr, default = “sqeuclidean”
String denoting the metric type, valid values for metric are [“sqeuclidean”, “inner_product”], where:

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

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

intermediate_graph_degree : int, default = 128 graph_degree : int, default = 64 build_algo: str, default = “ivf_pq”

string denoting the graph building algorithm to use. Valid values for algo: [“ivf_pq”, “nn_descent”, “iterative_cagra_search”], where

ivf_pq will use the IVF-PQ algorithm for building the knn graph

nn_descent (experimental) will use the NN-Descent algorithm for building the knn graph. It is expected to be generally faster than ivf_pq.

iterative_cagra_search will iteratively build the knn graph using CAGRA’s search() and optimize()

compression: CompressionParams, optional
If compression is desired should be a CompressionParams object. If None compression will be disabled.

ivf_pq_build_params: cuvs.neighbors.ivf_pq.IndexParams, optional
Parameters for IVF-PQ algorithm. If provided, it will be used for building the graph.

ivf_pq_search_params: cuvs.neighbors.ivf_pq.SearchParams, optional
Parameters for IVF-PQ search. If provided, it will be used for searching the graph.

refinement_rate: float, default = 1.0

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.cagra.SearchParams#

SearchParams(max_queries=0, *, itopk_size=64, max_iterations=0, algo=u’auto’, team_size=0, search_width=1, min_iterations=0, thread_block_size=0, hashmap_mode=u’auto’, hashmap_min_bitlen=0, hashmap_max_fill_rate=0.5, num_random_samplings=1, rand_xor_mask=0x128394, persistent=False, persistent_lifetime=None, persistent_device_usage=None)

CAGRA search parameters

max_queries: int, default = 0
Maximum number of queries to search at the same time (batch size). Auto select when 0.

itopk_size: int, default = 64
Number of intermediate search results retained during the search. This is the main knob to adjust trade off between accuracy and search speed. Higher values improve the search accuracy.

max_iterations: int, default = 0
Upper limit of search iterations. Auto select when 0.

algo: str, default = “auto”
String denoting the search algorithm to use Valid values for algo: [“auto”, “single_cta”, “multi_cta”], where:

auto will automatically select the best value based on query size

single_cta is better when query contains larger number of vectors (e.g >10)

multi_cta is better when query contains only a few vectors

team_size: int, default = 0
Number of threads used to calculate a single distance. 4, 8, 16, or 32.

search_width: int, default = 1
Number of graph nodes to select as the starting point for the search in each iteration.

min_iterations: int, default = 0
Lower limit of search iterations.

thread_block_size: int, default = 0
Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0.

hashmap_mode: str, default = “auto”
String denoting the type of hash map to use. It’s usually better to allow the algorithm to select this value, Valid values for hashmap_mode: [“auto”, “small”, “hash”], where:

auto will automatically select the best value based on algo

small will use the small shared memory hash table with resetting.

hash will use a single hash table in global memory.

hashmap_min_bitlen: int, default = 0
Upper limit of hashmap fill rate. More than 0.1, less than 0.9.

hashmap_max_fill_rate: float, default = 0.5
Upper limit of hashmap fill rate. More than 0.1, less than 0.9.

num_random_samplings: int, default = 1
Number of iterations of initial random seed node selection. 1 or more.

rand_xor_mask: int, default = 0x128394
Bit mask used for initial random seed node selection.

persistent: bool, default = false
Whether to use the persistent version of the kernel

persistent_lifetime: float
Persistent kernel: time in seconds before the kernel stops if no requests are received.

persistent_device_usagefloat
Sets the fraction of maximum grid size used by persistent kernel.

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

static __repr__(*args, **kwargs)#

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

get_handle()#: SearchParams.get_handle(self)

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

build(IndexParams index_params, dataset, resources=None)

Build the CAGRA index from the dataset for efficient search.

The build performs two different steps- first an intermediate knn-graph is constructed, then it’s optimized it to create the final graph. The index_params object controls the node degree of these graphs.

It is required that both the dataset and the optimized graph fit the GPU memory.

The following distance metrics are supported:

L2

InnerProduct

index_params : IndexParams object dataset : CUDA array interface compliant matrix shape (n_samples, dim)

Supported dtype [float32, float16, 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.cagra.Index
>>> import cupy as cp
>>> from cuvs.neighbors import cagra
>>> 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 = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                      index, dataset,
...                                      k)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

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

extend(ExtendParams params, Index index, additional_dataset, resources=None)

Extend a CAGRA index with additional vectors

params : ExtendParams object index: Index

Existing cagra index to extend

additional_datasetCUDA array interface compliant matrix shape
Supported dtype [float32, float16, 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.

cuvs.neighbors.cagra.from_graph(*args, resources=None, **kwargs)#

from_graph(graph, dataset, metric=u’sqeuclidean’, resources=None)

Construct a cagra index from an existing graph and dataset

graphArray interface compliant matrix with shape (n_samples,
graph_degree)

datasetArray interface compliant matrix shape (n_samples, dim)
Supported dtype [float32, float16, int8, uint8]

metric : str resources : Optional 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.cagra.Index

cuvs.neighbors.cagra.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.cagra.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 CAGRA index.

include_datasetbool
Whether or not to write out the dataset along with the index. Including the dataset in the serialized index will use extra disk space, and might not be desired if you already have a copy of the dataset on disk. If this option is set to false, you will have to call index.update_dataset(dataset) after loading the 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 cagra
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> index = cagra.build(cagra.IndexParams(), dataset)
>>> # Serialize and deserialize the cagra index built
>>> cagra.save("my_index.bin", index)
>>> index_loaded = cagra.load("my_index.bin")

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

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

Find the k nearest neighbors for each query.

search_params : SearchParams index : Index

Trained CAGRA index.

queriesCUDA array interface compliant matrix shape (n_samples, dim)
Supported dtype [float32, 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)

filter: Optional cuvs.neighbors.cuvsFilter can be used to filter
neighbors based on a given bitset. (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 cagra
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> index = cagra.build(cagra.IndexParams(), dataset)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> search_params = cagra.SearchParams(
...     max_queries=100,
...     itopk_size=64
... )
>>> # Using a pooling allocator reduces overhead of temporary array
>>> # creation during search. This is useful if multiple searches
>>> # are performed with same query size.
>>> distances, neighbors = cagra.search(search_params, index, queries,
...                                     k)
>>> neighbors = cp.asarray(neighbors)
>>> distances = cp.asarray(distances)