Attention

The vector search and clustering algorithms in RAFT are being migrated to a new library dedicated to vector search called cuVS. We will continue to support the vector search algorithms in RAFT during this move, but will no longer update them after the RAPIDS 24.06 (June) release. We plan to complete the migration by RAPIDS 24.08 (August) release.

Neighbors#

This page provides pylibraft class references for the publicly-exposed elements of the neighbors package.

Brute Force#

pylibraft.neighbors.brute_force.knn(dataset, queries, k=None, indices=None, distances=None, metric='sqeuclidean', metric_arg=2.0, global_id_offset=0, handle=None)[source]#

Perform a brute-force nearest neighbors search.

Parameters:

datasetarray interface compliant matrix, row-major layout,: shape (n_samples, dim). Supported dtype [float]
queriesarray interface compliant matrix, row-major layout,: shape (n_queries, dim) Supported dtype [float]
kint: Number of neighbors to search (k <= 2048). Optional if indices or distances arrays are given (in which case their second dimension is k).
indicesOptional array interface compliant matrix shape: (n_queries, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None)

Supported dtype uint64
distancesOptional array interface compliant matrix shape: (n_queries, k), dtype float. If supplied, neighbor indices will be written here in-place. (default None)
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

indices: array interface compliant object containing resulting indices: shape (n_queries, k)
distances: array interface compliant object containing resulting distances: shape (n_queries, k)

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors.brute_force import knn
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 40
>>> distances, neighbors = knn(dataset, queries, k)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

CAGRA#

class pylibraft.neighbors.cagra.IndexParams(metric='sqeuclidean', *, intermediate_graph_degree=128, graph_degree=64, build_algo='ivf_pq')#

Parameters to build index for CAGRA nearest neighbor search

Parameters:

metricstring denoting the metric type, default=”sqeuclidean”

Valid values for metric: [“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 is the dot product between two vectors i.e.: distance(a, b) = sum_i (a_i * b_i)

intermediate_graph_degreeint, default = 128

graph_degreeint, default = 64

build_algo: string denoting the graph building algorithm to use, default = “ivf_pq”

Valid values for algo: [“ivf_pq”, “nn_descent”], 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.

Attributes:

graph_degree
intermediate_graph_degree
metric

pylibraft.neighbors.cagra.build(IndexParams index_params, dataset, handle=None)[source]#

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
inner_product

Parameters:

index_paramsIndexParams object
datasetCUDA array interface compliant matrix shape (n_samples, dim): Supported dtype [float, int8, uint8]
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

index: cagra.Index

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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)
>>> handle = DeviceResources()
>>> build_params = cagra.IndexParams(metric="sqeuclidean")
>>> index = cagra.build(build_params, dataset, handle=handle)
>>> distances, neighbors = cagra.search(cagra.SearchParams(),
...                                      index, dataset,
...                                      k, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

class pylibraft.neighbors.cagra.SearchParams(max_queries=0, *, itopk_size=64, max_iterations=0, algo='auto', team_size=0, search_width=1, min_iterations=0, thread_block_size=0, hashmap_mode='auto', hashmap_min_bitlen=0, hashmap_max_fill_rate=0.5, num_random_samplings=1, rand_xor_mask=0x128394)#

CAGRA search parameters

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: string denoting the search algorithm to use, default = “auto”

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: string denoting the type of hash map to use.

It’s usually better to allow the algorithm to select this value, default = “auto”. 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.

Attributes:

algo
hashmap_max_fill_rate
hashmap_min_bitlen
hashmap_mode
itopk_size
max_iterations
max_queries
min_iterations
num_random_samplings
rand_xor_mask
search_width
team_size
thread_block_size

pylibraft.neighbors.cagra.search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, handle=None)[source]#

Find the k nearest neighbors for each query.

Parameters:

