cuvs.neighbors.hnsw#

9 min read time

Applies to Linux

Submodules#

Classes#

ExtendParams

ExtendParams(num_threads=0, *)

Index

HNSW index object. This object stores the trained HNSW index state

IndexParams

IndexParams(hierarchy=u'none', *, ef_construction=200, num_threads=0)

SearchParams

SearchParams(ef=200, *, num_threads=0)

Functions#

extend(*args[, resources])

extend(ExtendParams extend_params, Index index, data, resources=None)

from_cagra(*args[, resources])

from_cagra(IndexParams index_params, Index cagra_index, temporary_index_path=None, resources=None)

load(*args[, resources])

load(IndexParams index_params, filename, dim, dtype, metric=u'sqeuclidean', resources=None)

save(*args[, resources])

save(filename, Index index, resources=None)

search(*args[, resources])

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

Package Contents#

class cuvs.neighbors.hnsw.ExtendParams#

ExtendParams(num_threads=0, *)

Parameters to extend the HNSW index with new data

num_threadsint, default = 0 (optional)

Number of CPU threads used to increase construction parallelism. When set to 0, the number of threads is automatically determined.

static __reduce__(*args, **kwargs)#

ExtendParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#

ExtendParams.__setstate_cython__(self, __pyx_state)

class cuvs.neighbors.hnsw.Index#

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

class cuvs.neighbors.hnsw.IndexParams#

IndexParams(hierarchy=u’none’, *, ef_construction=200, num_threads=0)

Parameters to build index for HNSW nearest neighbor search

hierarchystring, default = “none” (optional)

The hierarchy of the HNSW index. Valid values are [“none”, “cpu”]. - “none”: No hierarchy is built. - “cpu”: Hierarchy is built using CPU. - “gpu”: Hierarchy is built using GPU.

ef_constructionint, default = 200 (optional)

Maximum number of candidate list size used during construction when hierarchy is cpu.

num_threadsint, default = 0 (optional)

Number of CPU threads used to increase construction parallelism when hierarchy is cpu or gpu. When the value is 0, the number of threads is automatically determined to the maximum number of threads available. NOTE: When hierarchy is gpu, while the majority of the work is done on the GPU, initialization of the HNSW index itself and some other work is parallelized with the help of CPU threads.

static __reduce__(*args, **kwargs)#

IndexParams.__reduce_cython__(self)

static __setstate__(*args, **kwargs)#

IndexParams.__setstate_cython__(self, __pyx_state)

class cuvs.neighbors.hnsw.SearchParams#

SearchParams(ef=200, *, num_threads=0)

HNSW search parameters

ef: int, default = 200

Maximum number of candidate list size used during search.

num_threads: int, default = 0

Number of CPU threads used to increase search parallelism. When set to 0, the number of threads is automatically determined using OpenMP’s omp_get_max_threads().

static __reduce__(*args, **kwargs)#

SearchParams.__reduce_cython__(self)

static __repr__(*args, **kwargs)#
static __setstate__(*args, **kwargs)#

SearchParams.__setstate_cython__(self, __pyx_state)

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

extend(ExtendParams extend_params, Index index, data, resources=None)

Extends the HNSW index with new data.

extend_params : ExtendParams index : Index

Trained HNSW index.

dataHost array interface compliant matrix shape (n_samples, dim)

