极速阅览

Python module: cuvs.neighbors.ivf_pq

Index

1	cdef class Index

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

Members

NameKind

trained	property
n_lists	property
dim	property
pq_dim	property
pq_len	property
pq_bits	property
centers	property
centers_padded	property
pq_centers	property
centers_rot	property
rotation_matrix	property
list_sizes	property
lists	method
list_data	method
list_indices	method

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

pq_dim

1	def pq_dim(self)

The dimensionality of an encoded vector after compression by PQ

pq_len

1	def pq_len(self)

The dimensionality of a subspace, i.e. the number of vector components mapped to a subspace

pq_bits

1	def pq_bits(self)

The bit length of an encoded vector element after compression by PQ.

centers

1	def centers(self)

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

centers_padded

1	def centers_padded(self)

Get the padded cluster centers [n_lists, dim_ext] where dim_ext = round_up(dim + 1, 8). This returns contiguous data suitable for build_precomputed.

pq_centers

1	def pq_centers(self)

Get the PQ cluster centers

centers_rot

1	def centers_rot(self)

Get the rotated cluster centers [n_lists, rot_dim] where rot_dim = pq_len * pq_dim

rotation_matrix

1	def rotation_matrix(self)

Get the rotation matrix [rot_dim, dim] Transform matrix (original space -> rotated padded space)

list_sizes

1	def list_sizes(self)

Get the sizes of each list

lists

1	def lists(self, resources=None)

Iterates through the pq-encoded list data

This function returns an iterator over each list, with each value being the pq-encoded data for the entire list

Parameters

NameTypeDescription

resources

cuvs.common.Resources, optional

list_data

def list_data(self, label, n_rows=0, offset=0, out_codes=None, resources=None)

Gets unpacked list data for a single list (cluster)

Parameters

NameTypeDescription

label, int		The cluster to get data for
n_rows, int		The number of rows to return for the cluster (0 is all rows)
offset, int		The row to start getting data at out_codes, CAI Optional buffer to hold memory. Will be created if None
resources	cuvs.common.Resources, optional

list_indices

1	def list_indices(self, label, n_rows=0)

Gets indices for a single cluster (list)

Parameters

NameTypeDescription

label, int

The cluster to get data for n_rows, int, optional Number of rows in the list

IndexParams

1	cdef class IndexParams

Parameters to build index for IvfPq nearest neighbor search

Parameters

NameTypeDescription

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).
kmeans_n_iters	int, default = 20	The number of iterations searching for kmeans centers during index building.
kmeans_trainset_fraction	int, 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_bits	int, default = 8	The bit length of the vector element after quantization.
pq_dim	int, 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_kind	string, default = "subspace"	Valid values [“subspace”, “cluster”]
force_random_rotation	bool, 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_build	bool, 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.
conservative_memory_allocation	bool, 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.
max_train_points_per_pq_code	int, default = 256	The max number of data points to use per PQ code during PQ codebook training. Using more data points per PQ code may increase the quality of PQ codebook but may also increase the build time. The parameter is applied to both PQ codebook generation methods, i.e., PER_SUBSPACE and PER_CLUSTER. In both cases, we will use pq_book_size * max_train_points_per_pq_code training points to train each codebook.
codes_layout	string, default = "interleaved"	Memory layout of the IVF-PQ list data. Valid values [“flat”, “interleaved”] - flat: Codes are stored contiguously, one vector’s codes after another. - interleaved: Codes are interleaved for optimized search performance. This is the default and recommended for search workloads.

Constructor

def __init__(self, *, n_lists=1024, metric="sqeuclidean", metric_arg=2.0, 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, max_train_points_per_pq_code=256, codes_layout="interleaved")

Members

NameKind

get_handle	method
metric	property
metric_arg	property
add_data_on_build	property
n_lists	property
kmeans_n_iters	property
kmeans_trainset_fraction	property
pq_bits	property
pq_dim	property
codebook_kind	property
force_random_rotation	property
add_data_on_build	property
conservative_memory_allocation	property
max_train_points_per_pq_code	property
codes_layout	property
get_handle	method

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)

kmeans_trainset_fraction

1	def kmeans_trainset_fraction(self)

pq_bits

1	def pq_bits(self)

pq_dim

1	def pq_dim(self)

codebook_kind

1	def codebook_kind(self)

force_random_rotation

1	def force_random_rotation(self)

add_data_on_build

1	def add_data_on_build(self)

conservative_memory_allocation

1	def conservative_memory_allocation(self)

max_train_points_per_pq_code

1	def max_train_points_per_pq_code(self)

codes_layout

1	def codes_layout(self)

get_handle

1	def get_handle(self)

SearchParams

1	cdef class SearchParams