search_paramsSearchParams
indexIndex: Trained CAGRA 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)
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>> # 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 performad with same query size.
>>> distances, neighbors = cagra.search(search_params, index, queries,
...                                     k, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()
>>> neighbors = cp.asarray(neighbors)
>>> distances = cp.asarray(distances)

Serializer Methods#

pylibraft.neighbors.cagra.save(filename, Index index, bool include_dataset=True, handle=None)[source]#

Saves the index to a file.

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

Parameters:

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.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import cagra
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>> # Serialize and deserialize the cagra index built
>>> cagra.save("my_index.bin", index, handle=handle)
>>> index_loaded = cagra.load("my_index.bin", handle=handle)

pylibraft.neighbors.cagra.load(filename, handle=None)[source]#

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 raft is not guaranteed to work.

Parameters:

filenamestring: Name of the file.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

indexIndex

HNSW#

class pylibraft.neighbors.hnsw.SearchParams(ef=200, num_threads=1)#

Hnswlib search parameters

Parameters:

ef: int, default=200: Size of list from which final neighbors k will be selected. ef should be greater than or equal to k.
num_threads: int, default=1: Number of host threads to use to search the hnswlib index and increase concurrency

Attributes:

ef
num_threads

pylibraft.neighbors.hnsw.from_cagra(Index index, handle=None)[source]#

Returns an hnswlib base-layer-only index from a CAGRA index.

NOTE: This method uses the filesystem to write the CAGRA index in: /tmp/<random_number>.bin before reading it as an hnswlib index, then deleting the temporary file.

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

Parameters:

indexIndex: Trained CAGRA index.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import cagra
>>> from pylibraft.neighbors import hnsw
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>> # Serialize the CAGRA index to hnswlib base layer only index format
>>> hnsw_index = hnsw.from_cagra(index, handle=handle)

pylibraft.neighbors.hnsw.search(SearchParams search_params, HnswIndex index, queries, k, neighbors=None, distances=None, handle=None)[source]#

Find the k nearest neighbors for each query.

Parameters:

search_paramsSearchParams
indexHnswIndex: Trained CAGRA index saved as base-layer-only hnswlib index.
queriesarray interface compliant matrix shape (n_samples, dim): Supported dtype [float, int8, uint8]
kint: The number of neighbors.
neighborsOptional array interface compliant matrix shape: (n_queries, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None)
distancesOptional array interface compliant matrix shape: (n_queries, k) If supplied, the distances to the neighbors will be written here in-place. (default None)
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> import numpy as np
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import cagra
>>> from pylibraft.neighbors import hnsw
>>> n_samples = 50000
>>> n_features = 50
>>> n_queries = 1000
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>>
>>> # Load saved base-layer-only hnswlib index from CAGRA index
>>> hnsw_index = hnsw.from_cagra(index, handle=handle)
>>>
>>> # Search hnswlib using the loaded index
>>> queries = np.random.random_sample((n_queries, n_features)).astype(
...     np.float32)
>>> k = 10
>>> search_params = hnsw.SearchParams(
...     ef=20,
...     num_threads=5
... )
>>> distances, neighbors = hnsw.search(search_params, hnsw_index,
...                                    queries, k, handle=handle)

Serializer Methods#

pylibraft.neighbors.hnsw.save(filename, Index index, handle=None)[source]#

Saves the CAGRA index as an hnswlib base-layer-only index to a file.

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

Parameters:

filenamestring: Name of the file.
indexIndex: Trained CAGRA index.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import cagra
>>> from pylibraft.neighbors import hnsw
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>> # Serialize the CAGRA index to hnswlib base layer only index format
>>> hnsw.save("my_index.bin", index, handle=handle)

pylibraft.neighbors.hnsw.load(filename, dim, dtype, metric='sqeuclidean', handle=None)[source]#

Loads base-layer-only hnswlib index from file, which was originally saved as a built CAGRA index.

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

Parameters:

filenamestring

Name of the file.

dimint

Dimensions of the training dataest

dtypenp.dtype of the saved index

Valid values for dtype: [np.float32, np.byte, np.ubyte]

metricstring denoting the metric type, default=”sqeuclidean”

Valid values for metric: [“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.

handleOptional RAFT resource handle for reusing CUDA resources.

If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

indexHnswIndex

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import cagra
>>> from pylibraft.neighbors import hnsw
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = cagra.build(cagra.IndexParams(), dataset, handle=handle)
>>> # Serialize the CAGRA index to hnswlib base layer only index format
>>> hnsw.save("my_index.bin", index, handle=handle)
>>> index = hnsw.load("my_index.bin", n_features, np.float32,
...                   "sqeuclidean")

IVF-Flat#

class pylibraft.neighbors.ivf_flat.IndexParams(n_lists=1024, *, metric=u'sqeuclidean', kmeans_n_iters=20, kmeans_trainset_fraction=0.5, add_data_on_build=True, bool adaptive_centers=False)#

Parameters to build index for IVF-FLAT nearest neighbor search

Parameters:

n_listint, default = 1024

The number of clusters used in the coarse quantizer.

metricstring denoting the metric type, default=”sqeuclidean”

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.

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
kmeans_n_iters
kmeans_trainset_fraction
metric
n_lists

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

Builds an IVF-FLAT index that can be used for nearest neighbor search.

Parameters:

index_paramsIndexParams object
datasetCUDA array interface compliant matrix shape (n_samples, dim): Supported dtype [float, int8, uint8]
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

index: ivf_flat.Index

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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)
>>> handle = DeviceResources()
>>> index_params = ivf_flat.IndexParams(
...     n_lists=1024,
...     metric="sqeuclidean")
>>> index = ivf_flat.build(index_params, dataset, handle=handle)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
...                                        index, queries, k,
...                                        handle=handle)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()

