IVF Flat

View as Markdown

Python module: cuvs.neighbors.ivf_flat

Index

1cdef class Index

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

Members

NameKind
trainedproperty
n_listsproperty
dimproperty
centersproperty

trained

1def trained(self)

n_lists

1def n_lists(self)

The number of inverted lists (clusters)

dim

1def dim(self)

dimensionality of the cluster centers

centers

1def centers(self)

Get the cluster centers corresponding to the lists in the original space

IndexParams

1cdef class IndexParams

Parameters to build index for IvfFlat nearest neighbor search

Parameters

NameTypeDescription
n_listsint, default = 1024The 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 = 20The 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.5If 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 = TrueAfter 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 = FalseBy 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.

Constructor

1def __init__(self, *, 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)

Members

NameKind
get_handlemethod
metricproperty
metric_argproperty
add_data_on_buildproperty
n_listsproperty
kmeans_n_itersproperty
kmeans_trainset_fractionproperty
adaptive_centersproperty
conservative_memory_allocationproperty

get_handle

1def get_handle(self)

metric

1def metric(self)

metric_arg

1def metric_arg(self)

add_data_on_build

1def add_data_on_build(self)

n_lists

1def n_lists(self)

kmeans_n_iters

1def kmeans_n_iters(self)

kmeans_trainset_fraction

1def kmeans_trainset_fraction(self)

adaptive_centers

1def adaptive_centers(self)

conservative_memory_allocation

1def conservative_memory_allocation(self)

SearchParams

1cdef class SearchParams

Supplemental parameters to search IVF-Flat index

Parameters

NameTypeDescription
n_probesintThe number of clusters to search.

Constructor

1def __init__(self, *, n_probes=20)

Members

NameKind
get_handlemethod
n_probesproperty

get_handle

1def get_handle(self)

n_probes

1def n_probes(self)

build

@auto_sync_resources

1def build(IndexParams index_params, dataset, resources=None)

Build the IvfFlat index from the dataset for efficient search.

Parameters

NameTypeDescription
index_paramscuvs.neighbors.ivf_flat.IndexParams
datasetCUDA array interface compliant matrix shape (n_samples, dim)Supported dtype [float32, float16, int8, uint8]
resourcescuvs.common.Resources, optional

Returns

NameTypeDescription
indexcuvs.neighbors.ivf_flat.Index

Examples

1>>> import cupy as cp
2>>> from cuvs.neighbors import ivf_flat
3>>> n_samples = 50000
4>>> n_features = 50
5>>> n_queries = 1000
6>>> k = 10
7>>> dataset = cp.random.random_sample((n_samples, n_features),
8... dtype=cp.float32)
9>>> build_params = ivf_flat.IndexParams(metric="sqeuclidean")
10>>> index = ivf_flat.build(build_params, dataset)
11>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
12... index, dataset,
13... k)
14>>> distances = cp.asarray(distances)
15>>> neighbors = cp.asarray(neighbors)

extend

@auto_sync_resources

1def 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.

Parameters

NameTypeDescription
indexivf_flat.IndexTrained ivf_flat object.
new_vectorsarray interface compliant matrix shape (n_samples, dim)Supported dtype [float32, float16, int8, uint8]
new_indicesarray interface compliant vector shape (n_samples)Supported dtype [int64]
resourcescuvs.common.Resources, optional

Returns

NameTypeDescription
indexcuvs.neighbors.ivf_flat.Index

Examples

1>>> import cupy as cp
2>>> from cuvs.neighbors import ivf_flat
3>>> n_samples = 50000
4>>> n_features = 50
5>>> n_queries = 1000
6>>> dataset = cp.random.random_sample((n_samples, n_features),
7... dtype=cp.float32)
8>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset)
9>>> n_rows = 100
10>>> more_data = cp.random.random_sample((n_rows, n_features),
11... dtype=cp.float32)
12>>> indices = n_samples + cp.arange(n_rows, dtype=cp.int64)
13>>> index = ivf_flat.extend(index, more_data, indices)
14>>> # Search using the built index
15>>> queries = cp.random.random_sample((n_queries, n_features),
16... dtype=cp.float32)
17>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
18... index, queries,
19... k=10)

load

@auto_sync_resources

1def 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.

Parameters

NameTypeDescription
filenamestringName of the file.
resourcescuvs.common.Resources, optional

Returns

NameTypeDescription
indexIndex

save

@auto_sync_resources

1def 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.

Parameters

NameTypeDescription
filenamestringName of the file.
indexIndexTrained IVF-Flat index.
resourcescuvs.common.Resources, optional

Examples

1>>> import cupy as cp
2>>> from cuvs.neighbors import ivf_flat
3>>> n_samples = 50000
4>>> n_features = 50
5>>> dataset = cp.random.random_sample((n_samples, n_features),
6... dtype=cp.float32)
7>>> # Build index
8>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset)
9>>> # Serialize and deserialize the ivf_flat index built
10>>> ivf_flat.save("my_index.bin", index)
11>>> index_loaded = ivf_flat.load("my_index.bin")

@auto_sync_resources @auto_convert_output

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

Find the k nearest neighbors for each query.

Parameters

NameTypeDescription
search_paramscuvs.neighbors.ivf_flat.SearchParams
indexcuvs.neighbors.ivf_flat.IndexTrained IvfFlat index.
queriesCUDA array interface compliant matrix shape (n_samples, dim)Supported dtype [float32, float16, int8, uint8]
kintThe 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)
filterOptional cuvs.neighbors.cuvsFilter can be used to filterneighbors based on a given bitset. (default None)
resourcescuvs.common.Resources, optional

Examples

1>>> import cupy as cp
2>>> from cuvs.neighbors import ivf_flat
3>>> n_samples = 50000
4>>> n_features = 50
5>>> n_queries = 1000
6>>> dataset = cp.random.random_sample((n_samples, n_features),
7... dtype=cp.float32)
8>>> # Build the index
9>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset)
10>>>
11>>> # Search using the built index
12>>> queries = cp.random.random_sample((n_queries, n_features),
13... dtype=cp.float32)
14>>> k = 10
15>>> search_params = ivf_flat.SearchParams(n_probes=20)
16>>>
17>>> distances, neighbors = ivf_flat.search(search_params, index, queries,
18... k)