Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/tmp/faiss/IndexIVFPQ.cpp
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 #include "IndexIVFPQ.h"
12
13 #include <cmath>
14 #include <cstdio>
15 #include <cassert>
16
17 #include <algorithm>
18
19 #include "Heap.h"
20 #include "utils.h"
21
22 #include "Clustering.h"
23 #include "IndexFlat.h"
24
25 #include "hamming.h"
26
27 #include "FaissAssert.h"
28
29 #include "AuxIndexStructures.h"
30
31 namespace faiss {
32
33
34 /*****************************************
35 * IndexIVFPQ implementation
36 ******************************************/
37
38 IndexIVFPQ::IndexIVFPQ (Index * quantizer, size_t d, size_t nlist,
39 size_t M, size_t nbits_per_idx):
40 IndexIVF (quantizer, d, nlist, 0, METRIC_L2),
41 pq (d, M, nbits_per_idx)
42 {
43 FAISS_THROW_IF_NOT (nbits_per_idx <= 8);
44 code_size = pq.code_size;
45 invlists->code_size = code_size;
46 is_trained = false;
47 by_residual = true;
48 use_precomputed_table = 0;
49 scan_table_threshold = 0;
50
51 polysemous_training = nullptr;
52 do_polysemous_training = false;
53 polysemous_ht = 0;
54
55 }
56
57
58 /****************************************************************
59 * training */
60
61 void IndexIVFPQ::train_residual (idx_t n, const float *x)
62 {
63 train_residual_o (n, x, nullptr);
64 }
65
66
67 void IndexIVFPQ::train_residual_o (idx_t n, const float *x, float *residuals_2)
68 {
69 const float * x_in = x;
70
71 x = fvecs_maybe_subsample (
72 d, (size_t*)&n, pq.cp.max_points_per_centroid * pq.ksub,
73 x, verbose, pq.cp.seed);
74
75 ScopeDeleter<float> del_x (x_in == x ? nullptr : x);
76
77 const float *trainset;
78 ScopeDeleter<float> del_residuals;
79 if (by_residual) {
80 if(verbose) printf("computing residuals\n");
81 idx_t * assign = new idx_t [n]; // assignement to coarse centroids
82 ScopeDeleter<idx_t> del (assign);
83 quantizer->assign (n, x, assign);
84 float *residuals = new float [n * d];
85 del_residuals.set (residuals);
86 for (idx_t i = 0; i < n; i++)
87 quantizer->compute_residual (x + i * d, residuals+i*d, assign[i]);
88
89 trainset = residuals;
90 } else {
91 trainset = x;
92 }
93 if (verbose)
94 printf ("training %zdx%zd product quantizer on %ld vectors in %dD\n",
95 pq.M, pq.ksub, n, d);
96 pq.verbose = verbose;
97 pq.train (n, trainset);
98
99 if (do_polysemous_training) {
100 if (verbose)
101 printf("doing polysemous training for PQ\n");
102 PolysemousTraining default_pt;
103 PolysemousTraining *pt = polysemous_training;
104 if (!pt) pt = &default_pt;
105 pt->optimize_pq_for_hamming (pq, n, trainset);
106 }
107
108 // prepare second-level residuals for refine PQ
109 if (residuals_2) {
110 uint8_t *train_codes = new uint8_t [pq.code_size * n];
111 ScopeDeleter<uint8_t> del (train_codes);
112 pq.compute_codes (trainset, train_codes, n);
113
114 for (idx_t i = 0; i < n; i++) {
115 const float *xx = trainset + i * d;
116 float * res = residuals_2 + i * d;
117 pq.decode (train_codes + i * pq.code_size, res);
118 for (int j = 0; j < d; j++)
119 res[j] = xx[j] - res[j];
120 }
121
122 }
123
124 if (by_residual) {
125 precompute_table ();
126 }
127
128 }
129
130
131
132
133
134
135 /****************************************************************
136 * IVFPQ as codec */
137
138
139 /* produce a binary signature based on the residual vector */
140 void IndexIVFPQ::encode (long key, const float * x, uint8_t * code) const
141 {
142 if (by_residual) {
143 float residual_vec[d];
144 quantizer->compute_residual (x, residual_vec, key);
145 pq.compute_code (residual_vec, code);
146 }
147 else pq.compute_code (x, code);
148 }
149
150 void IndexIVFPQ::encode_multiple (size_t n, long *keys,
151 const float * x, uint8_t * xcodes,
152 bool compute_keys) const
153 {
154 if (compute_keys)
155 quantizer->assign (n, x, keys);
156
157 encode_vectors (n, x, keys, xcodes);
158 }
159
160 void IndexIVFPQ::decode_multiple (size_t n, const long *keys,
161 const uint8_t * xcodes, float * x) const
162 {
163 pq.decode (xcodes, x, n);
164 if (by_residual) {
165 std::vector<float> centroid (d);
166 for (size_t i = 0; i < n; i++) {
167 quantizer->reconstruct (keys[i], centroid.data());
168 float *xi = x + i * d;
169 for (size_t j = 0; j < d; j++) {
170 xi [j] += centroid [j];
171 }
172 }
173 }
174 }
175
176
177
178
179 /****************************************************************
180 * add */
181
182
183 void IndexIVFPQ::add_with_ids (idx_t n, const float * x, const long *xids)
184 {
185 add_core_o (n, x, xids, nullptr);
186 }
187
188
189 static float * compute_residuals (
190 const Index *quantizer,
191 Index::idx_t n, const float* x,
192 const Index::idx_t *list_nos)
193 {
194 size_t d = quantizer->d;
195 float *residuals = new float [n * d];
196 // TODO: parallelize?
197 for (size_t i = 0; i < n; i++) {
198 if (list_nos[i] < 0)
199 memset (residuals + i * d, 0, sizeof(*residuals) * d);
200 else
201 quantizer->compute_residual (
202 x + i * d, residuals + i * d, list_nos[i]);
203 }
204 return residuals;
205 }
206
207 void IndexIVFPQ::encode_vectors(idx_t n, const float* x,
208 const idx_t *list_nos,
209 uint8_t * codes) const
210 {
211 if (by_residual) {
212 float *to_encode = compute_residuals (quantizer, n, x, list_nos);
213 ScopeDeleter<float> del (to_encode);
214 pq.compute_codes (to_encode, codes, n);
215 } else {
216 pq.compute_codes (x, codes, n);
217 }
218 }
219
220
221 void IndexIVFPQ::add_core_o (idx_t n, const float * x, const long *xids,
222 float *residuals_2, const long *precomputed_idx)
223 {
224
225 idx_t bs = 32768;
226 if (n > bs) {
227 for (idx_t i0 = 0; i0 < n; i0 += bs) {
228 idx_t i1 = std::min(i0 + bs, n);
229 if (verbose) {
230 printf("IndexIVFPQ::add_core_o: adding %ld:%ld / %ld\n",
231 i0, i1, n);
232 }
233 add_core_o (i1 - i0, x + i0 * d,
234 xids ? xids + i0 : nullptr,
235 residuals_2 ? residuals_2 + i0 * d : nullptr,
236 precomputed_idx ? precomputed_idx + i0 : nullptr);
237 }
238 return;
239 }
240
241 FAISS_THROW_IF_NOT (is_trained);
242 double t0 = getmillisecs ();
243 const long * idx;
244 ScopeDeleter<long> del_idx;
245
246 if (precomputed_idx) {
247 idx = precomputed_idx;
248 } else {
249 long * idx0 = new long [n];
250 del_idx.set (idx0);
251 quantizer->assign (n, x, idx0);
252 idx = idx0;
253 }
254
255 double t1 = getmillisecs ();
256 uint8_t * xcodes = new uint8_t [n * code_size];
257 ScopeDeleter<uint8_t> del_xcodes (xcodes);
258
259 const float *to_encode = nullptr;
260 ScopeDeleter<float> del_to_encode;
261
262 if (by_residual) {
263 to_encode = compute_residuals (quantizer, n, x, idx);
264 del_to_encode.set (to_encode);
265 } else {
266 to_encode = x;
267 }
268 pq.compute_codes (to_encode, xcodes, n);
269
270 double t2 = getmillisecs ();
271 // TODO: parallelize?
272 size_t n_ignore = 0;
273 for (size_t i = 0; i < n; i++) {
274 idx_t key = idx[i];
275 if (key < 0) {
276 n_ignore ++;
277 if (residuals_2)
278 memset (residuals_2, 0, sizeof(*residuals_2) * d);
279 continue;
280 }
281 idx_t id = xids ? xids[i] : ntotal + i;
282
283 uint8_t *code = xcodes + i * code_size;
284 size_t offset = invlists->add_entry (key, id, code);
285
286 if (residuals_2) {
287 float *res2 = residuals_2 + i * d;
288 const float *xi = to_encode + i * d;
289 pq.decode (code, res2);
290 for (int j = 0; j < d; j++)
291 res2[j] = xi[j] - res2[j];
292 }
293
294 if (maintain_direct_map)
295 direct_map.push_back (key << 32 | offset);
296 }
297
298
299 double t3 = getmillisecs ();
300 if(verbose) {
301 char comment[100] = {0};
302 if (n_ignore > 0)
303 snprintf (comment, 100, "(%ld vectors ignored)", n_ignore);
304 printf(" add_core times: %.3f %.3f %.3f %s\n",
305 t1 - t0, t2 - t1, t3 - t2, comment);
306 }
307 ntotal += n;
308 }
309
310
311 void IndexIVFPQ::reconstruct_from_offset (long list_no, long offset,
312 float* recons) const
313 {
314 const uint8_t* code = invlists->get_single_code (list_no, offset);
315
316 if (by_residual) {
317 std::vector<float> centroid(d);
318 quantizer->reconstruct (list_no, centroid.data());
319
320 pq.decode (code, recons);
321 for (int i = 0; i < d; ++i) {
322 recons[i] += centroid[i];
323 }
324 } else {
325 pq.decode (code, recons);
326 }
327 }
328
329
330
331 /// 2G by default, accommodates tables up to PQ32 w/ 65536 centroids
332 size_t IndexIVFPQ::precomputed_table_max_bytes = ((size_t)1) << 31;
333
334 /** Precomputed tables for residuals
335 *
336 * During IVFPQ search with by_residual, we compute
337 *
338 * d = || x - y_C - y_R ||^2
339 *
340 * where x is the query vector, y_C the coarse centroid, y_R the
341 * refined PQ centroid. The expression can be decomposed as:
342 *
343 * d = || x - y_C ||^2 + || y_R ||^2 + 2 * (y_C|y_R) - 2 * (x|y_R)
344 * --------------- --------------------------- -------
345 * term 1 term 2 term 3
346 *
347 * When using multiprobe, we use the following decomposition:
348 * - term 1 is the distance to the coarse centroid, that is computed
349 * during the 1st stage search.
350 * - term 2 can be precomputed, as it does not involve x. However,
351 * because of the PQ, it needs nlist * M * ksub storage. This is why
352 * use_precomputed_table is off by default
353 * - term 3 is the classical non-residual distance table.
354 *
355 * Since y_R defined by a product quantizer, it is split across
356 * subvectors and stored separately for each subvector. If the coarse
357 * quantizer is a MultiIndexQuantizer then the table can be stored
358 * more compactly.
359 *
360 * At search time, the tables for term 2 and term 3 are added up. This
361 * is faster when the length of the lists is > ksub * M.
362 */
363
364 void IndexIVFPQ::precompute_table ()
365 {
366 if (use_precomputed_table == -1)
367 return;
368
369 if (use_precomputed_table == 0) { // then choose the type of table
370 if (quantizer->metric_type == METRIC_INNER_PRODUCT) {
371 if (verbose) {
372 printf("IndexIVFPQ::precompute_table: precomputed "
373 "tables not needed for inner product quantizers\n");
374 }
375 return;
376 }
377 const MultiIndexQuantizer *miq =
378 dynamic_cast<const MultiIndexQuantizer *> (quantizer);
379 if (miq && pq.M % miq->pq.M == 0)
380 use_precomputed_table = 2;
381 else {
382 size_t table_size = pq.M * pq.ksub * nlist * sizeof(float);
383 if (table_size > precomputed_table_max_bytes) {
384 if (verbose) {
385 printf(
386 "IndexIVFPQ::precompute_table: not precomputing table, "
387 "it would be too big: %ld bytes (max %ld)\n",
388 table_size, precomputed_table_max_bytes);
389 use_precomputed_table = 0;
390 }
391 return;
392 }
393 use_precomputed_table = 1;
394 }
395 } // otherwise assume user has set appropriate flag on input
396
397 if (verbose) {
398 printf ("precomputing IVFPQ tables type %d\n",
399 use_precomputed_table);
400 }
401
402 // squared norms of the PQ centroids
403 std::vector<float> r_norms (pq.M * pq.ksub, NAN);
404 for (int m = 0; m < pq.M; m++)
405 for (int j = 0; j < pq.ksub; j++)
406 r_norms [m * pq.ksub + j] =
407 fvec_norm_L2sqr (pq.get_centroids (m, j), pq.dsub);
408
409 if (use_precomputed_table == 1) {
410
411 precomputed_table.resize (nlist * pq.M * pq.ksub);
412 std::vector<float> centroid (d);
413
414 for (size_t i = 0; i < nlist; i++) {
415 quantizer->reconstruct (i, centroid.data());
416
417 float *tab = &precomputed_table[i * pq.M * pq.ksub];
418 pq.compute_inner_prod_table (centroid.data(), tab);
419 fvec_madd (pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
420 }
421 } else if (use_precomputed_table == 2) {
422 const MultiIndexQuantizer *miq =
423 dynamic_cast<const MultiIndexQuantizer *> (quantizer);
424 FAISS_THROW_IF_NOT (miq);
425 const ProductQuantizer &cpq = miq->pq;
426 FAISS_THROW_IF_NOT (pq.M % cpq.M == 0);
427
428 precomputed_table.resize(cpq.ksub * pq.M * pq.ksub);
429
430 // reorder PQ centroid table
431 std::vector<float> centroids (d * cpq.ksub, NAN);
432
433 for (int m = 0; m < cpq.M; m++) {
434 for (size_t i = 0; i < cpq.ksub; i++) {
435 memcpy (centroids.data() + i * d + m * cpq.dsub,
436 cpq.get_centroids (m, i),
437 sizeof (*centroids.data()) * cpq.dsub);
438 }
439 }
440
441 pq.compute_inner_prod_tables (cpq.ksub, centroids.data (),
442 precomputed_table.data ());
443
444 for (size_t i = 0; i < cpq.ksub; i++) {
445 float *tab = &precomputed_table[i * pq.M * pq.ksub];
446 fvec_madd (pq.M * pq.ksub, r_norms.data(), 2.0, tab, tab);
447 }
448
449 }
450 }
451
452 namespace {
453
454 using idx_t = Index::idx_t;
455
456 static uint64_t get_cycles () {
457 #ifdef __x86_64__
458 uint32_t high, low;
459 asm volatile("rdtsc \n\t"
460 : "=a" (low),
461 "=d" (high));
462 return ((uint64_t)high << 32) | (low);
463 #else
464 return 0;
465 #endif
466 }
467
468 #define TIC t0 = get_cycles()
469 #define TOC get_cycles () - t0
470
471
472
473 /** QueryTables manages the various ways of searching an
474 * IndexIVFPQ. The code contains a lot of branches, depending on:
475 * - metric_type: are we computing L2 or Inner product similarity?
476 * - by_residual: do we encode raw vectors or residuals?
477 * - use_precomputed_table: are x_R|x_C tables precomputed?
478 * - polysemous_ht: are we filtering with polysemous codes?
479 */
480 struct QueryTables {
481
482 /*****************************************************
483 * General data from the IVFPQ
484 *****************************************************/
485
486 const IndexIVFPQ & ivfpq;
487 const IVFSearchParameters *params;
488
489 // copied from IndexIVFPQ for easier access
490 int d;
491 const ProductQuantizer & pq;
492 MetricType metric_type;
493 bool by_residual;
494 int use_precomputed_table;
495 int polysemous_ht;
496
497 // pre-allocated data buffers
498 float * sim_table, * sim_table_2;
499 float * residual_vec, *decoded_vec;
500
501 // single data buffer
502 std::vector<float> mem;
503
504 // for table pointers
505 std::vector<const float *> sim_table_ptrs;
506
507 explicit QueryTables (const IndexIVFPQ & ivfpq,
508 const IVFSearchParameters *params):
509 ivfpq(ivfpq),
510 d(ivfpq.d),
511 pq (ivfpq.pq),
512 metric_type (ivfpq.metric_type),
513 by_residual (ivfpq.by_residual),
514 use_precomputed_table (ivfpq.use_precomputed_table)
515 {
516 mem.resize (pq.ksub * pq.M * 2 + d * 2);
517 sim_table = mem.data ();
518 sim_table_2 = sim_table + pq.ksub * pq.M;
519 residual_vec = sim_table_2 + pq.ksub * pq.M;
520 decoded_vec = residual_vec + d;
521
522 // for polysemous
523 polysemous_ht = ivfpq.polysemous_ht;
524 if (auto ivfpq_params =
525 dynamic_cast<const IVFPQSearchParameters *>(params)) {
526 polysemous_ht = ivfpq_params->polysemous_ht;
527 }
528 if (polysemous_ht != 0) {
529 q_code.resize (pq.code_size);
530 }
531 init_list_cycles = 0;
532 sim_table_ptrs.resize (pq.M);
533 }
534
535 /*****************************************************
536 * What we do when query is known
537 *****************************************************/
538
539 // field specific to query
540 const float * qi;
541
542 // query-specific intialization
543 void init_query (const float * qi) {
544 this->qi = qi;
545 if (metric_type == METRIC_INNER_PRODUCT)
546 init_query_IP ();
547 else
548 init_query_L2 ();
549 if (!by_residual && polysemous_ht != 0)
550 pq.compute_code (qi, q_code.data());
551 }
552
553 void init_query_IP () {
554 // precompute some tables specific to the query qi
555 pq.compute_inner_prod_table (qi, sim_table);
556 }
557
558 void init_query_L2 () {
559 if (!by_residual) {
560 pq.compute_distance_table (qi, sim_table);
561 } else if (use_precomputed_table) {
562 pq.compute_inner_prod_table (qi, sim_table_2);
563 }
564 }
565
566 /*****************************************************
567 * When inverted list is known: prepare computations
568 *****************************************************/
569
570 // fields specific to list
571 Index::idx_t key;
572 float coarse_dis;
573 std::vector<uint8_t> q_code;
574
575 uint64_t init_list_cycles;
576
577 /// once we know the query and the centroid, we can prepare the
578 /// sim_table that will be used for accumulation
579 /// and dis0, the initial value
580 float precompute_list_tables () {
581 float dis0 = 0;
582 uint64_t t0; TIC;
583 if (by_residual) {
584 if (metric_type == METRIC_INNER_PRODUCT)
585 dis0 = precompute_list_tables_IP ();
586 else
587 dis0 = precompute_list_tables_L2 ();
588 }
589 init_list_cycles += TOC;
590 return dis0;
591 }
592
593 float precompute_list_table_pointers () {
594 float dis0 = 0;
595 uint64_t t0; TIC;
596 if (by_residual) {
597 if (metric_type == METRIC_INNER_PRODUCT)
598 FAISS_THROW_MSG ("not implemented");
599 else
600 dis0 = precompute_list_table_pointers_L2 ();
601 }
602 init_list_cycles += TOC;
603 return dis0;
604 }
605
606 /*****************************************************
607 * compute tables for inner prod
608 *****************************************************/
609
610 float precompute_list_tables_IP ()
611 {
612 // prepare the sim_table that will be used for accumulation
613 // and dis0, the initial value
614 ivfpq.quantizer->reconstruct (key, decoded_vec);
615 // decoded_vec = centroid
616 float dis0 = fvec_inner_product (qi, decoded_vec, d);
617
618 if (polysemous_ht) {
619 for (int i = 0; i < d; i++) {
620 residual_vec [i] = qi[i] - decoded_vec[i];
621 }
622 pq.compute_code (residual_vec, q_code.data());
623 }
624 return dis0;
625 }
626
627
628 /*****************************************************
629 * compute tables for L2 distance
630 *****************************************************/
631
632 float precompute_list_tables_L2 ()
633 {
634 float dis0 = 0;
635
636 if (use_precomputed_table == 0 || use_precomputed_table == -1) {
637 ivfpq.quantizer->compute_residual (qi, residual_vec, key);
638 pq.compute_distance_table (residual_vec, sim_table);
639 } else if (use_precomputed_table == 1) {
640 dis0 = coarse_dis;
641
642 fvec_madd (pq.M * pq.ksub,
643 &ivfpq.precomputed_table [key * pq.ksub * pq.M],
644 -2.0, sim_table_2,
645 sim_table);
646 } else if (use_precomputed_table == 2) {
647 dis0 = coarse_dis;
648
649 const MultiIndexQuantizer *miq =
650 dynamic_cast<const MultiIndexQuantizer *> (ivfpq.quantizer);
651 FAISS_THROW_IF_NOT (miq);
652 const ProductQuantizer &cpq = miq->pq;
653 int Mf = pq.M / cpq.M;
654
655 const float *qtab = sim_table_2; // query-specific table
656 float *ltab = sim_table; // (output) list-specific table
657
658 long k = key;
659 for (int cm = 0; cm < cpq.M; cm++) {
660 // compute PQ index
661 int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
662 k >>= cpq.nbits;
663
664 // get corresponding table
665 const float *pc = &ivfpq.precomputed_table
666 [(ki * pq.M + cm * Mf) * pq.ksub];
667
668 if (polysemous_ht == 0) {
669
670 // sum up with query-specific table
671 fvec_madd (Mf * pq.ksub,
672 pc,
673 -2.0, qtab,
674 ltab);
675 ltab += Mf * pq.ksub;
676 qtab += Mf * pq.ksub;
677 } else {
678 for (int m = cm * Mf; m < (cm + 1) * Mf; m++) {
679 q_code[m] = fvec_madd_and_argmin
680 (pq.ksub, pc, -2, qtab, ltab);
681 pc += pq.ksub;
682 ltab += pq.ksub;
683 qtab += pq.ksub;
684 }
685 }
686
687 }
688 }
689
690 return dis0;
691 }
692
693 float precompute_list_table_pointers_L2 ()
694 {
695 float dis0 = 0;
696
697 if (use_precomputed_table == 1) {
698 dis0 = coarse_dis;
699
700 const float * s = &ivfpq.precomputed_table [key * pq.ksub * pq.M];
701 for (int m = 0; m < pq.M; m++) {
702 sim_table_ptrs [m] = s;
703 s += pq.ksub;
704 }
705 } else if (use_precomputed_table == 2) {
706 dis0 = coarse_dis;
707
708 const MultiIndexQuantizer *miq =
709 dynamic_cast<const MultiIndexQuantizer *> (ivfpq.quantizer);
710 FAISS_THROW_IF_NOT (miq);
711 const ProductQuantizer &cpq = miq->pq;
712 int Mf = pq.M / cpq.M;
713
714 long k = key;
715 int m0 = 0;
716 for (int cm = 0; cm < cpq.M; cm++) {
717 int ki = k & ((uint64_t(1) << cpq.nbits) - 1);
718 k >>= cpq.nbits;
719
720 const float *pc = &ivfpq.precomputed_table
721 [(ki * pq.M + cm * Mf) * pq.ksub];
722
723 for (int m = m0; m < m0 + Mf; m++) {
724 sim_table_ptrs [m] = pc;
725 pc += pq.ksub;
726 }
727 m0 += Mf;
728 }
729 } else {
730 FAISS_THROW_MSG ("need precomputed tables");
731 }
732
733 if (polysemous_ht) {
734 FAISS_THROW_MSG ("not implemented");
735 // Not clear that it makes sense to implemente this,
736 // because it costs M * ksub, which is what we wanted to
737 // avoid with the tables pointers.
738 }
739
740 return dis0;
741 }
742
743
744 };
745
746
747 /*****************************************************
748 * Scaning the codes.
749 * The scanning functions call their favorite precompute_*
750 * function to precompute the tables they need.
751 *****************************************************/
752 template <typename IDType, bool store_pairs, class C, MetricType METRIC_TYPE>
753 struct IVFPQScannerT: QueryTables {
754
755 const uint8_t * list_codes;
756 const IDType * list_ids;
757 size_t list_size;
758
759 explicit IVFPQScannerT (const IndexIVFPQ & ivfpq,
760 const IVFSearchParameters *params):
761 QueryTables (ivfpq, params)
762 {
763 FAISS_THROW_IF_NOT (pq.byte_per_idx == 1);
764 assert(METRIC_TYPE == metric_type);
765 }
766
767 float dis0;
768
769 void init_list (idx_t list_no, float coarse_dis,
770 int mode) {
771 this->key = list_no;
772 this->coarse_dis = coarse_dis;
773
774 if (mode == 2) {
775 dis0 = precompute_list_tables ();
776 } else if (mode == 1) {
777 dis0 = precompute_list_table_pointers ();
778 }
779 }
780
781 /*****************************************************
782 * Scaning the codes: simple PQ scan.
783 *****************************************************/
784
785 /// version of the scan where we use precomputed tables
786 size_t scan_list_with_table (
787 size_t ncode, const uint8_t *codes, const idx_t *ids,
788 size_t k, float * heap_sim, long * heap_ids) const
789 {
790 int nup = 0;
791 for (size_t j = 0; j < ncode; j++) {
792
793 float dis = dis0;
794 const float *tab = sim_table;
795
796 for (size_t m = 0; m < pq.M; m++) {
797 dis += tab[*codes++];
798 tab += pq.ksub;
799 }
800
801 if (C::cmp (heap_sim[0], dis)) {
802 heap_pop<C> (k, heap_sim, heap_ids);
803 long id = store_pairs ? (key << 32 | j) : ids[j];
804 heap_push<C> (k, heap_sim, heap_ids, dis, id);
805 nup++;
806 }
807 }
808 return nup;
809 }
810
811
812 /// tables are not precomputed, but pointers are provided to the
813 /// relevant X_c|x_r tables
814 size_t scan_list_with_pointer (
815 size_t ncode, const uint8_t *codes, const idx_t *ids,
816 size_t k, float * heap_sim, long * heap_ids) const
817 {
818 size_t nup = 0;
819 for (size_t j = 0; j < ncode; j++) {
820
821 float dis = dis0;
822 const float *tab = sim_table_2;
823
824 for (size_t m = 0; m < pq.M; m++) {
825 int ci = *codes++;
826 dis += sim_table_ptrs [m][ci] - 2 * tab [ci];
827 tab += pq.ksub;
828 }
829
830 if (C::cmp (heap_sim[0], dis)) {
831 heap_pop<C> (k, heap_sim, heap_ids);
832 long id = store_pairs ? (key << 32 | j) : ids[j];
833 heap_push<C> (k, heap_sim, heap_ids, dis, id);
834 nup++;
835 }
836 }
837 return nup;
838 }
839
840
841 /// nothing is precomputed: access residuals on-the-fly
842 size_t scan_on_the_fly_dist (
843 size_t ncode, const uint8_t *codes, const idx_t *ids,
844 size_t k, float * heap_sim, long * heap_ids) const
845 {
846 const float *dvec;
847 float dis0 = 0;
848 size_t nup = 0;
849 if (by_residual) {
850 if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
851 ivfpq.quantizer->reconstruct (key, residual_vec);
852 dis0 = fvec_inner_product (residual_vec, qi, d);
853 } else {
854 ivfpq.quantizer->compute_residual (qi, residual_vec, key);
855 }
856 dvec = residual_vec;
857 } else {
858 dvec = qi;
859 dis0 = 0;
860 }
861
862 for (size_t j = 0; j < ncode; j++) {
863
864 pq.decode (codes, decoded_vec);
865 codes += pq.code_size;
866
867 float dis;
868 if (METRIC_TYPE == METRIC_INNER_PRODUCT) {
869 dis = dis0 + fvec_inner_product (decoded_vec, qi, d);
870 } else {
871 dis = fvec_L2sqr (decoded_vec, dvec, d);
872 }
873
874 if (C::cmp (heap_sim[0], dis)) {
875 heap_pop<C> (k, heap_sim, heap_ids);
876 long id = store_pairs ? (key << 32 | j) : ids[j];
877 heap_push<C> (k, heap_sim, heap_ids, dis, id);
878 nup++;
879 }
880 }
881 return nup;
882 }
883
884 /*****************************************************
885 * Scanning codes with polysemous filtering
886 *****************************************************/
887
888 template <class HammingComputer>
889 size_t scan_list_polysemous_hc (
890 size_t ncode, const uint8_t *codes, const idx_t *ids,
891 size_t k, float * heap_sim, long * heap_ids) const
892 {
893 int ht = ivfpq.polysemous_ht;
894 size_t n_hamming_pass = 0, nup = 0;
895
896 int code_size = pq.code_size;
897
898 HammingComputer hc (q_code.data(), code_size);
899
900 for (size_t j = 0; j < ncode; j++) {
901 const uint8_t *b_code = codes;
902 int hd = hc.hamming (b_code);
903 if (hd < ht) {
904 n_hamming_pass ++;
905
906 float dis = dis0;
907 const float *tab = sim_table;
908
909 for (size_t m = 0; m < pq.M; m++) {
910 dis += tab[*b_code++];
911 tab += pq.ksub;
912 }
913
914 if (C::cmp (heap_sim[0], dis)) {
915 heap_pop<C> (k, heap_sim, heap_ids);
916 long id = store_pairs ? (key << 32 | j) : ids[j];
917 heap_push<C> (k, heap_sim, heap_ids, dis, id);
918 nup++;
919 }
920 }
921 codes += code_size;
922 }
923 #pragma omp critical
924 {
925 indexIVFPQ_stats.n_hamming_pass += n_hamming_pass;
926 }
927 return nup;
928 }
929
930 size_t scan_list_polysemous (
931 size_t ncode, const uint8_t *codes, const idx_t *ids,
932 size_t k, float * heap_sim, long * heap_ids) const
933 {
934 switch (pq.code_size) {
935 #define HANDLE_CODE_SIZE(cs) \
936 case cs: \
937 return scan_list_polysemous_hc <HammingComputer ## cs> \
938 (ncode, codes, ids, k, heap_sim, heap_ids); \
939 break
940 HANDLE_CODE_SIZE(4);
941 HANDLE_CODE_SIZE(8);
942 HANDLE_CODE_SIZE(16);
943 HANDLE_CODE_SIZE(20);
944 HANDLE_CODE_SIZE(32);
945 HANDLE_CODE_SIZE(64);
946 #undef HANDLE_CODE_SIZE
947 default:
948 if (pq.code_size % 8 == 0)
949 return scan_list_polysemous_hc <HammingComputerM8>
950 (ncode, codes, ids, k, heap_sim, heap_ids);
951 else
952 return scan_list_polysemous_hc <HammingComputerM4>
953 (ncode, codes, ids, k, heap_sim, heap_ids);
954 break;
955 }
956 }
957
958 };
959
960
961 /* We put as many parameters as possible in template. Hopefully the
962 * gain in runtime is worth the code bloat. C is the comparator < or
963 * >, it is directly related to METRIC_TYPE. precompute_mode is how
964 * much we precompute (2 = precompute distance tables, 1 = precompute
965 * pointers to distances, 0 = compute distances one by one). Currently
966 * only 2 is supported. */
967 template<MetricType METRIC_TYPE, bool store_pairs, class C,
968 int precompute_mode>
969 struct IVFPQScanner:
970 IVFPQScannerT<Index::idx_t, store_pairs, C, METRIC_TYPE>,
971 InvertedListScanner
972 {
973
974 IVFPQScanner(const IndexIVFPQ & ivfpq):
975 IVFPQScannerT<Index::idx_t, store_pairs, C, METRIC_TYPE>(ivfpq, nullptr)
976 {
977 }
978
979 void set_query (const float *query) override {
980 this->init_query (query);
981 }
982
983 void set_list (idx_t list_no, float coarse_dis) override {
984 this->init_list (list_no, coarse_dis, precompute_mode);
985 }
986
987 float distance_to_code (const uint8_t *code) const override {
988 assert(precompute_mode == 2);
989 float dis = this->dis0;
990 const float *tab = this->sim_table;
991
992 for (size_t m = 0; m < this->pq.M; m++) {
993 dis += tab[*code++];
994 tab += this->pq.ksub;
995 }
996 return dis;
997 }
998
999 size_t scan_codes (size_t ncode,
1000 const uint8_t *codes,
1001 const idx_t *ids,
1002 float *heap_sim, idx_t *heap_ids,
1003 size_t k) const override
1004 {
1005 if (this->polysemous_ht > 0) {
1006 assert(precompute_mode == 2);
1007 this->scan_list_polysemous
1008 (ncode, codes, ids, k, heap_sim, heap_ids);
1009 } else if (precompute_mode == 2) {
1010 this->scan_list_with_table
1011 (ncode, codes, ids, k, heap_sim, heap_ids);
1012 } else if (precompute_mode == 1) {
1013 this->scan_list_with_pointer
1014 (ncode, codes, ids, k, heap_sim, heap_ids);
1015 } else if (precompute_mode == 0) {
1016 this->scan_on_the_fly_dist
1017 (ncode, codes, ids, k, heap_sim, heap_ids);
1018 } else {
1019 FAISS_THROW_MSG("bad precomp mode");
1020 }
1021 return 0;
1022 }
1023
1024 };
1025
1026
1027
1028
1029 } // anonymous namespace
1030
1031 InvertedListScanner *
1032 IndexIVFPQ::get_InvertedListScanner (bool store_pairs) const
1033 {
1034 if (metric_type == METRIC_INNER_PRODUCT) {
1035 if (store_pairs) {
1036 return new IVFPQScanner<
1037 METRIC_INNER_PRODUCT, true, CMin<float, long>, 2> (*this);
1038 } else {
1039 return new IVFPQScanner<
1040 METRIC_INNER_PRODUCT, false, CMin<float, long>, 2>(*this);
1041 }
1042 } else if (metric_type == METRIC_L2) {
1043 if (store_pairs) {
1044 return new IVFPQScanner<
1045 METRIC_L2, true, CMax<float, long>, 2> (*this);
1046 } else {
1047 return new IVFPQScanner<
1048 METRIC_L2, false, CMax<float, long>, 2>(*this);
1049 }
1050 }
1051 return nullptr;
1052
1053 }
1054
1055
1056
1057 IndexIVFPQStats indexIVFPQ_stats;
1058
1059 void IndexIVFPQStats::reset () {
1060 memset (this, 0, sizeof (*this));
1061 }
1062
1063
1064
1065 IndexIVFPQ::IndexIVFPQ ()
1066 {
1067 // initialize some runtime values
1068 use_precomputed_table = 0;
1069 scan_table_threshold = 0;
1070 do_polysemous_training = false;
1071 polysemous_ht = 0;
1072 polysemous_training = nullptr;
1073 }
1074
1075
1076 struct CodeCmp {
1077 const uint8_t *tab;
1078 size_t code_size;
1079 bool operator () (int a, int b) const {
1080 return cmp (a, b) > 0;
1081 }
1082 int cmp (int a, int b) const {
1083 return memcmp (tab + a * code_size, tab + b * code_size,
1084 code_size);
1085 }
1086 };
1087
1088
1089 size_t IndexIVFPQ::find_duplicates (idx_t *dup_ids, size_t *lims) const
1090 {
1091 size_t ngroup = 0;
1092 lims[0] = 0;
1093 for (size_t list_no = 0; list_no < nlist; list_no++) {
1094 size_t n = invlists->list_size (list_no);
1095 std::vector<int> ord (n);
1096 for (int i = 0; i < n; i++) ord[i] = i;
1097 InvertedLists::ScopedCodes codes (invlists, list_no);
1098 CodeCmp cs = { codes.get(), code_size };
1099 std::sort (ord.begin(), ord.end(), cs);
1100
1101 InvertedLists::ScopedIds list_ids (invlists, list_no);
1102 int prev = -1; // all elements from prev to i-1 are equal
1103 for (int i = 0; i < n; i++) {
1104 if (prev >= 0 && cs.cmp (ord [prev], ord [i]) == 0) {
1105 // same as previous => remember
1106 if (prev + 1 == i) { // start new group
1107 ngroup++;
1108 lims[ngroup] = lims[ngroup - 1];
1109 dup_ids [lims [ngroup]++] = list_ids [ord [prev]];
1110 }
1111 dup_ids [lims [ngroup]++] = list_ids [ord [i]];
1112 } else { // not same as previous.
1113 prev = i;
1114 }
1115 }
1116 }
1117 return ngroup;
1118 }
1119
1120
1121
1122
1123 /*****************************************
1124 * IndexIVFPQR implementation
1125 ******************************************/
1126
1127 IndexIVFPQR::IndexIVFPQR (
1128 Index * quantizer, size_t d, size_t nlist,
1129 size_t M, size_t nbits_per_idx,
1130 size_t M_refine, size_t nbits_per_idx_refine):
1131 IndexIVFPQ (quantizer, d, nlist, M, nbits_per_idx),
1132 refine_pq (d, M_refine, nbits_per_idx_refine),
1133 k_factor (4)
1134 {
1135 by_residual = true;
1136 }
1137
1138 IndexIVFPQR::IndexIVFPQR ():
1139 k_factor (1)
1140 {
1141 by_residual = true;
1142 }
1143
1144
1145
1146 void IndexIVFPQR::reset()
1147 {
1148 IndexIVFPQ::reset();
1149 refine_codes.clear();
1150 }
1151
1152
1153
1154
1155 void IndexIVFPQR::train_residual (idx_t n, const float *x)
1156 {
1157
1158 float * residual_2 = new float [n * d];
1159 ScopeDeleter <float> del(residual_2);
1160
1161 train_residual_o (n, x, residual_2);
1162
1163 if (verbose)
1164 printf ("training %zdx%zd 2nd level PQ quantizer on %ld %dD-vectors\n",
1165 refine_pq.M, refine_pq.ksub, n, d);
1166
1167 refine_pq.cp.max_points_per_centroid = 1000;
1168 refine_pq.cp.verbose = verbose;
1169
1170 refine_pq.train (n, residual_2);
1171
1172 }
1173
1174
1175 void IndexIVFPQR::add_with_ids (idx_t n, const float *x, const long *xids) {
1176 add_core (n, x, xids, nullptr);
1177 }
1178
1179 void IndexIVFPQR::add_core (idx_t n, const float *x, const long *xids,
1180 const long *precomputed_idx) {
1181
1182 float * residual_2 = new float [n * d];
1183 ScopeDeleter <float> del(residual_2);
1184
1185 idx_t n0 = ntotal;
1186
1187 add_core_o (n, x, xids, residual_2, precomputed_idx);
1188
1189 refine_codes.resize (ntotal * refine_pq.code_size);
1190
1191 refine_pq.compute_codes (
1192 residual_2, &refine_codes[n0 * refine_pq.code_size], n);
1193
1194
1195 }
1196
1197
1198 void IndexIVFPQR::search_preassigned (idx_t n, const float *x, idx_t k,
1199 const idx_t *idx,
1200 const float *L1_dis,
1201 float *distances, idx_t *labels,
1202 bool store_pairs,
1203 const IVFSearchParameters *params
1204 ) const
1205 {
1206 uint64_t t0;
1207 TIC;
1208 size_t k_coarse = long(k * k_factor);
1209 idx_t *coarse_labels = new idx_t [k_coarse * n];
1210 ScopeDeleter<idx_t> del1 (coarse_labels);
1211 { // query with quantizer levels 1 and 2.
1212 float *coarse_distances = new float [k_coarse * n];
1213 ScopeDeleter<float> del(coarse_distances);
1214
1215 IndexIVFPQ::search_preassigned (
1216 n, x, k_coarse,
1217 idx, L1_dis, coarse_distances, coarse_labels,
1218 true, params);
1219 }
1220
1221
1222 indexIVFPQ_stats.search_cycles += TOC;
1223
1224 TIC;
1225
1226 // 3rd level refinement
1227 size_t n_refine = 0;
1228 #pragma omp parallel reduction(+ : n_refine)
1229 {
1230 // tmp buffers
1231 float *residual_1 = new float [2 * d];
1232 ScopeDeleter<float> del (residual_1);
1233 float *residual_2 = residual_1 + d;
1234 #pragma omp for
1235 for (idx_t i = 0; i < n; i++) {
1236 const float *xq = x + i * d;
1237 const long * shortlist = coarse_labels + k_coarse * i;
1238 float * heap_sim = distances + k * i;
1239 long * heap_ids = labels + k * i;
1240 maxheap_heapify (k, heap_sim, heap_ids);
1241
1242 for (int j = 0; j < k_coarse; j++) {
1243 long sl = shortlist[j];
1244
1245 if (sl == -1) continue;
1246
1247 int list_no = sl >> 32;
1248 int ofs = sl & 0xffffffff;
1249
1250 assert (list_no >= 0 && list_no < nlist);
1251 assert (ofs >= 0 && ofs < invlists->list_size (list_no));
1252
1253 // 1st level residual
1254 quantizer->compute_residual (xq, residual_1, list_no);
1255
1256 // 2nd level residual
1257 const uint8_t * l2code =
1258 invlists->get_single_code (list_no, ofs);
1259
1260 pq.decode (l2code, residual_2);
1261 for (int l = 0; l < d; l++)
1262 residual_2[l] = residual_1[l] - residual_2[l];
1263
1264 // 3rd level residual's approximation
1265 idx_t id = invlists->get_single_id (list_no, ofs);
1266 assert (0 <= id && id < ntotal);
1267 refine_pq.decode (&refine_codes [id * refine_pq.code_size],
1268 residual_1);
1269
1270 float dis = fvec_L2sqr (residual_1, residual_2, d);
1271
1272 if (dis < heap_sim[0]) {
1273 maxheap_pop (k, heap_sim, heap_ids);
1274 long id_or_pair = store_pairs ? sl : id;
1275 maxheap_push (k, heap_sim, heap_ids, dis, id_or_pair);
1276 }
1277 n_refine ++;
1278 }
1279 maxheap_reorder (k, heap_sim, heap_ids);
1280 }
1281 }
1282 indexIVFPQ_stats.nrefine += n_refine;
1283 indexIVFPQ_stats.refine_cycles += TOC;
1284 }
1285
1286 void IndexIVFPQR::reconstruct_from_offset (long list_no, long offset,
1287 float* recons) const
1288 {
1289 IndexIVFPQ::reconstruct_from_offset (list_no, offset, recons);
1290
1291 idx_t id = invlists->get_single_id (list_no, offset);
1292 assert (0 <= id && id < ntotal);
1293
1294 std::vector<float> r3(d);
1295 refine_pq.decode (&refine_codes [id * refine_pq.code_size], r3.data());
1296 for (int i = 0; i < d; ++i) {
1297 recons[i] += r3[i];
1298 }
1299 }
1300
1301 void IndexIVFPQR::merge_from (IndexIVF &other_in, idx_t add_id)
1302 {
1303 IndexIVFPQR *other = dynamic_cast<IndexIVFPQR *> (&other_in);
1304 FAISS_THROW_IF_NOT(other);
1305
1306 IndexIVF::merge_from (other_in, add_id);
1307
1308 refine_codes.insert (refine_codes.end(),
1309 other->refine_codes.begin(),
1310 other->refine_codes.end());
1311 other->refine_codes.clear();
1312 }
1313
1314 long IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
1315 FAISS_THROW_MSG("not implemented");
1316 return 0;
1317 }
1318
1319 /*************************************
1320 * Index2Layer implementation
1321 *************************************/
1322
1323
1324 Index2Layer::Index2Layer (Index * quantizer, size_t nlist,
1325 int M,
1326 MetricType metric):
1327 Index (quantizer->d, metric),
1328 q1 (quantizer, nlist),
1329 pq (quantizer->d, M, 8)
1330 {
1331 is_trained = false;
1332 for (int nbyte = 0; nbyte < 7; nbyte++) {
1333 if ((1L << (8 * nbyte)) >= nlist) {
1334 code_size_1 = nbyte;
1335 break;
1336 }
1337 }
1338 code_size_2 = pq.code_size;
1339 code_size = code_size_1 + code_size_2;
1340 }
1341
1342 Index2Layer::Index2Layer ()
1343 {
1344 code_size = code_size_1 = code_size_2 = 0;
1345 }
1346
1347 Index2Layer::~Index2Layer ()
1348 {}
1349
1350 void Index2Layer::train(idx_t n, const float* x)
1351 {
1352 if (verbose) {
1353 printf ("training level-1 quantizer %ld vectors in %dD\n",
1354 n, d);
1355 }
1356
1357 q1.train_q1 (n, x, verbose, metric_type);
1358
1359 if (verbose) {
1360 printf("computing residuals\n");
1361 }
1362
1363 const float * x_in = x;
1364
1365 x = fvecs_maybe_subsample (
1366 d, (size_t*)&n, pq.cp.max_points_per_centroid * pq.ksub,
1367 x, verbose, pq.cp.seed);
1368
1369 ScopeDeleter<float> del_x (x_in == x ? nullptr : x);
1370
1371 std::vector<idx_t> assign(n); // assignement to coarse centroids
1372 q1.quantizer->assign (n, x, assign.data());
1373 std::vector<float> residuals(n * d);
1374 for (idx_t i = 0; i < n; i++) {
1375 q1.quantizer->compute_residual (
1376 x + i * d, residuals.data() + i * d, assign[i]);
1377 }
1378
1379 if (verbose)
1380 printf ("training %zdx%zd product quantizer on %ld vectors in %dD\n",
1381 pq.M, pq.ksub, n, d);
1382 pq.verbose = verbose;
1383 pq.train (n, residuals.data());
1384
1385 is_trained = true;
1386 }
1387
1388 void Index2Layer::add(idx_t n, const float* x)
1389 {
1390 idx_t bs = 32768;
1391 if (n > bs) {
1392 for (idx_t i0 = 0; i0 < n; i0 += bs) {
1393 idx_t i1 = std::min(i0 + bs, n);
1394 if (verbose) {
1395 printf("Index2Layer::add: adding %ld:%ld / %ld\n",
1396 i0, i1, n);
1397 }
1398 add (i1 - i0, x + i0 * d);
1399 }
1400 return;
1401 }
1402
1403 std::vector<idx_t> codes1 (n);
1404 q1.quantizer->assign (n, x, codes1.data());
1405 std::vector<float> residuals(n * d);
1406 for (idx_t i = 0; i < n; i++) {
1407 q1.quantizer->compute_residual (
1408 x + i * d, residuals.data() + i * d, codes1[i]);
1409 }
1410 std::vector<uint8_t> codes2 (n * code_size_2);
1411
1412 pq.compute_codes (residuals.data(), codes2.data(), n);
1413
1414 codes.resize ((ntotal + n) * code_size);
1415 uint8_t *wp = &codes[ntotal * code_size];
1416
1417 {
1418 int i = 0x11223344;
1419 const char *ip = (char*)&i;
1420 FAISS_THROW_IF_NOT_MSG (ip[0] == 0x44,
1421 "works only on a little-endian CPU");
1422 }
1423
1424 // copy to output table
1425 for (idx_t i = 0; i < n; i++) {
1426 memcpy (wp, &codes1[i], code_size_1);
1427 wp += code_size_1;
1428 memcpy (wp, &codes2[i * code_size_2], code_size_2);
1429 wp += code_size_2;
1430 }
1431
1432 ntotal += n;
1433
1434 }
1435
1436 void Index2Layer::search(
1437 idx_t /*n*/,
1438 const float* /*x*/,
1439 idx_t /*k*/,
1440 float* /*distances*/,
1441 idx_t* /*labels*/) const {
1442 FAISS_THROW_MSG("not implemented");
1443 }
1444
1445
1446 void Index2Layer::reconstruct_n(idx_t i0, idx_t ni, float* recons) const
1447 {
1448 float recons1[d];
1449 FAISS_THROW_IF_NOT (i0 >= 0 && i0 + ni <= ntotal);
1450 const uint8_t *rp = &codes[i0 * code_size];
1451
1452 for (idx_t i = 0; i < ni; i++) {
1453 idx_t key = 0;
1454 memcpy (&key, rp, code_size_1);
1455 q1.quantizer->reconstruct (key, recons1);
1456 rp += code_size_1;
1457 pq.decode (rp, recons);
1458 for (idx_t j = 0; j < d; j++) {
1459 recons[j] += recons1[j];
1460 }
1461 rp += code_size_2;
1462 recons += d;
1463 }
1464 }
1465
1466 void Index2Layer::transfer_to_IVFPQ (IndexIVFPQ & other) const
1467 {
1468 FAISS_THROW_IF_NOT (other.nlist == q1.nlist);
1469 FAISS_THROW_IF_NOT (other.code_size == code_size_2);
1470 FAISS_THROW_IF_NOT (other.ntotal == 0);
1471
1472 const uint8_t *rp = codes.data();
1473
1474 for (idx_t i = 0; i < ntotal; i++) {
1475 idx_t key = 0;
1476 memcpy (&key, rp, code_size_1);
1477 rp += code_size_1;
1478 other.invlists->add_entry (key, i, rp);
1479 rp += code_size_2;
1480 }
1481
1482 other.ntotal = ntotal;
1483
1484 }
1485
1486
1487
1488 void Index2Layer::reconstruct(idx_t key, float* recons) const
1489 {
1490 reconstruct_n (key, 1, recons);
1491 }
1492
1493 void Index2Layer::reset()
1494 {
1495 ntotal = 0;
1496 codes.clear ();
1497 }
1498
1499
1500 } // namespace faiss
faiss::IndexIVF
Definition: IndexIVF.h:91
faiss::IndexIVFPQ::precompute_table
void precompute_table()
build precomputed table
Definition: IndexIVFPQ.cpp:364
faiss::IndexIVFPQR::merge_from
void merge_from(IndexIVF &other, idx_t add_id) override
Definition: IndexIVFPQ.cpp:1301
faiss::Index2Layer::transfer_to_IVFPQ
void transfer_to_IVFPQ(IndexIVFPQ &other) const
transfer the flat codes to an IVFPQ index
Definition: IndexIVFPQ.cpp:1466
faiss::Index2Layer::code_size_2
size_t code_size_2
size of the code for the second level
Definition: IndexIVFPQ.h:219
faiss::IndexIVFPQR::reconstruct_from_offset
void reconstruct_from_offset(long list_no, long offset, float *recons) const override
Definition: IndexIVFPQ.cpp:1286
faiss::IndexIVF::search_preassigned
virtual 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
Definition: IndexIVF.cpp:189
faiss::IndexIVFPQR
Definition: IndexIVFPQ.h:156
faiss::ProductQuantizer::decode
void decode(const uint8_t *code, float *x) const
decode a vector from a given code (or n vectors if third argument)
Definition: ProductQuantizer.cpp:335
faiss::IndexIVFPQR::refine_pq
ProductQuantizer refine_pq
3rd level quantizer
Definition: IndexIVFPQ.h:157
faiss::fvec_L2sqr
float fvec_L2sqr(const float *x, const float *y, size_t d)
Squared L2 distance between two vectors.
Definition: utils_simd.cpp:502
faiss::IndexIVFPQ::polysemous_training
PolysemousTraining * polysemous_training
if NULL, use default
Definition: IndexIVFPQ.h:41
faiss::CMax
Definition: Heap.h:48
faiss::Index2Layer::reconstruct_n
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexIVFPQ.cpp:1446
faiss::Index2Layer::code_size
size_t code_size
code_size_1 + code_size_2
Definition: IndexIVFPQ.h:222
faiss::IndexIVFPQStats
Definition: IndexIVFPQ.h:135
faiss::fvecs_maybe_subsample
const float * fvecs_maybe_subsample(size_t d, size_t *n, size_t nmax, const float *x, bool verbose, long seed)
Definition: utils.cpp:1528
faiss::Index2Layer::reset
void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1493
faiss::Index::assign
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
Definition: Index.cpp:35
faiss::IDSelector
Definition: AuxIndexStructures.h:54
faiss::InvertedLists::list_size
virtual size_t list_size(size_t list_no) const =0
get the size of a list
faiss::IndexIVFPQ::decode_multiple
void decode_multiple(size_t n, const long *keys, const uint8_t *xcodes, float *x) const
inverse of encode_multiple
Definition: IndexIVFPQ.cpp:160
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:67
faiss::IndexIVFPQ::do_polysemous_training
bool do_polysemous_training
reorder PQ centroids after training?
Definition: IndexIVFPQ.h:40
faiss::IndexIVFPQ::scan_table_threshold
size_t scan_table_threshold
use table computation or on-the-fly?
Definition: IndexIVFPQ.h:44
faiss::IndexIVFPQR::train_residual
void train_residual(idx_t n, const float *x) override
trains the two product quantizers
Definition: IndexIVFPQ.cpp:1155
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:1179
faiss::ScopeDeleter
Definition: FaissException.h:39
faiss::ProductQuantizer::dsub
size_t dsub
dimensionality of each subvector
Definition: ProductQuantizer.h:31
faiss::ClusteringParameters::seed
int seed
seed for the random number generator
Definition: Clustering.h:35
faiss::IndexIVFPQ::precomputed_table
std::vector< float > precomputed_table
Definition: IndexIVFPQ.h:60
faiss::Index2Layer::q1
Level1Quantizer q1
first level quantizer
Definition: IndexIVFPQ.h:207
faiss::fvec_madd
void fvec_madd(size_t n, const float *a, float bf, const float *b, float *c)
Definition: utils_simd.cpp:589
faiss::IndexIVFPQ::encode_vectors
void encode_vectors(idx_t n, const float *x, const idx_t *list_nos, uint8_t *codes) const override
Definition: IndexIVFPQ.cpp:207
faiss::IndexIVFPQ::polysemous_ht
int polysemous_ht
Hamming thresh for polysemous filtering.
Definition: IndexIVFPQ.h:45
faiss::IndexIVFPQ::get_InvertedListScanner
InvertedListScanner * get_InvertedListScanner(bool store_pairs) const override
get a scanner for this index (store_pairs means ignore labels)
Definition: IndexIVFPQ.cpp:1032
faiss::IndexIVFPQR::reset
void reset() override
removes all elements from the database.
Definition: IndexIVFPQ.cpp:1146
faiss::IndexIVFPQ::add_with_ids
void add_with_ids(idx_t n, const float *x, const long *xids=nullptr) override
Definition: IndexIVFPQ.cpp:183
faiss::IVFSearchParameters
Definition: IndexIVF.h:61
faiss::ProductQuantizer::compute_codes
void compute_codes(const float *x, uint8_t *codes, size_t n) const
same as compute_code for several vectors
Definition: ProductQuantizer.cpp:382
faiss::InvertedLists::get_single_id
virtual idx_t get_single_id(size_t list_no, size_t offset) const
Definition: InvertedLists.cpp:36
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::InvertedLists::code_size
size_t code_size
code size per vector in bytes
Definition: InvertedLists.h:36
faiss::CodeCmp
Definition: IndexIVFPQ.cpp:1076
faiss::InvertedLists::get_single_code
virtual const uint8_t * get_single_code(size_t list_no, size_t offset) const
Definition: InvertedLists.cpp:53
faiss::IndexIVFPQR::refine_codes
std::vector< uint8_t > refine_codes
corresponding codes
Definition: IndexIVFPQ.h:158
faiss::ProductQuantizer::code_size
size_t code_size
byte per indexed vector
Definition: ProductQuantizer.h:33
faiss::IndexIVFPQ::train_residual
void train_residual(idx_t n, const float *x) override
trains the product quantizer
Definition: IndexIVFPQ.cpp:61
faiss::IndexIVFPQ::precomputed_table_max_bytes
static size_t precomputed_table_max_bytes
2G by default, accommodates tables up to PQ32 w/ 65536 centroids
Definition: IndexIVFPQ.h:56
faiss::IndexIVFPQ::reconstruct_from_offset
void reconstruct_from_offset(long list_no, long offset, float *recons) const override
Definition: IndexIVFPQ.cpp:311
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:35
faiss::IndexIVFPQ::encode_multiple
void encode_multiple(size_t n, long *keys, const float *x, uint8_t *codes, bool compute_keys=false) const
Definition: IndexIVFPQ.cpp:150
faiss::InvertedLists::add_entry
virtual size_t add_entry(size_t list_no, idx_t theid, const uint8_t *code)
add one entry to an inverted list
Definition: InvertedLists.cpp:60
faiss::ProductQuantizer::ksub
size_t ksub
number of centroids for each subquantizer
Definition: ProductQuantizer.h:34
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::ProductQuantizer::compute_code
void compute_code(const float *x, uint8_t *code) const
Quantize one vector with the product quantizer.
Definition: ProductQuantizer.cpp:308
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::Index2Layer::train
void train(idx_t n, const float *x) override
Definition: IndexIVFPQ.cpp:1350
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:68
faiss::IndexIVF::reset
void reset() override
removes all elements from the database.
Definition: IndexIVF.cpp:360
faiss::getmillisecs
double getmillisecs()
ms elapsed since some arbitrary epoch
Definition: utils.cpp:70
faiss::PolysemousTraining
optimizes the order of indices in a ProductQuantizer
Definition: PolysemousTraining.h:122
faiss::Index2Layer::add
void add(idx_t n, const float *x) override
Definition: IndexIVFPQ.cpp:1388
faiss::MultiIndexQuantizer
Definition: IndexPQ.h:136
faiss::fvec_norm_L2sqr
float fvec_norm_L2sqr(const float *x, size_t d)
Definition: utils_simd.cpp:516
faiss::Level1Quantizer::train_q1
void train_q1(size_t n, const float *x, bool verbose, MetricType metric_type)
Trains the quantizer and calls train_residual to train sub-quantizers.
Definition: IndexIVF.cpp:57
faiss::ProductQuantizer::cp
ClusteringParameters cp
parameters used during clustering
Definition: ProductQuantizer.h:48
faiss::Index2Layer::search
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
not implemented
Definition: IndexIVFPQ.cpp:1436
faiss::IndexIVFPQ::by_residual
bool by_residual
Encode residual or plain vector?
Definition: IndexIVFPQ.h:36
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::IndexIVFPQ::pq
ProductQuantizer pq
produces the codes
Definition: IndexIVFPQ.h:38
faiss::IndexIVF::invlists
InvertedLists * invlists
Acess to the actual data.
Definition: IndexIVF.h:93
faiss::ProductQuantizer::M
size_t M
number of subquantizers
Definition: ProductQuantizer.h:27
faiss::Index2Layer::code_size_1
size_t code_size_1
size of the code for the first level (ceil(log8(q1.nlist)))
Definition: IndexIVFPQ.h:216
faiss::Index
Definition: Index.h:62
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:221
faiss::fvec_madd_and_argmin
int fvec_madd_and_argmin(size_t n, const float *a, float bf, const float *b, float *c)
Definition: utils_simd.cpp:676
faiss::IndexIVFPQR::search_preassigned
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
Definition: IndexIVFPQ.cpp:1198
faiss::IndexIVFPQR::remove_ids
long remove_ids(const IDSelector &sel) override
Dataset manipulation functions.
Definition: IndexIVFPQ.cpp:1314
faiss::Level1Quantizer::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:33
faiss::ProductQuantizer
Definition: ProductQuantizer.h:24
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:71
faiss::Index::compute_residual
void compute_residual(const float *x, float *residual, idx_t key) const
Definition: Index.cpp:87
faiss::InvertedLists::ScopedIds
Definition: InvertedLists.h:125
faiss::ProductQuantizer::get_centroids
float * get_centroids(size_t m, size_t i)
return the centroids associated with subvector m
Definition: ProductQuantizer.h:58
faiss::Index::reconstruct
virtual void reconstruct(idx_t key, float *recons) const
Definition: Index.cpp:55
faiss::Index2Layer::pq
ProductQuantizer pq
second level quantizer is always a PQ
Definition: IndexIVFPQ.h:210
faiss::IndexIVF::maintain_direct_map
bool maintain_direct_map
map for direct access to the elements. Enables reconstruct().
Definition: IndexIVF.h:102
faiss::PolysemousTraining::optimize_pq_for_hamming
void optimize_pq_for_hamming(ProductQuantizer &pq, size_t n, const float *x) const
Definition: PolysemousTraining.cpp:935
faiss::ClusteringParameters::max_points_per_centroid
int max_points_per_centroid
to limit size of dataset
Definition: Clustering.h:33
faiss::ProductQuantizer::verbose
bool verbose
verbose during training?
Definition: ProductQuantizer.h:35
faiss::IndexIVFPQR::add_with_ids
void add_with_ids(idx_t n, const float *x, const long *xids) override
Definition: IndexIVFPQ.cpp:1175
faiss::Index2Layer::reconstruct
void reconstruct(idx_t key, float *recons) const override
Definition: IndexIVFPQ.cpp:1488
faiss::IndexIVF::merge_from
virtual void merge_from(IndexIVF &other, idx_t add_id)
Definition: IndexIVF.cpp:472
faiss::InvertedLists::ScopedCodes
Definition: InvertedLists.h:144
faiss::Level1Quantizer::nlist
size_t nlist
number of possible key values
Definition: IndexIVF.h:34
faiss::IndexIVF::code_size
size_t code_size
code size per vector in bytes
Definition: IndexIVF.h:96
faiss::IndexIVFPQ::find_duplicates
size_t find_duplicates(idx_t *ids, size_t *lims) const
Definition: IndexIVFPQ.cpp:1089
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:45
faiss::IndexIVFPQR::k_factor
float k_factor
factor between k requested in search and the k requested from the IVFPQ
Definition: IndexIVFPQ.h:161
faiss::CMin
Definition: Heap.h:53
faiss::IndexIVFPQ::use_precomputed_table
int use_precomputed_table
if by_residual, build precompute tables
Definition: IndexIVFPQ.h:55