pylibraft.neighbors.ivf_flat.extend(Index index, new_vectors, new_indices, handle=None)[source]#

Extend an existing index with new vectors.

Parameters:

indexivf_flat.Index: Trained ivf_flat object.
new_vectorsCUDA array interface compliant matrix shape (n_samples, dim): Supported dtype [float, int8, uint8]
new_indicesCUDA array interface compliant vector shape (n_samples): Supported dtype [int64]
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

index: ivf_flat.Index

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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)
>>> handle = DeviceResources()
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset,
...                        handle=handle)
>>> n_rows = 100
>>> more_data = cp.random.random_sample((n_rows, n_features),
...                                     dtype=cp.float32)
>>> indices = index.size + cp.arange(n_rows, dtype=cp.int64)
>>> index = ivf_flat.extend(index, more_data, indices)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
...                                      index, queries,
...                                      k, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()

>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

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

IVF-FLAT search parameters

Parameters:

n_probes: int, default = 1024: The number of course clusters to select for the fine search.

Attributes:

n_probes

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

Find the k nearest neighbors for each query.

Parameters:

search_paramsSearchParams
indexIndex: Trained IVF-FLAT 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)
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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 index
>>> handle = DeviceResources()
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset,
...                        handle=handle)
>>> # 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, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()
>>> neighbors = cp.asarray(neighbors)
>>> distances = cp.asarray(distances)

Serializer Methods#

pylibraft.neighbors.ivf_flat.save(filename, Index index, handle=None)[source]#

Saves the index to a file.

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

Parameters:

filenamestring: Name of the file.
indexIndex: Trained IVF-Flat index.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import ivf_flat
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset,
...                        handle=handle)
>>> ivf_flat.save("my_index.bin", index, handle=handle)

pylibraft.neighbors.ivf_flat.load(filename, handle=None)[source]#

Loads index from a file.

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

Parameters:

filenamestring: Name of the file.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

indexIndex

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import ivf_flat
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build and save index
>>> handle = DeviceResources()
>>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset,
...                        handle=handle)
>>> ivf_flat.save("my_index.bin", index, handle=handle)
>>> del index
>>> n_queries = 100
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> handle = DeviceResources()
>>> index = ivf_flat.load("my_index.bin", handle=handle)
>>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(),
...                                        index, queries, k=10,
...                                        handle=handle)

IVF-PQ#

class pylibraft.neighbors.ivf_pq.IndexParams(n_lists=1024, *, metric='sqeuclidean', kmeans_n_iters=20, kmeans_trainset_fraction=0.5, pq_bits=8, pq_dim=0, codebook_kind='subspace', force_random_rotation=False, add_data_on_build=True, conservative_memory_allocation=False)#

Parameters to build index for IVF-PQ nearest neighbor search

Parameters:

n_listint, default = 1024

The number of clusters used in the coarse quantizer.

metricstring denoting the metric type, default=”sqeuclidean”

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

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.

Attributes:

add_data_on_build
codebook_kind
conservative_memory_allocation
force_random_rotation
kmeans_n_iters
kmeans_trainset_fraction
metric
n_lists
pq_bits
pq_dim

pylibraft.neighbors.ivf_pq.build(IndexParams index_params, dataset, handle=None)[source]#

Builds an IVF-PQ index that can be later used for nearest neighbor search.

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

Parameters:

index_paramsIndexParams object
datasetarray interface compliant matrix shape (n_samples, dim): Supported dtype [float, int8, uint8]
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

