Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IndexIVFPQ.h
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11 // -*- c++ -*-
12
13 #ifndef FAISS_INDEX_IVFPQ_H
14 #define FAISS_INDEX_IVFPQ_H
15
16
17 #include <vector>
18
19 #include "IndexIVF.h"
20 #include "IndexPQ.h"
21
22
23 namespace faiss {
24
25
26
27 /** Inverted file with Product Quantizer encoding. Each residual
28 * vector is encoded as a product quantizer code.
29 */
30 struct IndexIVFPQ: IndexIVF {
31 bool by_residual; ///< Encode residual or plain vector?
32 int use_precomputed_table; ///< if by_residual, build precompute tables
33 size_t code_size; ///< code size per vector in bytes
34 ProductQuantizer pq; ///< produces the codes
35
36 bool do_polysemous_training; ///< reorder PQ centroids after training?
37 PolysemousTraining *polysemous_training; ///< if NULL, use default
38
39 // search-time parameters
40 size_t scan_table_threshold; ///< use table computation or on-the-fly?
41 size_t max_codes; ///< max nb of codes to visit to do a query
42 int polysemous_ht; ///< Hamming thresh for polysemous filtering
43
44 std::vector < std::vector<uint8_t> > codes; // binary codes, size nlist
45
46 /// if use_precompute_table
47 /// size nlist * pq.M * pq.ksub
48 std::vector <float> precomputed_table;
49
50 IndexIVFPQ (
51 Index * quantizer, size_t d, size_t nlist,
52 size_t M, size_t nbits_per_idx);
53
54 virtual void set_typename () override;
55
56 virtual void add_with_ids (
57 idx_t n, const float *x,
58 const long *xids = nullptr) override;
59
60 /// same as add_core, also:
61 /// - output 2nd level residuals if residuals_2 != NULL
62 /// - use precomputed list numbers if precomputed_idx != NULL
63 void add_core_o (idx_t n, const float *x,
64 const long *xids, float *residuals_2,
65 const long *precomputed_idx = nullptr);
66
67 virtual void search (
68 idx_t n, const float *x, idx_t k,
69 float *distances, idx_t *labels) const override;
70
71 virtual void reset () override;
72
73 virtual long remove_ids (const IDSelector & sel) override;
74
75 /// trains the product quantizer
76 virtual void train_residual(idx_t n, const float *x) override;
77
78 /// same as train_residual, also output 2nd level residuals
79 void train_residual_o (idx_t n, const float *x, float *residuals_2);
80
81
82 /** Reconstruct a subset of the indexed vectors
83 *
84 * @param i0 first vector to reconstruct
85 * @param ni nb of vectors to reconstruct
86 * @param recons output array of reconstructed vectors, size ni * d
87 */
88 virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons)
89 const override;
90
91 virtual void reconstruct (idx_t key, float * recons)
92 const override;
93
94 /** Find exact duplicates in the dataset.
95 *
96 * the duplicates are returned in pre-allocated arrays (see the
97 * max sizes).
98 *
99 * @params lims limits between groups of duplicates
100 * (max size ntotal / 2 + 1)
101 * @params ids ids[lims[i]] : ids[lims[i+1]-1] is a group of
102 * duplicates (max size ntotal)
103 * @return n number of groups found
104 */
105 size_t find_duplicates (idx_t *ids, size_t *lims) const;
106
107 // map a vector to a binary code knowning the index
108 void encode (long key, const float * x, uint8_t * code) const;
109
110 /// same as encode, for multiple points at once
111 void encode_multiple (size_t n, const long *keys,
112 const float * x, uint8_t * codes) const;
113
114 /** search a set of vectors, that are pre-quantized by the IVF
115 * quantizer. Fill in the corresponding heaps with the query
116 * results.
117 *
118 * @param nx nb of vectors to query
119 * @param qx query vectors, size nx * d
120 * @param keys coarse quantization indices, size nx * nprobe
121 * @param coarse_dis
122 * distances to coarse centroids, size nx * nprobe
123 * @param res heaps for all the results, gives the nprobe
124 * @param store_pairs store inv list index + inv list offset
125 * instead in upper/lower 32 bit of result,
126 * instead of ids (used for reranking).
127 */
128 virtual void search_knn_with_key (
129 size_t nx,
130 const float * qx,
131 const long * keys,
132 const float * coarse_dis,
133 float_maxheap_array_t* res,
134 bool store_pairs = false) const;
135
136 /// build precomputed table
137 void precompute_table ();
138
139 /// used to implement merging
140 virtual void merge_from_residuals (IndexIVF &other) override;
141
142
143 /** copy a subset of the entries index to the other index
144 *
145 * if subset_type == 0: copies ids in [a1, a2)
146 * if subset_type == 1: copies ids if id % a1 == a2
147 */
148 void copy_subset_to (IndexIVFPQ & other, int subset_type,
149 long a1, long a2) const;
150
151 IndexIVFPQ ();
152
153 };
154
155
156 /// statistics are robust to internal threading, but not if
157 /// IndexIVFPQ::search is called by multiple threads
158 struct IndexIVFPQStats {
159 size_t nq; // nb of queries run
160 size_t nlist; // nb of inverted lists scanned
161 size_t ncode; // nb of codes visited
162 size_t nrefine; // nb of refines (IVFPQR)
163
164 size_t n_hamming_pass;
165 // nb of passed Hamming distance tests (for polysemous)
166
167 // timings measured with the CPU RTC
168 // on all threads
169 size_t assign_cycles;
170 size_t search_cycles;
171 size_t refine_cycles; // only for IVFPQR
172
173 // single thread (double-counted with search_cycles)
174 size_t init_query_cycles;
175 size_t init_list_cycles;
176 size_t scan_cycles;
177 size_t heap_cycles;
178
179 IndexIVFPQStats () {reset (); }
180 void reset ();
181 };
182
183 // global var that collects them all
184 extern IndexIVFPQStats indexIVFPQ_stats;
185
186
187
188 /** Index with an additional level of PQ refinement */
189 struct IndexIVFPQR: IndexIVFPQ {
190 ProductQuantizer refine_pq; ///< 3rd level quantizer
191 std::vector <uint8_t> refine_codes; ///< corresponding codes
192
193 /// factor between k requested in search and the k requested from the IVFPQ
194 float k_factor;
195
196 IndexIVFPQR (
197 Index * quantizer, size_t d, size_t nlist,
198 size_t M, size_t nbits_per_idx,
199 size_t M_refine, size_t nbits_per_idx_refine);
200
201 virtual void set_typename () override;
202
203 virtual void reset() override;
204
205 virtual long remove_ids (const IDSelector & sel) override;
206
207 /// trains the two product quantizers
208 virtual void train_residual (idx_t n, const float *x) override;
209
210 virtual void add_with_ids (idx_t n, const float *x, const long *xids)
211 override;
212
213 /// same as add_with_ids, but optionally use the precomputed list ids
214 void add_core (idx_t n, const float *x, const long *xids,
215 const long *precomputed_idx = nullptr);
216
217
218 virtual void reconstruct_n (idx_t i0, idx_t ni, float *recons)
219 const override;
220
221 virtual void search (
222 idx_t n, const float *x, idx_t k,
223 float *distances, idx_t *labels) const override;
224
225 virtual void merge_from_residuals (IndexIVF &other) override;
226
227 IndexIVFPQR();
228 };
229
230
231 /** Index with 32-bit ids and flat tables. Must be constructed from an
232 * exisiting IndexIVFPQ. Cannot be copy-constructed/assigned. The
233 * actual data is stored in the compact_* tables, the ids and codes
234 * tables are not used. */
235 struct IndexIVFPQCompact: IndexIVFPQ {
236
237 explicit IndexIVFPQCompact (const IndexIVFPQ &other);
238
239 /// how were the compact tables allocated?
240 enum Alloc_type_t {
241 Alloc_type_none, ///< alloc from outside
242 Alloc_type_new, ///< was allocated with new
243 Alloc_type_mmap ///< was mmapped
244 };
245
246 Alloc_type_t alloc_type;
247
248 uint32_t *limits; ///< size nlist + 1
249 uint32_t *compact_ids; ///< size ntotal
250 uint8_t *compact_codes; ///< size ntotal * code_size
251
252 // file and buffer this was mmapped (will be unmapped when object
253 // is deleted)
254 char * mmap_buffer;
255 long mmap_length;
256
257 virtual void search_knn_with_key (
258 size_t nx,
259 const float * qx,
260 const long * keys,
261 const float * coarse_dis,
262 float_maxheap_array_t * res,
263 bool store_pairs = false) const override;
264
265 /// the three following functions will fail at runtime
266 virtual void add (idx_t, const float *) override;
267 virtual void reset () override;
268 virtual void train (idx_t, const float *) override;
269
270 virtual ~IndexIVFPQCompact ();
271
272 IndexIVFPQCompact ();
273
274 };
275
276
277
278 } // namespace faiss
279
280
281
282
283
284 #endif
faiss::IndexIVFPQCompact::compact_ids
uint32_t * compact_ids
size ntotal
Definition: IndexIVFPQ.h:249
faiss::IndexIVF
Definition: IndexIVF.h:46
faiss::IndexIVFPQCompact::compact_codes
uint8_t * compact_codes
size ntotal * code_size
Definition: IndexIVFPQ.h:250
faiss::IndexIVFPQ::precompute_table
void precompute_table()
build precomputed table
Definition: IndexIVFPQ.cpp:376
faiss::IndexIVFPQ::copy_subset_to
void copy_subset_to(IndexIVFPQ &other, int subset_type, long a1, long a2) const
Definition: IndexIVFPQ.cpp:317
faiss::IndexIVFPQ::reconstruct
virtual void reconstruct(idx_t key, float *recons) const override
Definition: IndexIVFPQ.cpp:287
faiss::IndexIVFPQR
Definition: IndexIVFPQ.h:189
faiss::IndexIVFPQR::refine_pq
ProductQuantizer refine_pq
3rd level quantizer
Definition: IndexIVFPQ.h:190
faiss::IndexIVFPQ::polysemous_training
PolysemousTraining * polysemous_training
if NULL, use default
Definition: IndexIVFPQ.h:37
faiss::IndexIVFPQCompact::add
virtual void add(idx_t, const float *) override
the three following functions will fail at runtime
Definition: IndexIVFPQ.cpp:1445
faiss::IndexIVFPQStats
Definition: IndexIVFPQ.h:158
faiss::IndexIVFPQCompact::search_knn_with_key
virtual void search_knn_with_key(size_t nx, const float *qx, const long *keys, const float *coarse_dis, float_maxheap_array_t *res, bool store_pairs=false) const override
Definition: IndexIVFPQ.cpp:1473
faiss::IndexIVFPQR::reconstruct_n
virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexIVFPQ.cpp:1336
faiss::IDSelector
Definition: AuxIndexStructures.h:58
faiss::IndexIVFPQCompact::Alloc_type_mmap
was mmapped
Definition: IndexIVFPQ.h:243
faiss::IndexIVFPQ::merge_from_residuals
virtual void merge_from_residuals(IndexIVF &other) override
used to implement merging
Definition: IndexIVFPQ.cpp:307
faiss::IndexIVFPQ::train_residual_o
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:83
faiss::IndexIVFPQ::do_polysemous_training
bool do_polysemous_training
reorder PQ centroids after training?
Definition: IndexIVFPQ.h:36
faiss::IndexIVFPQ::scan_table_threshold
size_t scan_table_threshold
use table computation or on-the-fly?
Definition: IndexIVFPQ.h:40
faiss::IndexIVFPQR::train_residual
virtual void train_residual(idx_t n, const float *x) override
trains the two product quantizers
Definition: IndexIVFPQ.cpp:1206
faiss::IndexIVFPQR::add_core
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
Definition: IndexIVFPQ.cpp:1230
faiss::IndexIVFPQCompact::limits
uint32_t * limits
size nlist + 1
Definition: IndexIVFPQ.h:248
faiss::IndexIVFPQ::precomputed_table
std::vector< float > precomputed_table
Definition: IndexIVFPQ.h:48
faiss::IndexIVFPQ::polysemous_ht
int polysemous_ht
Hamming thresh for polysemous filtering.
Definition: IndexIVFPQ.h:42
faiss::IndexIVFPQ::search_knn_with_key
virtual void search_knn_with_key(size_t nx, const float *qx, const long *keys, const float *coarse_dis, float_maxheap_array_t *res, bool store_pairs=false) const
Definition: IndexIVFPQ.cpp:946
faiss::IndexIVFPQR::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1197
faiss::IndexIVFPQ::add_with_ids
virtual void add_with_ids(idx_t n, const float *x, const long *xids=nullptr) override
Definition: IndexIVFPQ.cpp:170
faiss::IndexIVFPQCompact
Definition: IndexIVFPQ.h:235
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:56
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::IndexIVF::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:50
faiss::IndexIVFPQ::max_codes
size_t max_codes
max nb of codes to visit to do a query
Definition: IndexIVFPQ.h:41
faiss::IndexIVFPQCompact::Alloc_type_t
Alloc_type_t
how were the compact tables allocated?
Definition: IndexIVFPQ.h:240
faiss::IndexIVFPQR::refine_codes
std::vector< uint8_t > refine_codes
corresponding codes
Definition: IndexIVFPQ.h:191
faiss::IndexIVFPQ::remove_ids
virtual long remove_ids(const IDSelector &sel) override
Definition: IndexIVFPQ.cpp:1073
faiss::IndexIVFPQ::train_residual
virtual void train_residual(idx_t n, const float *x) override
trains the product quantizer
Definition: IndexIVFPQ.cpp:77
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:30
faiss::IndexIVFPQCompact::train
virtual void train(idx_t, const float *) override
Trains the quantizer and calls train_residual to train sub-quantizers.
Definition: IndexIVFPQ.cpp:1453
faiss::HeapArray
Definition: Heap.h:351
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::IndexIVFPQ::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1065
faiss::PolysemousTraining
optimizes the order of indices in a ProductQuantizer
Definition: PolysemousTraining.h:127
faiss::IndexIVFPQR::merge_from_residuals
virtual void merge_from_residuals(IndexIVF &other) override
used to implement merging
Definition: IndexIVFPQ.cpp:1353
faiss::IndexIVFPQ::by_residual
bool by_residual
Encode residual or plain vector?
Definition: IndexIVFPQ.h:31
faiss::IndexIVFPQCompact::Alloc_type_none
alloc from outside
Definition: IndexIVFPQ.h:241
faiss::IndexIVFPQ::pq
ProductQuantizer pq
produces the codes
Definition: IndexIVFPQ.h:34
faiss::IndexIVF::nlist
size_t nlist
number of possible key values
Definition: IndexIVF.h:47
faiss::Index
Definition: Index.h:61
faiss::IndexIVFPQ::reconstruct_n
virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexIVFPQ.cpp:258
faiss::IndexIVFPQ::add_core_o
void add_core_o(idx_t n, const float *x, const long *xids, float *residuals_2, const long *precomputed_idx=nullptr)
Definition: IndexIVFPQ.cpp:176
faiss::IndexIVFPQ::code_size
size_t code_size
code size per vector in bytes
Definition: IndexIVFPQ.h:33
faiss::IndexIVFPQR::remove_ids
virtual long remove_ids(const IDSelector &sel) override
Definition: IndexIVFPQ.cpp:1362
faiss::ProductQuantizer
Definition: ProductQuantizer.h:26
faiss::IndexIVFPQCompact::reset
virtual void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1449
faiss::IndexIVFPQCompact::Alloc_type_new
was allocated with new
Definition: IndexIVFPQ.h:242
faiss::IndexIVFPQ::encode_multiple
void encode_multiple(size_t n, const long *keys, const float *x, uint8_t *codes) const
same as encode, for multiple points at once
Definition: IndexIVFPQ.cpp:153
faiss::IndexIVFPQR::add_with_ids
virtual void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: IndexIVFPQ.cpp:1226
faiss::IndexIVFPQR::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexIVFPQ.cpp:1249
faiss::IndexIVFPQ::search
virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexIVFPQ.cpp:1045
faiss::IndexIVFPQ::find_duplicates
size_t find_duplicates(idx_t *ids, size_t *lims) const
Definition: IndexIVFPQ.cpp:1130
faiss::IndexIVFPQR::k_factor
float k_factor
factor between k requested in search and the k requested from the IVFPQ
Definition: IndexIVFPQ.h:194
faiss::IndexIVFPQ::use_precomputed_table
int use_precomputed_table
if by_residual, build precompute tables
Definition: IndexIVFPQ.h:32