Preprocessing Quantize PQ

Source header: cuvs/preprocessing/quantize/pq.hpp

Product Quantizer utilities

kmeans_params_variant

Alias for the variant holding either balanced or regular k-means parameters.

1 using kmeans_params_variant =
2 std::variant<cuvs::cluster::kmeans::balanced_params, cuvs::cluster::kmeans::params>;

preprocessing::quantize::pq::params

Product Quantizer parameters.

1 struct params { ... };

Fields

Name	Type	Description
`pq_bits`	`uint32_t`	The bit length of the vector element after compression by PQ. Possible value range: [4-16]. Hint: the smaller the ‘pq_bits’, the smaller the index size and the faster the fit/transform time, but the lower the recall.
`pq_dim`	`uint32_t`	The dimensionality of the vector after compression by PQ. When zero, dim / 4 is used as default. TODO: at the moment `dim` must be a multiple `pq_dim`.
`use_subspaces`	`bool`	Whether to use subspaces for product quantization (PQ). When true, one PQ codebook is used for each subspace. Otherwise, a single PQ codebook is used.
`use_vq`	`bool`	Whether to use Vector Quantization (KMeans) before product quantization (PQ). When true, VQ is used and PQ is trained on the residuals.
`vq_n_centers`	`uint32_t`	Vector Quantization (VQ) codebook size - number of “coarse cluster centers”. When zero, an optimal value is selected using a heuristic. (sqrt(n_rows))
`kmeans_params`	`kmeans_params_variant`	K-means parameters for PQ codebook training. Set to cuvs::cluster::kmeans::balanced_params for balanced k-means (default), or cuvs::cluster::kmeans::params for regular k-means. The active variant type selects the algorithm; balanced k-means tends to be faster for PQ training where cluster sizes are approximately equal. Only L2Expanded metric is supported. The number of clusters is always set to 1 << pq_bits.
`max_train_points_per_pq_code`	`uint32_t`	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. We will use `pq_n_centers * max_train_points_per_pq_code` training points to train each PQ codebook.
`max_train_points_per_vq_cluster`	`uint32_t`	The max number of data points to use per VQ cluster during training.

preprocessing::quantize::pq::params::params

Simplified constructor that will build an appropriate kmeans params object.

1 params(uint32_t pq_bits,
2 uint32_t pq_dim,
3 bool use_subspaces,
4 bool use_vq,
5 uint32_t vq_n_centers,
6 uint32_t kmeans_n_iters,
7 cuvs::cluster::kmeans::kmeans_type pq_kmeans_type =
8 cuvs::cluster::kmeans::kmeans_type::KMeansBalanced,
9 uint32_t max_train_points_per_pq_code    = 256,
10 uint32_t max_train_points_per_vq_cluster = 1024)
11 : pq_bits(pq_bits),

Parameters

Name	Type	Description
`pq_bits`	`uint32_t`
`pq_dim`	`uint32_t`
`use_subspaces`	`bool`
`use_vq`	`bool`
`vq_n_centers`	`uint32_t`
`kmeans_n_iters`	`uint32_t`
`pq_kmeans_type`	`cuvs::cluster::kmeans::kmeans_type`	Default: `cuvs::cluster::kmeans::kmeans_type::KMeansBalanced`.
`max_train_points_per_pq_code`	`uint32_t`	Default: `256`.
`max_train_points_per_vq_cluster`	`uint32_t`	Default: `1024`.

Returns

void

preprocessing::quantize::pq::quantizer

Defines and stores VPQ codebooks upon training

1 template <typename T>
2 struct quantizer { ... };

Fields

Name	Type	Description
`params_quantizer`	`params`	Parameters used to build this quantizer.
`vpq_codebooks`	`cuvs::neighbors::vpq_dataset<T, int64_t>`	VPQ codebooks produced during training.

preprocessing::quantize::pq::build

Initializes a product quantizer to be used later for quantizing the dataset.

1 quantizer<float> build(raft::resources const& res,
2 const params params,
3 raft::device_matrix_view<const float, int64_t> dataset);