Supplemental parameters to search IVF-Pq index

Parameters

NameTypeDescription

n_probes	int	The number of clusters to 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]
coarse_search_dtype	default = np.float32	[Experimental] The data type to use as the GEMM element type when searching the clusters to probe. Possible values: [np.float32, np.float16, np.int8]. - Legacy default: np.float32 - Recommended for performance: np.float16 (half) - Experimental/low-precision: np.int8
max_internal_batch_size	default = 4096	Set the internal batch size to improve GPU utilization at the cost of larger memory footprint.

Constructor

def __init__(self, *, n_probes=20, lut_dtype=np.float32, internal_distance_dtype=np.float32, coarse_search_dtype=np.float32, max_internal_batch_size=4096)

Members

NameKind

get_handle	method
n_probes	property
lut_dtype	property
internal_distance_dtype	property
coarse_search_dtype	property
max_internal_batch_size	property
get_handle	method

get_handle

1	def get_handle(self)

n_probes

1	def n_probes(self)

lut_dtype

1	def lut_dtype(self)

internal_distance_dtype

1	def internal_distance_dtype(self)

coarse_search_dtype

1	def coarse_search_dtype(self)

max_internal_batch_size

1	def max_internal_batch_size(self)

get_handle

1	def get_handle(self)

build

@auto_sync_resources

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

Build the IvfPq index from the dataset for efficient search.

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

Parameters

NameTypeDescription

index_params	cuvs.neighbors.ivf_pq.IndexParams	Parameters on how to build the index
dataset	Array interface compliant matrix shape (n_samples, dim)	Supported dtype [float32, float16, int8, uint8]
resources	cuvs.common.Resources, optional

Returns

NameTypeDescription

index

cuvs.neighbors.ivf_pq.Index

Examples

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

build_precomputed

@auto_sync_resources

def build_precomputed(IndexParams index_params, uint32_t dim, pq_centers, centers, centers_rot, rotation_matrix, resources=None)

Build a view-type IVF-PQ index from precomputed centroids and codebook.

This function creates a non-owning index that stores a reference to the provided device data. All parameters must be provided with correct extents. The caller is responsible for ensuring the lifetime of the input data exceeds the lifetime of the returned index.

The index_params must be consistent with the provided matrices. Specifically:

index_params.codebook_kind determines the expected shape of pq_centers
index_params.metric will be stored in the index
index_params.conservative_memory_allocation will be stored in the index

Parameters

NameTypeDescription

index_params	cuvs.neighbors.ivf_pq.IndexParams	Parameters that must be consistent with the provided matrices
dim	int	Dimensionality of the input data
pq_centers	CUDA array interface compliant tensor	PQ codebook on device memory with required shape: - codebook_kind “subspace”: [pq_dim, pq_len, pq_book_size] - codebook_kind “cluster”: [n_lists, pq_len, pq_book_size] Supported dtype: float32
centers	CUDA array interface compliant matrix	Cluster centers in the original space [n_lists, dim_ext] where dim_ext = round_up(dim + 1, 8). Supported dtype: float32
centers_rot	CUDA array interface compliant matrix	Rotated cluster centers [n_lists, rot_dim] where rot_dim = pq_len * pq_dim. Supported dtype: float32
rotation_matrix	CUDA array interface compliant matrix	Transform matrix (original space -> rotated padded space) [rot_dim, dim]. Supported dtype: float32
resources	cuvs.common.Resources, optional

Returns

NameTypeDescription

index

cuvs.neighbors.ivf_pq.Index

Examples

>>> import cupy as cp
>>> from cuvs.neighbors import ivf_pq
>>> n_lists = 100
>>> dim = 128
>>> pq_dim = 16
>>> pq_bits = 8
>>> pq_len = (dim + pq_dim - 1) // pq_dim  # ceil division
>>> pq_book_size = 1 << pq_bits
>>> rot_dim = pq_len * pq_dim
>>> dim_ext = ((dim + 1 + 7) // 8) * 8  # round_up(dim + 1, 8)
>>>
>>> # Prepare precomputed matrices (example with random data)
>>> pq_centers = cp.random.random((pq_dim, pq_len, pq_book_size),
...                               dtype=cp.float32)
>>> centers = cp.random.random((n_lists, dim_ext), dtype=cp.float32)
>>> centers_rot = cp.random.random((n_lists, rot_dim), dtype=cp.float32)
>>> rotation_matrix = cp.random.random((rot_dim, dim), dtype=cp.float32)
>>>
>>> # Build index from precomputed data
>>> build_params = ivf_pq.IndexParams(n_lists=n_lists, pq_dim=pq_dim,
...                                    pq_bits=pq_bits,
...                                    codebook_kind="subspace")
>>> index = ivf_pq.build_precomputed(build_params, dim, pq_centers,
...                                   centers, centers_rot, rotation_matrix)

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

NameTypeDescription

index	ivf_pq.Index	Trained ivf_pq object.
new_vectors	array interface compliant matrix shape (n_samples, dim)	Supported dtype [float, int8, uint8]
new_indices	array interface compliant vector shape (n_samples)	Supported dtype [int64]
resources	cuvs.common.Resources, optional

Returns

NameTypeDescription

index

cuvs.neighbors.ivf_pq.Index

Examples

>>> import cupy as cp
>>> from cuvs.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)
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
>>> n_rows = 100
>>> more_data = cp.random.random_sample((n_rows, n_features),
...                                     dtype=cp.float32)
>>> indices = n_samples + 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)
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
...                                      index, queries,
...                                      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

