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
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <climits>
#include <cstdio>
#include <memory>
#include <faiss/IVFlib.h>
#include <faiss/IndexAdditiveQuantizer.h>
#include <faiss/IndexIVFAdditiveQuantizer.h>
#include <faiss/MetricType.h>
#include <faiss/utils/distances.h>
#include <faiss/utils/hamming.h>
#include <faiss/utils/random.h>
#include <faiss/utils/utils.h>
/* This demo file shows how to:
* - use a DistanceComputer to compute distances with encoded vectors
* - in the context of an IVF, how to split an additive quantizer into an
* AdditiveCoarseQuantizer and a ResidualQuantizer, in two different ways, with
* and without storing the prefix.
*/
int main() {
/******************************************
* Generate a test dataset
******************************************/
using idx_t = faiss::idx_t;
size_t d = 128;
size_t nt = 10000;
size_t nb = 10000;
size_t nq = 100;
double t0 = faiss::getmillisecs();
auto tic = [t0]() {
printf("[%.3f s] ", (faiss::getmillisecs() - t0) / 1000);
};
tic();
printf("samping dataset of %zd dim vectors, Q %zd B %zd T %zd\n",
d,
nq,
nb,
nt);
std::vector<float> buf(d * (nq + nt + nb));
faiss::rand_smooth_vectors(nq + nt + nb, d, buf.data(), 1234);
const float* xt = buf.data();
const float* xb = buf.data() + nt * d;
const float* xq = buf.data() + (nt + nb) * d;
idx_t k = 10;
std::vector<idx_t> gt(k * nq);
std::vector<float> unused(k * nq);
tic();
printf("compute ground truth, k=%zd\n", k);
faiss::knn_L2sqr(xq, xb, d, nq, nb, k, unused.data(), gt.data());
// a function to compute the accuracy
auto accuracy = [&](const idx_t* I) {
idx_t accu = 0;
for (idx_t q = 0; q < nq; q++) {
accu += faiss::ranklist_intersection_size(
k, gt.data() + q * k, k, I + q * k);
}
return double(accu) / (k * nq);
};
/******************************************
* Prepare the residual quantizer
******************************************/
faiss::ResidualQuantizer rq(
d, 7, 6, faiss::AdditiveQuantizer::ST_norm_qint8);
// do cheap an inaccurate training
rq.cp.niter = 5;
rq.max_beam_size = 5;
rq.train_type = 0;
tic();
printf("training the residual quantizer beam_size=%d\n", rq.max_beam_size);
rq.train(nt, xt);
tic();
printf("encoding the database, code_size=%zd\n", rq.code_size);
size_t code_size = rq.code_size;
std::vector<uint8_t> raw_codes(nb * code_size);
rq.compute_codes(xb, raw_codes.data(), nb);
/****************************************************************
* Make an index that uses that residual quantizer
* Verify that a distance computer gives the same distances
****************************************************************/
{
faiss::IndexResidualQuantizer index(
rq.d, rq.nbits, faiss::METRIC_L2, rq.search_type);
// override trained index
index.rq = rq;
index.is_trained = true;
// override vectors
index.codes = faiss::MaybeOwnedVector<uint8_t>(raw_codes);
index.ntotal = nb;
tic();
printf("IndexResidualQuantizer ready, searching\n");
std::vector<float> D(k * nq);
std::vector<idx_t> I(k * nq);
index.search(nq, xq, k, D.data(), I.data());
tic();
printf("Accuracy (intersection @ %zd): %.3f\n", k, accuracy(I.data()));
std::unique_ptr<faiss::FlatCodesDistanceComputer> dc(
index.get_FlatCodesDistanceComputer());
float max_diff12 = 0, max_diff13 = 0;
for (idx_t q = 0; q < nq; q++) {
const float* query = xq + q * d;
dc->set_query(query);
for (int i = 0; i < k; i++) {
// 3 ways of computing the same distance
// distance returned by the index
float dis1 = D[q * k + i];
// distance returned by the DistanceComputer that accesses the
// index
idx_t db_index = I[q * k + i];
float dis2 = (*dc)(db_index);
// distance computer from a code that does not belong to the
// index
const uint8_t* code = raw_codes.data() + code_size * db_index;
float dis3 = dc->distance_to_code(code);
max_diff12 = std::max(std::abs(dis1 - dis2), max_diff12);
max_diff13 = std::max(std::abs(dis1 - dis3), max_diff13);
}
}
tic();
printf("Max DistanceComputer discrepancy 1-2: %g 1-3: %g\n",
max_diff12,
max_diff13);
}
/****************************************************************
* Make an IVF index that uses the first 2 levels as a coarse quantizer
* The IVF codes contain the full code (ie. redundant with the coarse
*quantizer code)
****************************************************************/
{
// build a coarse quantizer from the 2 first levels of the RQ
std::vector<size_t> nbits(2);
std::copy(rq.nbits.begin(), rq.nbits.begin() + 2, nbits.begin());
faiss::ResidualCoarseQuantizer rcq(rq.d, nbits);
// set the coarse quantizer from the 2 first quantizers
rcq.rq.initialize_from(rq);
rcq.is_trained = true;
rcq.ntotal = (idx_t)1 << rcq.rq.tot_bits;
// settings for exhaustive search in RCQ
rcq.centroid_norms.resize(rcq.ntotal);
rcq.aq->compute_centroid_norms(rcq.centroid_norms.data());
rcq.beam_factor = -1.0; // use exact search
size_t nlist = rcq.ntotal;
tic();
printf("RCQ nlist = %zd tot_bits=%zd\n", nlist, rcq.rq.tot_bits);
// build a IVFResidualQuantizer from that
faiss::IndexIVFResidualQuantizer index(
&rcq, rcq.d, nlist, rq.nbits, faiss::METRIC_L2, rq.search_type);
index.by_residual = false;
index.rq = rq;
index.is_trained = true;
// there are 3 ways of filling up the index...
for (std::string filled_with : {"add", "manual", "derived"}) {
tic();
printf("filling up the index with %s, code_size=%zd\n",
filled_with.c_str(),
index.code_size);
index.reset();
if (filled_with == "add") {
// standard add method
index.add(nb, xb);
} else if (filled_with == "manual") {
// compute inverted lists and add elements manually
// fill in the inverted index manually
faiss::InvertedLists& invlists = *index.invlists;
// assign vectors to inverted lists
std::vector<idx_t> listnos(nb);
std::vector<float> unused(nb);
rcq.search(nb, xb, 1, unused.data(), listnos.data());
// populate inverted lists
for (idx_t i = 0; i < nb; i++) {
invlists.add_entry(
listnos[i], i, &raw_codes[i * code_size]);
}
index.ntotal = nb;
} else if (filled_with == "derived") {
// Since we have the raw codes precomputed, their prefix is the
// inverted list index, so let's use that.
faiss::InvertedLists& invlists = *index.invlists;
// populate inverted lists
for (idx_t i = 0; i < nb; i++) {
const uint8_t* code = &raw_codes[i * code_size];
faiss::BitstringReader rd(code, code_size);
idx_t list_no =
rd.read(rcq.rq.tot_bits); // read the list number
invlists.add_entry(list_no, i, code);
}
index.ntotal = nb;
}
tic();
printf("Index filled in\n");
for (int nprobe : {1, 4, 16, 64, int(nlist)}) {
printf("setting nprobe=%-4d", nprobe);
index.nprobe = nprobe;
std::vector<float> D(k * nq);
std::vector<idx_t> I(k * nq);
index.search(nq, xq, k, D.data(), I.data());
tic();
printf("Accuracy (intersection @ %zd): %.3f\n",
k,
accuracy(I.data()));
}
}
}
/****************************************************************
* Make an IVF index that uses the first 2 levels as a coarse
* quantizer, but this time does not store the code prefix from the index
****************************************************************/
{
// build a coarse quantizer from the 2 first levels of the RQ
int nlevel = 2;
std::unique_ptr<faiss::IndexIVFResidualQuantizer> index(
faiss::ivflib::ivf_residual_from_quantizer(rq, nlevel));
// there are 2 ways of filling up the index...
for (std::string filled_with : {"add", "derived"}) {
tic();
printf("filling up the IVF index with %s, code_size=%zd\n",
filled_with.c_str(),
index->code_size);
index->reset();
if (filled_with == "add") {
// standard add method
index->add(nb, xb);
} else if (filled_with == "derived") {
faiss::ivflib::ivf_residual_add_from_flat_codes(
index.get(), nb, raw_codes.data(), rq.code_size);
}
tic();
printf("Index filled in\n");
for (int nprobe : {1, 4, 16, 64, int(index->nlist)}) {
printf("setting nprobe=%-4d", nprobe);
index->nprobe = nprobe;
std::vector<float> D(k * nq);
std::vector<idx_t> I(k * nq);
index->search(nq, xq, k, D.data(), I.data());
tic();
printf("Accuracy (intersection @ %zd): %.3f\n",
k,
accuracy(I.data()));
}
}
}
return 0;
}