The use of a pool memory resource is recommended for more consistent training performance.

Usage example:

Parameters

Name	Direction	Type	Description
`res`	in	`raft::resources const&`	raft resource
`params`	in	`const params`	configure product quantizer, e.g. quantile
`dataset`	in	`raft::device_matrix_view<const float, int64_t>`	a row-major matrix view on device or host

Returns

quantizer<float>

Additional overload: preprocessing::quantize::pq::build

1 quantizer<float> build(raft::resources const& res,
2 const params params,
3 raft::host_matrix_view<const float, int64_t> dataset);

Parameters

Name	Direction	Type	Description
`res`		`raft::resources const&`
`params`		`const params`
`dataset`		`raft::host_matrix_view<const float, int64_t>`

Returns

quantizer<float>

preprocessing::quantize::pq::transform

Applies quantization transform to given dataset

1 void transform(raft::resources const& res,
2 const quantizer<float>& quant,
3 raft::device_matrix_view<const float, int64_t> dataset,
4 raft::device_matrix_view<uint8_t, int64_t> codes_out,
5 std::optional<raft::device_vector_view<uint32_t, int64_t>> vq_labels = std::nullopt);

Usage example:

used, optional

Parameters

Name	Direction	Type	Description
`res`	in	`raft::resources const&`	raft resource
`quant`	in	`const quantizer<float>&`	a product quantizer
`dataset`	in	`raft::device_matrix_view<const float, int64_t>`	a row-major matrix view on device or host
`codes_out`	out	`raft::device_matrix_view<uint8_t, int64_t>`	a row-major matrix view on device containing the PQ codes
`vq_labels`	out	`std::optional<raft::device_vector_view<uint32_t, int64_t>>`	a vector view on device containing the VQ labels when VQ is Default: `std::nullopt`.

Returns

void

Additional overload: preprocessing::quantize::pq::transform

1 void transform(raft::resources const& res,
2 const quantizer<float>& quant,
3 raft::host_matrix_view<const float, int64_t> dataset,
4 raft::device_matrix_view<uint8_t, int64_t> codes_out,
5 std::optional<raft::device_vector_view<uint32_t, int64_t>> vq_labels = std::nullopt);

Parameters

Name	Type	Description
`res`	`raft::resources const&`
`quant`	`const quantizer<float>&`
`dataset`	`raft::host_matrix_view<const float, int64_t>`
`codes_out`	`raft::device_matrix_view<uint8_t, int64_t>`
`vq_labels`	`std::optional<raft::device_vector_view<uint32_t, int64_t>>`	Default: `std::nullopt`.

Returns

void

preprocessing::quantize::pq::get_quantized_dim

Get the dimension of the quantized dataset (in bytes)

1 inline int64_t get_quantized_dim(const params& config);

Parameters

Name	Direction	Type	Description
`config`	in	`const params&`	product quantizer parameters

Returns

inline int64_t

the dimension of the quantized dataset

preprocessing::quantize::pq::inverse_transform

Applies inverse quantization transform to given dataset

1 void inverse_transform(
2 raft::resources const& res,
3 const quantizer<float>& quant,
4 raft::device_matrix_view<const uint8_t, int64_t> pq_codes,
5 raft::device_matrix_view<float, int64_t> out,
6 std::optional<raft::device_vector_view<const uint32_t, int64_t>> vq_labels = std::nullopt);

Parameters

Name	Direction	Type	Description
`res`	in	`raft::resources const&`	raft resource
`quant`	in	`const quantizer<float>&`	a product quantizer
`pq_codes`	in	`raft::device_matrix_view<const uint8_t, int64_t>`	a row-major matrix view on device containing the PQ codes
`out`	out	`raft::device_matrix_view<float, int64_t>`	a row-major matrix view on device
`vq_labels`	in	`std::optional<raft::device_vector_view<const uint32_t, int64_t>>`	a vector view on device containing the VQ labels when VQ is used, optional Default: `std::nullopt`.

Returns

void