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