Supported dtype [float32, 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.

>>> import numpy as np
>>> from cuvs.neighbors import hnsw, cagra
>>>
>>> n_samples = 50000
>>> n_features = 50
>>> dataset = np.random.random_sample((n_samples, n_features))
>>>
>>> # Build index
>>> index = cagra.build(hnsw.IndexParams(), dataset)
>>> # Load index
>>> hnsw_index = hnsw.from_cagra(hnsw.IndexParams(hierarchy="cpu"), index)
>>> # Extend the index with new data
>>> new_data = np.random.random_sample((n_samples, n_features))
>>> hnsw.extend(hnsw.ExtendParams(), hnsw_index, new_data)
cuvs.neighbors.hnsw.from_cagra(*args, resources=None, **kwargs)#

from_cagra(IndexParams index_params, Index cagra_index, temporary_index_path=None, resources=None)

Returns an HNSW index from a CAGRA index.

NOTE: When index_params.hierarchy is:
  1. NONE: 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. The returned index is immutable and can only be searched by the hnswlib wrapper in cuVS, as the format is not

    compatible with the original hnswlib.

  2. CPU: The returned index is mutable and can be extended with

    additional vectors. The serialized index is also compatible with the original hnswlib library.

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

index_paramsIndexParams

Parameters to convert the CAGRA index to HNSW index.

cagra_indexcagra.Index

Trained CAGRA index.

temporary_index_pathstring, default = None

Path to save the temporary index file. If None, the temporary file will be saved in /tmp/<random_number>.bin.

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
>>> from cuvs.neighbors import hnsw
>>> 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 the CAGRA index to hnswlib base layer only index format
>>> hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), index)
cuvs.neighbors.hnsw.load(*args, resources=None, **kwargs)#

load(IndexParams index_params, filename, dim, dtype, metric=u’sqeuclidean’, resources=None)

Loads an HNSW index. If the index was constructed with hnsw.IndexParams(hierarchy=”none”), then the loaded index is immutable and can only be searched by the hnswlib wrapper in cuVS, as the format is not compatible with the original hnswlib. However, if the index was constructed with hnsw.IndexParams(hierarchy=”cpu”), then the loaded index is mutable and compatible with the original hnswlib.

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.

index_paramsIndexParams

Parameters that were used to convert CAGRA index to HNSW index.

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.

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 : HnswIndex

>>> import cupy as cp
>>> from cuvs.neighbors import cagra
>>> from cuvs.neighbors import hnsw
>>> 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 the CAGRA index to hnswlib base layer only index format
>>> hnsw.save("my_index.bin", index)
>>> index = hnsw.load("my_index.bin", n_features, np.float32,
...                   "sqeuclidean")
cuvs.neighbors.hnsw.save(*args, resources=None, **kwargs)#

save(filename, Index index, resources=None)

Saves the CAGRA index to a file as an hnswlib index. If the index was constructed with hnsw.IndexParams(hierarchy=”none”), then the saved index is immutable and can only be searched by the hnswlib wrapper in cuVS, as the format is not compatible with the original hnswlib. However, if the index was constructed with hnsw.IndexParams(hierarchy=”cpu”), then the saved index is mutable and compatible with the original hnswlib.

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

filenamestring

Name of the file.

indexIndex

Trained HNSW 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
>>> cagra_index = cagra.build(cagra.IndexParams(), dataset)
>>> # Serialize and deserialize the cagra index built
>>> hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), cagra_index)
>>> hnsw.save("my_index.bin", hnsw_index)
cuvs.neighbors.hnsw.search(*args, resources=None, **kwargs)#

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

Find the k nearest neighbors for each query.

search_params : SearchParams index : Index

Trained HNSW index.

queriesCPU array interface compliant matrix shape (n_samples, dim)

Supported dtype [float, int]

kint

The number of neighbors.

neighborsOptional CPU array interface compliant matrix shape

(n_queries, k), dtype uint64_t. If supplied, neighbor indices will be written here in-place. (default None)

distancesOptional CPU 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.

>>> import cupy as cp
>>> from cuvs.neighbors import cagra, hnsw
>>> 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 = hnsw.SearchParams(
...     ef=200,
...     num_threads=0
... )
>>> # Convert CAGRA index to HNSW
>>> hnsw_index = hnsw.from_cagra(hnsw.IndexParams(), index)
>>> # 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 = hnsw.search(search_params, index, queries,
...                                     k)
>>> neighbors = cp.asarray(neighbors)
>>> distances = cp.asarray(distances)