NameTypeDescription

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

Returns

NameTypeDescription

index

Index

save

@auto_sync_resources

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

filename	string	Name of the file.
index	Index	Trained IVF-PQ index.
resources	cuvs.common.Resources, optional

Examples

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

search

@auto_sync_resources @auto_convert_output

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

Find the k nearest neighbors for each query.

Parameters

NameTypeDescription

search_params	cuvs.neighbors.ivf_pq.SearchParams	Parameters on how to search the index
index	cuvs.neighbors.ivf_pq.Index	Trained IvfPq index.
queries	CUDA array interface compliant matrix shape (n_samples, dim)	Supported dtype [float, int8, uint8]
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)
resources	cuvs.common.Resources, optional

Examples

>>> import cupy as cp
>>> from cuvs.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 the index
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
>>>
>>> # 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)
>>>
>>> distances, neighbors = ivf_pq.search(search_params, index, queries,
...                                     k)

transform

@auto_sync_resources

def transform(Index index, input_dataset, output_labels=None, output_dataset=None, resources=None)

Transform a dataset by applying pq-encoding to the vectors.

Parameters

NameTypeDescription

index	ivf_pq.Index	Trained ivf_pq object.
input_dataset	array interface compliant matrix shape (n_samples, dim)	Supported dtype [float]
new_indices	Optional array interface compliant vector shape (n_samples)	Supported dtype [uint32]
output_dataset	Optional array interface compliant matrix shape (n_samples, pq_dim)	Supported dtype [uint8]
resources	cuvs.common.Resources, optional

Returns

NameTypeDescription

output_labels, output_dataset

The cluster that each point in the dataset belongs to, and the transformed dataset

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

1	>>> import cupy as cp
2	>>> from cuvs.neighbors import ivf_pq
3	>>> n_lists = 100
4	>>> dim = 128
5	>>> pq_dim = 16
6	>>> pq_bits = 8
7	>>> pq_len = (dim + pq_dim - 1) // pq_dim # ceil division
8	>>> pq_book_size = 1 << pq_bits
9	>>> rot_dim = pq_len * pq_dim
10	>>> dim_ext = ((dim + 1 + 7) // 8) * 8 # round_up(dim + 1, 8)
11	>>>
12	>>> # Prepare precomputed matrices (example with random data)
13	>>> pq_centers = cp.random.random((pq_dim, pq_len, pq_book_size),
14	... dtype=cp.float32)
15	>>> centers = cp.random.random((n_lists, dim_ext), dtype=cp.float32)
16	>>> centers_rot = cp.random.random((n_lists, rot_dim), dtype=cp.float32)
17	>>> rotation_matrix = cp.random.random((rot_dim, dim), dtype=cp.float32)
18	>>>
19	>>> # Build index from precomputed data
20	>>> build_params = ivf_pq.IndexParams(n_lists=n_lists, pq_dim=pq_dim,
21	... pq_bits=pq_bits,
22	... codebook_kind="subspace")
23	>>> index = ivf_pq.build_precomputed(build_params, dim, pq_centers,
24	... centers, centers_rot, rotation_matrix)

IVF PQ

Index

trained

n_lists

dim

pq_dim

pq_len

pq_bits

centers

centers_padded

pq_centers

centers_rot

rotation_matrix

list_sizes

lists

list_data

list_indices

IndexParams

get_handle

metric

metric_arg

add_data_on_build

n_lists

kmeans_n_iters

kmeans_trainset_fraction

pq_bits

pq_dim

codebook_kind

force_random_rotation

add_data_on_build

conservative_memory_allocation

max_train_points_per_pq_code

codes_layout

get_handle

SearchParams

get_handle

n_probes

lut_dtype

internal_distance_dtype

coarse_search_dtype

max_internal_batch_size

get_handle

build

build_precomputed

extend

load

save

search

transform

Neighbors IVF Flat

Neighbors IVF SQ