Neighbors IVF SQ

Python module: cuvs.neighbors.ivf_sq

Index

1 cdef class Index

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

Members

Name	Kind
`trained`	property
`n_lists`	property
`dim`	property
`centers`	property

trained

1 def trained(self)

n_lists

1 def n_lists(self)

The number of inverted lists (clusters)

dim

1 def dim(self)

dimensionality of the cluster centers

centers

1 def centers(self)

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

IndexParams

1 cdef class IndexParams

Parameters to build index for IvfSq nearest neighbor search

Note: IVF-SQ currently uses fixed 8-bit residual scalar quantization. There are no additional SQ-specific tuning knobs.

Parameters

Name	Type	Description
`n_lists`	`int, default = 1024`	The number of clusters used in the coarse quantizer.
`metric`	`str, 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).
`metric_arg`	`float, default = 2.0`	Additional metric argument forwarded to cuVS distance computations.
`kmeans_n_iters`	`int, default = 20`	The number of iterations searching for kmeans centers during index building.
`max_train_points_per_cluster`	`int, default = 256`	The number of data vectors per cluster to use during iterative kmeans building. The index uses at most n_lists * max_train_points_per_cluster rows for training.
`add_data_on_build`	`bool, default = True`	After training the coarse clustering model and residual scalar quantization parameters, we 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_allocation`	`bool, default = False`	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 flag to `True`.

Constructor

1 def __init__(self, *, n_lists=1024, metric="sqeuclidean", metric_arg=2.0, kmeans_n_iters=20, max_train_points_per_cluster=256, add_data_on_build=True, conservative_memory_allocation=False)

Members

Name	Kind
`get_handle`	method
`metric`	property
`metric_arg`	property
`add_data_on_build`	property
`n_lists`	property
`kmeans_n_iters`	property
`max_train_points_per_cluster`	property
`conservative_memory_allocation`	property

get_handle

1 def get_handle(self)

metric

1 def metric(self)

metric_arg

1 def metric_arg(self)

add_data_on_build

1 def add_data_on_build(self)

n_lists

1 def n_lists(self)

kmeans_n_iters

1 def kmeans_n_iters(self)

max_train_points_per_cluster

1 def max_train_points_per_cluster(self)

conservative_memory_allocation

1 def conservative_memory_allocation(self)

SearchParams

1 cdef class SearchParams

Supplemental parameters to search IVF-SQ index

Parameters

Name	Type	Description
`n_probes`	`int`	The number of clusters to search.

Constructor

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

Members

Name	Kind
`get_handle`	method
`n_probes`	property

get_handle

1 def get_handle(self)

n_probes

1 def n_probes(self)

build

@auto_sync_resources

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

Build the IvfSq index from the dataset for efficient search.

IVF-SQ (Scalar Quantization) uses IVF partitioning together with per-dimension scalar quantization. Each vector’s residual is encoded as one byte per dimension, which can reduce vector-storage memory by about 4x vs IVF-Flat for float32 inputs (about 2x for float16 inputs), excluding IVF structural overhead. Recall and speed trade-offs versus IVF-PQ are dataset and tuning dependent.

Parameters

Name	Type	Description
`index_params`	`cuvs.neighbors.ivf_sq.IndexParams`
`dataset`	`CUDA array interface compliant matrix shape (n_samples, dim)`	Supported dtype [float32, float16]
`resources`	`cuvs.common.Resources, optional`

Returns

Name	Type	Description
`index`	`cuvs.neighbors.ivf_sq.Index`

Examples

1 >>> import cupy as cp
2 >>> from cuvs.neighbors import ivf_sq
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_sq.IndexParams(metric="sqeuclidean")
10 >>> index = ivf_sq.build(build_params, dataset)
11 >>> distances, neighbors = ivf_sq.search(ivf_sq.SearchParams(),
12 ...                                      index, dataset,
13 ...                                      k)
14 >>> distances = cp.asarray(distances)
15 >>> neighbors = cp.asarray(neighbors)

extend

@auto_sync_resources

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

Name	Type	Description
`index`	`ivf_sq.Index`	Trained ivf_sq object.
`new_vectors`	`array interface compliant matrix shape (n_samples, dim)`	Supported dtype [float32, float16]
`new_indices`	`array interface compliant vector shape (n_samples)`	Supported dtype [int64]
`resources`	`cuvs.common.Resources, optional`

Returns

Name	Type	Description
`index`	`cuvs.neighbors.ivf_sq.Index`

Examples

1 >>> import cupy as cp
2 >>> from cuvs.neighbors import ivf_sq
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_sq.build(ivf_sq.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_sq.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_sq.search(ivf_sq.SearchParams(),
18 ...                                      index, queries,
19 ...                                      k=10)

load

@auto_sync_resources

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

Name	Type	Description
`filename`	`string`	Name of the file.
`resources`	`cuvs.common.Resources, optional`

Returns

Name	Type	Description
`index`	`Index`

save

@auto_sync_resources

1 def save(filename, Index index, resources=None)

Saves the index to a file.

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

Parameters

Name	Type	Description
`filename`	`string`	Name of the file.
`index`	`Index`	Trained IVF-SQ index.
`resources`	`cuvs.common.Resources, optional`

Examples

1 >>> import cupy as cp
2 >>> from cuvs.neighbors import ivf_sq
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_sq.build(ivf_sq.IndexParams(), dataset)
9 >>> # Serialize and deserialize the ivf_sq index built
10 >>> ivf_sq.save("my_index.bin", index)
11 >>> index_loaded = ivf_sq.load("my_index.bin")

search

@auto_sync_resources @auto_convert_output

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

Name	Type	Description
`search_params`	`cuvs.neighbors.ivf_sq.SearchParams`
`index`	`cuvs.neighbors.ivf_sq.Index`	Trained IvfSq index.
`queries`	`CUDA array interface compliant matrix shape (n_samples, dim)`	Supported dtype [float32, float16]
`k`	`int`	The number of neighbors.
`neighbors`	`Optional CUDA array interface compliant matrix shape`	(n_queries, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None)
`distances`	`Optional 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)
`resources`	`cuvs.common.Resources, optional`

Examples

1 >>> import cupy as cp
2 >>> from cuvs.neighbors import ivf_sq
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_sq.build(ivf_sq.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_sq.SearchParams(n_probes=20)
16 >>>
17 >>> distances, neighbors = ivf_sq.search(search_params, index, queries,
18 ...                                     k)