Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/tmp/faiss/IndexIVFPQ.h
1 /**
2  * Copyright (c) 2015-present, Facebook, Inc.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD+Patents license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 // -*- c++ -*-
10 
11 #ifndef FAISS_INDEX_IVFPQ_H
12 #define FAISS_INDEX_IVFPQ_H
13 
14 
15 #include <vector>
16 
17 #include "IndexIVF.h"
18 #include "IndexPQ.h"
19 
20 
21 namespace faiss {
22 
24  size_t scan_table_threshold; ///< use table computation or on-the-fly?
25  int polysemous_ht; ///< Hamming thresh for polysemous filtering
27 };
28 
29 
30 
31 
32 /** Inverted file with Product Quantizer encoding. Each residual
33  * vector is encoded as a product quantizer code.
34  */
36  bool by_residual; ///< Encode residual or plain vector?
37 
38  ProductQuantizer pq; ///< produces the codes
39 
40  bool do_polysemous_training; ///< reorder PQ centroids after training?
41  PolysemousTraining *polysemous_training; ///< if NULL, use default
42 
43  // search-time parameters
44  size_t scan_table_threshold; ///< use table computation or on-the-fly?
45  int polysemous_ht; ///< Hamming thresh for polysemous filtering
46 
47  /** Precompute table that speed up query preprocessing at some
48  * memory cost
49  * =-1: force disable
50  * =0: decide heuristically (default: use tables only if they are
51  * < precomputed_tables_max_bytes)
52  * =1: tables that work for all quantizers (size 256 * nlist * M)
53  * =2: specific version for MultiIndexQuantizer (much more compact)
54  */
55  int use_precomputed_table; ///< if by_residual, build precompute tables
57 
58  /// if use_precompute_table
59  /// size nlist * pq.M * pq.ksub
60  std::vector <float> precomputed_table;
61 
62  IndexIVFPQ (
63  Index * quantizer, size_t d, size_t nlist,
64  size_t M, size_t nbits_per_idx);
65 
66  void add_with_ids(idx_t n, const float* x, const long* xids = nullptr)
67  override;
68 
69  void encode_vectors(idx_t n, const float* x,
70  const idx_t *list_nos,
71  uint8_t * codes) const override;
72 
73  /// same as add_core, also:
74  /// - output 2nd level residuals if residuals_2 != NULL
75  /// - use precomputed list numbers if precomputed_idx != NULL
76  void add_core_o (idx_t n, const float *x,
77  const long *xids, float *residuals_2,
78  const long *precomputed_idx = nullptr);
79 
80  /// trains the product quantizer
81  void train_residual(idx_t n, const float* x) override;
82 
83  /// same as train_residual, also output 2nd level residuals
84  void train_residual_o (idx_t n, const float *x, float *residuals_2);
85 
86  void reconstruct_from_offset (long list_no, long offset,
87  float* recons) const override;
88 
89  /** Find exact duplicates in the dataset.
90  *
91  * the duplicates are returned in pre-allocated arrays (see the
92  * max sizes).
93  *
94  * @params lims limits between groups of duplicates
95  * (max size ntotal / 2 + 1)
96  * @params ids ids[lims[i]] : ids[lims[i+1]-1] is a group of
97  * duplicates (max size ntotal)
98  * @return n number of groups found
99  */
100  size_t find_duplicates (idx_t *ids, size_t *lims) const;
101 
102  // map a vector to a binary code knowning the index
103  void encode (long key, const float * x, uint8_t * code) const;
104 
105  /** Encode multiple vectors
106  *
107  * @param n nb vectors to encode
108  * @param keys posting list ids for those vectors (size n)
109  * @param x vectors (size n * d)
110  * @param codes output codes (size n * code_size)
111  * @param compute_keys if false, assume keys are precomputed,
112  * otherwise compute them
113  */
114  void encode_multiple (size_t n, long *keys,
115  const float * x, uint8_t * codes,
116  bool compute_keys = false) const;
117 
118  /// inverse of encode_multiple
119  void decode_multiple (size_t n, const long *keys,
120  const uint8_t * xcodes, float * x) const;
121 
122  InvertedListScanner *get_InvertedListScanner (bool store_pairs)
123  const override;
124 
125  /// build precomputed table
126  void precompute_table ();
127 
128  IndexIVFPQ ();
129 
130 };
131 
132 
133 /// statistics are robust to internal threading, but not if
134 /// IndexIVFPQ::search_preassigned is called by multiple threads
136  size_t nrefine; // nb of refines (IVFPQR)
137 
138  size_t n_hamming_pass;
139  // nb of passed Hamming distance tests (for polysemous)
140 
141  // timings measured with the CPU RTC
142  // on all threads
143  size_t search_cycles;
144  size_t refine_cycles; // only for IVFPQR
145 
146  IndexIVFPQStats () {reset (); }
147  void reset ();
148 };
149 
150 // global var that collects them all
151 extern IndexIVFPQStats indexIVFPQ_stats;
152 
153 
154 
155 /** Index with an additional level of PQ refinement */
157  ProductQuantizer refine_pq; ///< 3rd level quantizer
158  std::vector <uint8_t> refine_codes; ///< corresponding codes
159 
160  /// factor between k requested in search and the k requested from the IVFPQ
161  float k_factor;
162 
163  IndexIVFPQR (
164  Index * quantizer, size_t d, size_t nlist,
165  size_t M, size_t nbits_per_idx,
166  size_t M_refine, size_t nbits_per_idx_refine);
167 
168  void reset() override;
169 
170  long remove_ids(const IDSelector& sel) override;
171 
172  /// trains the two product quantizers
173  void train_residual(idx_t n, const float* x) override;
174 
175  void add_with_ids(idx_t n, const float* x, const long* xids) override;
176 
177  /// same as add_with_ids, but optionally use the precomputed list ids
178  void add_core (idx_t n, const float *x, const long *xids,
179  const long *precomputed_idx = nullptr);
180 
181  void reconstruct_from_offset (long list_no, long offset,
182  float* recons) const override;
183 
184  void merge_from (IndexIVF &other, idx_t add_id) override;
185 
186 
187  void search_preassigned (idx_t n, const float *x, idx_t k,
188  const idx_t *assign,
189  const float *centroid_dis,
190  float *distances, idx_t *labels,
191  bool store_pairs,
192  const IVFSearchParameters *params=nullptr
193  ) const override;
194 
195  IndexIVFPQR();
196 };
197 
198 
199 
200 /** Same as an IndexIVFPQ without the inverted lists: codes are stored sequentially
201  *
202  * The class is mainly inteded to store encoded vectors that can be
203  * accessed randomly, the search function is not implemented.
204  */
206  /// first level quantizer
208 
209  /// second level quantizer is always a PQ
211 
212  /// Codes. Size ntotal * code_size.
213  std::vector<uint8_t> codes;
214 
215  /// size of the code for the first level (ceil(log8(q1.nlist)))
216  size_t code_size_1;
217 
218  /// size of the code for the second level
219  size_t code_size_2;
220 
221  /// code_size_1 + code_size_2
222  size_t code_size;
223 
224  Index2Layer (Index * quantizer, size_t nlist,
225  int M, MetricType metric = METRIC_L2);
226 
227  Index2Layer ();
228  ~Index2Layer ();
229 
230  void train(idx_t n, const float* x) override;
231 
232  void add(idx_t n, const float* x) override;
233 
234  /// not implemented
235  void search(
236  idx_t n,
237  const float* x,
238  idx_t k,
239  float* distances,
240  idx_t* labels) const override;
241 
242  void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;
243 
244  void reconstruct(idx_t key, float* recons) const override;
245 
246  void reset() override;
247 
248  /// transfer the flat codes to an IVFPQ index
249  void transfer_to_IVFPQ(IndexIVFPQ & other) const;
250 
251 };
252 
253 
254 } // namespace faiss
255 
256 
257 #endif
void precompute_table()
build precomputed table
Definition: IndexIVFPQ.cpp:364
int polysemous_ht
Hamming thresh for polysemous filtering.
Definition: IndexIVFPQ.h:25
void merge_from(IndexIVF &other, idx_t add_id) override
void transfer_to_IVFPQ(IndexIVFPQ &other) const
transfer the flat codes to an IVFPQ index
size_t code_size_2
size of the code for the second level
Definition: IndexIVFPQ.h:219
void reconstruct_from_offset(long list_no, long offset, float *recons) const override
ProductQuantizer refine_pq
3rd level quantizer
Definition: IndexIVFPQ.h:157
PolysemousTraining * polysemous_training
if NULL, use default
Definition: IndexIVFPQ.h:41
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
size_t code_size
code_size_1 + code_size_2
Definition: IndexIVFPQ.h:222
void reset() override
removes all elements from the database.
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
Definition: Index.cpp:35
void decode_multiple(size_t n, const long *keys, const uint8_t *xcodes, float *x) const
inverse of encode_multiple
Definition: IndexIVFPQ.cpp:160
void train_residual_o(idx_t n, const float *x, float *residuals_2)
same as train_residual, also output 2nd level residuals
Definition: IndexIVFPQ.cpp:67
bool do_polysemous_training
reorder PQ centroids after training?
Definition: IndexIVFPQ.h:40
size_t scan_table_threshold
use table computation or on-the-fly?
Definition: IndexIVFPQ.h:44
void train_residual(idx_t n, const float *x) override
trains the two product quantizers
void add_core(idx_t n, const float *x, const long *xids, const long *precomputed_idx=nullptr)
same as add_with_ids, but optionally use the precomputed list ids
std::vector< float > precomputed_table
Definition: IndexIVFPQ.h:60
Level1Quantizer q1
first level quantizer
Definition: IndexIVFPQ.h:207
void encode_vectors(idx_t n, const float *x, const idx_t *list_nos, uint8_t *codes) const override
Definition: IndexIVFPQ.cpp:207
int polysemous_ht
Hamming thresh for polysemous filtering.
Definition: IndexIVFPQ.h:45
InvertedListScanner * get_InvertedListScanner(bool store_pairs) const override
get a scanner for this index (store_pairs means ignore labels)
void reset() override
removes all elements from the database.
std::vector< uint8_t > codes
Codes. Size ntotal * code_size.
Definition: IndexIVFPQ.h:213
void add_with_ids(idx_t n, const float *x, const long *xids=nullptr) override
Definition: IndexIVFPQ.cpp:183
int d
vector dimension
Definition: Index.h:66
std::vector< uint8_t > refine_codes
corresponding codes
Definition: IndexIVFPQ.h:158
void train_residual(idx_t n, const float *x) override
trains the product quantizer
Definition: IndexIVFPQ.cpp:61
static size_t precomputed_table_max_bytes
2G by default, accommodates tables up to PQ32 w/ 65536 centroids
Definition: IndexIVFPQ.h:56
void reconstruct_from_offset(long list_no, long offset, float *recons) const override
Definition: IndexIVFPQ.cpp:311
void encode_multiple(size_t n, long *keys, const float *x, uint8_t *codes, bool compute_keys=false) const
Definition: IndexIVFPQ.cpp:150
long idx_t
all indices are this type
Definition: Index.h:64
size_t scan_table_threshold
use table computation or on-the-fly?
Definition: IndexIVFPQ.h:24
void train(idx_t n, const float *x) override
optimizes the order of indices in a ProductQuantizer
void add(idx_t n, const float *x) override
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
not implemented
bool by_residual
Encode residual or plain vector?
Definition: IndexIVFPQ.h:36
ProductQuantizer pq
produces the codes
Definition: IndexIVFPQ.h:38
size_t code_size_1
size of the code for the first level (ceil(log8(q1.nlist)))
Definition: IndexIVFPQ.h:216
void add_core_o(idx_t n, const float *x, const long *xids, float *residuals_2, const long *precomputed_idx=nullptr)
Definition: IndexIVFPQ.cpp:221
void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const IVFSearchParameters *params=nullptr) const override
long remove_ids(const IDSelector &sel) override
Dataset manipulation functions.
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:33
ProductQuantizer pq
second level quantizer is always a PQ
Definition: IndexIVFPQ.h:210
void add_with_ids(idx_t n, const float *x, const long *xids) override
void reconstruct(idx_t key, float *recons) const override
size_t nlist
number of possible key values
Definition: IndexIVF.h:34
size_t find_duplicates(idx_t *ids, size_t *lims) const
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:45
float k_factor
factor between k requested in search and the k requested from the IVFPQ
Definition: IndexIVFPQ.h:161
int use_precomputed_table
if by_residual, build precompute tables
Definition: IndexIVFPQ.h:55