index: ivf_pq.Index

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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)
>>> handle = DeviceResources()
>>> index_params = ivf_pq.IndexParams(
...     n_lists=1024,
...     metric="sqeuclidean",
...     pq_dim=10)
>>> index = ivf_pq.build(index_params, dataset, handle=handle)
>>> # Search using the built index
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> k = 10
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(), index,
...                                      queries, k, handle=handle)
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()

pylibraft.neighbors.ivf_pq.extend(Index index, new_vectors, new_indices, handle=None)[source]#

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:

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]
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

index: ivf_pq.Index

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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)
>>> handle = DeviceResources()
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset, handle=handle)
>>> n_rows = 100
>>> more_data = cp.random.random_sample((n_rows, n_features),
...                                     dtype=cp.float32)
>>> indices = index.size + 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)
>>> k = 10
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
...                                      index, queries,
...                                      k, handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()
>>> distances = cp.asarray(distances)
>>> neighbors = cp.asarray(neighbors)

class pylibraft.neighbors.ivf_pq.SearchParams(n_probes=20, *, lut_dtype=np.float32, internal_distance_dtype=np.float32)#

IVF-PQ search parameters

Parameters:

n_probes: int, default = 1024: The number of course clusters to select for the fine 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]

Attributes:

internal_distance_dtype
lut_dtype
n_probes

pylibraft.neighbors.ivf_pq.search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, DeviceMemoryResource memory_resource=None, handle=None)[source]#

Find the k nearest neighbors for each query.

Parameters:

search_paramsSearchParams
indexIndex: Trained IVF-PQ 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)
memory_resourceRMM DeviceMemoryResource object, optional: This can be used to explicitly manage the temporary memory allocation during search. Passing a pooling allocator can reduce memory allocation overhead. If not specified, then the memory resource from the raft handle is used.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.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 index
>>> handle = DeviceResources()
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset, handle=handle)
>>> # 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,
...     lut_dtype=cp.float16,
...     internal_distance_dtype=cp.float32
... )
>>> # Using a pooling allocator reduces overhead of temporary array
>>> # creation during search. This is useful if multiple searches
>>> # are performad with same query size.
>>> import rmm
>>> mr = rmm.mr.PoolMemoryResource(
...     rmm.mr.CudaMemoryResource(),
...     initial_pool_size=2**29,
...     maximum_pool_size=2**31
... )
>>> distances, neighbors = ivf_pq.search(search_params, index, queries,
...                                      k, memory_resource=mr,
...                                      handle=handle)
>>> # pylibraft functions are often asynchronous so the
>>> # handle needs to be explicitly synchronized
>>> handle.sync()
>>> neighbors = cp.asarray(neighbors)
>>> distances = cp.asarray(distances)

Serializer Methods#

pylibraft.neighbors.ivf_pq.save(filename, Index index, handle=None)[source]#

Saves the index to a file.

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

Parameters:

filenamestring: Name of the file.
indexIndex: Trained IVF-PQ index.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build index
>>> handle = DeviceResources()
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset, handle=handle)
>>> ivf_pq.save("my_index.bin", index, handle=handle)

pylibraft.neighbors.ivf_pq.load(filename, handle=None)[source]#

Loads index from a file.

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

Parameters:

filenamestring: Name of the file.
handleOptional RAFT resource handle for reusing CUDA resources.: If a handle isn’t supplied, CUDA resources will be allocated inside this function and synchronized before the function exits. If a handle is supplied, you will need to explicitly synchronize yourself by calling handle.sync() before accessing the output.

Returns:

indexIndex

Examples

>>> import cupy as cp
>>> from pylibraft.common import DeviceResources
>>> from pylibraft.neighbors import ivf_pq
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = cp.random.random_sample((n_samples, n_features),
...                                   dtype=cp.float32)
>>> # Build and save index
>>> handle = DeviceResources()
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset, handle=handle)
>>> ivf_pq.save("my_index.bin", index, handle=handle)
>>> del index
>>> n_queries = 100
>>> queries = cp.random.random_sample((n_queries, n_features),
...                                   dtype=cp.float32)
>>> handle = DeviceResources()
>>> index = ivf_pq.load("my_index.bin", handle=handle)
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(), index,
...                                      queries, k=10, handle=handle)

Neighbors#

Brute Force#

CAGRA#

Serializer Methods#

HNSW#

Serializer Methods#

IVF-Flat#

Serializer Methods#

IVF-PQ#

Serializer Methods#

Candidate Refinement#