IVF-Flat#

Index build parameters#

class cuvs.neighbors.ivf_flat.IndexParams(n_lists=1024, *, metric='sqeuclidean', metric_arg=2.0, kmeans_n_iters=20, kmeans_trainset_fraction=0.5, adaptive_centers=False, add_data_on_build=True, conservative_memory_allocation=False)#

Parameters to build index for IvfFlat nearest neighbor search

Parameters:

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”, “cosine”], 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.

cosine distance is defined as distance(a, b) = 1 - sum_i a_i * b_i / ( ||a||_2 * ||b||_2).

kmeans_n_itersint, default = 20

The number of iterations searching for kmeans centers during index building. The default setting is often fine, but this parameter can be decreased to improve training time wih larger trainset fractions (10M+ vectors) or increased for smaller trainset fractions (very small number of vectors) to improve recall.

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.

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.

adaptive_centersbool, default = False

By default (adaptive_centers = False), the cluster centers are trained in ivf_flat.build, and and never modified in ivf_flat.extend. The alternative behavior (adaptive_centers = true) is to update the cluster centers for new data when it is added. In this case, index.centers() are always exactly the centroids of the data in the corresponding clusters. The drawback of this behavior is that the centroids depend on the order of adding new data (through the classification of the added data); that is, index.centers() “drift” together with the changing distribution of the newly added data.

Attributes:

adaptive_centers
add_data_on_build
conservative_memory_allocation
kmeans_n_iters
kmeans_trainset_fraction
metric
metric_arg
n_lists

Index search parameters#

class cuvs.neighbors.ivf_flat.SearchParams(n_probes=20, *)#

Supplemental parameters to search IVF-Flat index

Parameters:

n_probes: int: The number of clusters to search.

Attributes:

n_probes

Index#

class cuvs.neighbors.ivf_flat.Index#

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

Attributes:

trained

Index build#

cuvs.neighbors.ivf_flat.build(IndexParams index_params, dataset, resources=None)[source]#

Build the IvfFlat index from the dataset for efficient search.

Parameters:

index_paramscuvs.neighbors.ivf_flat.IndexParams
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.

Returns:

index: py:class:cuvs.neighbors.ivf_flat.Index

Examples

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_flat
>>> 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_flat.IndexParams(metric="sqeuclidean")
>>> index = ivf_flat.build(build_params, dataset)
>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
...                                        index, dataset,
...                                        k)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

Index search#

cuvs.neighbors.ivf_flat.search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)[source]#

Find the k nearest neighbors for each query.

Parameters:

search_paramspy:class:cuvs.neighbors.ivf_flat.SearchParams
indexpy:class:cuvs.neighbors.ivf_flat.Index: Trained IvfFlat 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.

Examples

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_flat
>>> 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_flat.build(ivf_flat.IndexParams(), dataset)
>>>
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> search_params = ivf_flat.SearchParams(n_probes=20)
>>>
>>> distances, neighbors = ivf_flat.search(search_params, index, queries,
...                                     k)