File ProductQuantizer.h

namespace faiss

Implementation of k-means clustering with many variants.

Copyright (c) Facebook, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

IDSelector is intended to define a subset of vectors to handle (for removal or as subset to search)

PQ4 SIMD packing and accumulation functions

The basic kernel accumulates nq query vectors with bbs = nb * 2 * 16 vectors and produces an output matrix for that. It is interesting for nq * nb <= 4, otherwise register spilling becomes too large.

The implementation of these functions is spread over 3 cpp files to reduce parallel compile times. Templates are instantiated explicitly.

This file contains callbacks for kernels that compute distances.

Throughout the library, vectors are provided as float * pointers. Most algorithms can be optimized when several vectors are processed (added/searched) together in a batch. In this case, they are passed in as a matrix. When n vectors of size d are provided as float * x, component j of vector i is

x[ i * d + j ]

where 0 <= i < n and 0 <= j < d. In other words, matrices are always compact. When specifying the size of the matrix, we call it an n*d matrix, which implies a row-major storage.

I/O functions can read/write to a filename, a file handle or to an object that abstracts the medium.

The read functions return objects that should be deallocated with delete. All references within these objectes are owned by the object.

Definition of inverted lists + a few common classes that implement the interface.

Since IVF (inverted file) indexes are of so much use for large-scale use cases, we group a few functions related to them in this small library. Most functions work both on IndexIVFs and IndexIVFs embedded within an IndexPreTransform.

In this file are the implementations of extra metrics beyond L2 and inner product

Implements a few neural net layers, mainly to support QINCo

Defines a few objects that apply transformations to a set of vectors Often these are pre-processing steps.

Variables

FAISS_API int product_quantizer_compute_codes_bs
struct ProductQuantizer : public faiss::Quantizer
#include <ProductQuantizer.h>

Product Quantizer. PQ is trained using k-means, minimizing the L2 distance to centroids. PQ supports L2 and Inner Product search, however the quantization error is biased towards L2 distance.

Public Types

enum train_type_t

initialization

Values:

enumerator Train_default
enumerator Train_hot_start

the centroids are already initialized

enumerator Train_shared

share dictionary across PQ segments

enumerator Train_hypercube

initialize centroids with nbits-D hypercube

enumerator Train_hypercube_pca

initialize centroids with nbits-D hypercube

Public Functions

inline float *get_centroids(size_t m, size_t i)

return the centroids associated with subvector m

inline const float *get_centroids(size_t m, size_t i) const
virtual void train(size_t n, const float *x) override

Train the quantizer

Parameters:

x – training vectors, size n * d

ProductQuantizer(size_t d, size_t M, size_t nbits)
ProductQuantizer()
void set_derived_values()

compute derived values when d, M and nbits have been set

void set_params(const float *centroids, int m)

Define the centroids for subquantizer m.

void compute_code(const float *x, uint8_t *code) const

Quantize one vector with the product quantizer.

virtual void compute_codes(const float *x, uint8_t *codes, size_t n) const override

same as compute_code for several vectors

void compute_codes_with_assign_index(const float *x, uint8_t *codes, size_t n)

speed up code assignment using assign_index (non-const because the index is changed)

void decode(const uint8_t *code, float *x) const

decode a vector from a given code (or n vectors if third argument)

virtual void decode(const uint8_t *code, float *x, size_t n) const override

Decode a set of vectors

Parameters:
  • codes – input codes, size n * code_size

  • x – output vectors, size n * d

void compute_code_from_distance_table(const float *tab, uint8_t *code) const

If we happen to have the distance tables precomputed, this is more efficient to compute the codes.

void compute_distance_table(const float *x, float *dis_table) const

Compute distance table for one vector.

The distance table for x = [x_0 x_1 .. x_(M-1)] is a M * ksub matrix that contains

dis_table (m, j) = || x_m - c_(m, j)||^2 for m = 0..M-1 and j = 0 .. ksub - 1

where c_(m, j) is the centroid no j of sub-quantizer m.

Parameters:
  • x – input vector size d

  • dis_table – output table, size M * ksub

void compute_inner_prod_table(const float *x, float *dis_table) const
void compute_distance_tables(size_t nx, const float *x, float *dis_tables) const

compute distance table for several vectors

Parameters:
  • nx – nb of input vectors

  • x – input vector size nx * d

  • dis_table – output table, size nx * M * ksub

void compute_inner_prod_tables(size_t nx, const float *x, float *dis_tables) const
void search(const float *x, size_t nx, const uint8_t *codes, const size_t ncodes, float_maxheap_array_t *res, bool init_finalize_heap = true) const

perform a search (L2 distance)

Parameters:
  • x – query vectors, size nx * d

  • nx – nb of queries

  • codes – database codes, size ncodes * code_size

  • ncodes – nb of nb vectors

  • res – heap array to store results (nh == nx)

  • init_finalize_heap – initialize heap (input) and sort (output)?

void search_ip(const float *x, size_t nx, const uint8_t *codes, const size_t ncodes, float_minheap_array_t *res, bool init_finalize_heap = true) const

same as search, but with inner product similarity

void compute_sdc_table()
void search_sdc(const uint8_t *qcodes, size_t nq, const uint8_t *bcodes, const size_t ncodes, float_maxheap_array_t *res, bool init_finalize_heap = true) const
void sync_transposed_centroids()

Sync transposed centroids with regular centroids. This call is needed if centroids were edited directly.

void clear_transposed_centroids()

Clear transposed centroids table so ones are no longer used.

Public Members

size_t M

number of subquantizers

size_t nbits

number of bits per quantization index

size_t dsub

dimensionality of each subvector

size_t ksub

number of centroids for each subquantizer

bool verbose

verbose during training?

train_type_t train_type
ClusteringParameters cp

parameters used during clustering

Index *assign_index

if non-NULL, use this index for assignment (should be of size d / M)

std::vector<float> centroids

Centroid table, size M * ksub * dsub. Layout: (M, ksub, dsub)

std::vector<float> transposed_centroids

Transposed centroid table, size M * ksub * dsub. Layout: (dsub, M, ksub)

std::vector<float> centroids_sq_lengths

Squared lengths of centroids, size M * ksub Layout: (M, ksub)

std::vector<float> sdc_table

Symmetric Distance Table.

struct PQEncoderGeneric

Public Functions

inline PQEncoderGeneric(uint8_t *code, int nbits, uint8_t offset = 0)
inline void encode(uint64_t x)
inline ~PQEncoderGeneric()

Public Members

uint8_t *code

code for this vector

uint8_t offset
const int nbits

number of bits per subquantizer index

uint8_t reg
struct PQEncoder8

Public Functions

inline PQEncoder8(uint8_t *code, int nbits)
inline void encode(uint64_t x)

Public Members

uint8_t *code
struct PQEncoder16

Public Functions

inline PQEncoder16(uint8_t *code, int nbits)
inline void encode(uint64_t x)

Public Members

uint16_t *code
struct PQDecoderGeneric

Public Functions

inline PQDecoderGeneric(const uint8_t *code, int nbits)
inline uint64_t decode()

Public Members

const uint8_t *code
uint8_t offset
const int nbits
const uint64_t mask
uint8_t reg
struct PQDecoder8

Public Functions

inline PQDecoder8(const uint8_t *code, int nbits)
inline uint64_t decode()

Public Members

const uint8_t *code

Public Static Attributes

static const int nbits = 8
struct PQDecoder16

Public Functions

inline PQDecoder16(const uint8_t *code, int nbits)
inline uint64_t decode()

Public Members

const uint16_t *code

Public Static Attributes

static const int nbits = 16