Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/matthijs/github_faiss/faiss/IndexPQ.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 /* Copyright 2004-present Facebook. All Rights Reserved.
10 Index based on product quantiztion.
11 */
12
13 #include "IndexPQ.h"
14
15
16 #include <cstddef>
17 #include <cstring>
18 #include <cstdio>
19 #include <cmath>
20
21 #include <algorithm>
22
23 #include "FaissAssert.h"
24 #include "hamming.h"
25
26 namespace faiss {
27
28 /*********************************************************
29 * IndexPQ implementation
30 ********************************************************/
31
32
33 IndexPQ::IndexPQ (int d, size_t M, size_t nbits, MetricType metric):
34 Index(d, metric), pq(d, M, nbits)
35 {
36 is_trained = false;
37 do_polysemous_training = false;
38 polysemous_ht = nbits * M + 1;
39 search_type = ST_PQ;
40 encode_signs = false;
41 }
42
43 IndexPQ::IndexPQ ()
44 {
45 metric_type = METRIC_L2;
46 is_trained = false;
47 do_polysemous_training = false;
48 polysemous_ht = pq.nbits * pq.M + 1;
49 search_type = ST_PQ;
50 encode_signs = false;
51 }
52
53
54 void IndexPQ::train (idx_t n, const float *x)
55 {
56 if (!do_polysemous_training) { // standard training
57 pq.train(n, x);
58 } else {
59 idx_t ntrain_perm = polysemous_training.ntrain_permutation;
60
61 if (ntrain_perm > n / 4)
62 ntrain_perm = n / 4;
63 if (verbose) {
64 printf ("PQ training on %ld points, remains %ld points: "
65 "training polysemous on %s\n",
66 n - ntrain_perm, ntrain_perm,
67 ntrain_perm == 0 ? "centroids" : "these");
68 }
69 pq.train(n - ntrain_perm, x);
70
71 polysemous_training.optimize_pq_for_hamming (
72 pq, ntrain_perm, x + (n - ntrain_perm) * d);
73 }
74 is_trained = true;
75 }
76
77
78 void IndexPQ::add (idx_t n, const float *x)
79 {
80 FAISS_THROW_IF_NOT (is_trained);
81 codes.resize ((n + ntotal) * pq.code_size);
82 pq.compute_codes (x, &codes[ntotal * pq.code_size], n);
83 ntotal += n;
84 }
85
86
87
88 void IndexPQ::reset()
89 {
90 codes.clear();
91 ntotal = 0;
92 }
93
94 void IndexPQ::reconstruct_n (idx_t i0, idx_t ni, float *recons) const
95 {
96 FAISS_THROW_IF_NOT (ni == 0 || (i0 >= 0 && i0 + ni <= ntotal));
97 for (idx_t i = 0; i < ni; i++) {
98 const uint8_t * code = &codes[(i0 + i) * pq.code_size];
99 pq.decode (code, recons + i * d);
100 }
101 }
102
103
104 void IndexPQ::reconstruct (idx_t key, float * recons) const
105 {
106 FAISS_THROW_IF_NOT (key >= 0 && key < ntotal);
107 pq.decode (&codes[key * pq.code_size], recons);
108 }
109
110
111
112
113
114
115
116 /*****************************************
117 * IndexPQ polysemous search routines
118 ******************************************/
119
120
121
122
123
124 void IndexPQ::search (idx_t n, const float *x, idx_t k,
125 float *distances, idx_t *labels) const
126 {
127 FAISS_THROW_IF_NOT (is_trained);
128 if (search_type == ST_PQ) { // Simple PQ search
129
130 if (metric_type == METRIC_L2) {
131 float_maxheap_array_t res = {
132 size_t(n), size_t(k), labels, distances };
133 pq.search (x, n, codes.data(), ntotal, &res, true);
134 } else {
135 float_minheap_array_t res = {
136 size_t(n), size_t(k), labels, distances };
137 pq.search_ip (x, n, codes.data(), ntotal, &res, true);
138 }
139 indexPQ_stats.nq += n;
140 indexPQ_stats.ncode += n * ntotal;
141
142 } else if (search_type == ST_polysemous ||
143 search_type == ST_polysemous_generalize) {
144
145 FAISS_THROW_IF_NOT (metric_type == METRIC_L2);
146
147 search_core_polysemous (n, x, k, distances, labels);
148
149 } else { // code-to-code distances
150
151 uint8_t * q_codes = new uint8_t [n * pq.code_size];
152 ScopeDeleter<uint8_t> del (q_codes);
153
154
155 if (!encode_signs) {
156 pq.compute_codes (x, q_codes, n);
157 } else {
158 FAISS_THROW_IF_NOT (d == pq.nbits * pq.M);
159 memset (q_codes, 0, n * pq.code_size);
160 for (size_t i = 0; i < n; i++) {
161 const float *xi = x + i * d;
162 uint8_t *code = q_codes + i * pq.code_size;
163 for (int j = 0; j < d; j++)
164 if (xi[j] > 0) code [j>>3] |= 1 << (j & 7);
165 }
166 }
167
168 if (search_type == ST_SDC) {
169
170 float_maxheap_array_t res = {
171 size_t(n), size_t(k), labels, distances};
172
173 pq.search_sdc (q_codes, n, codes.data(), ntotal, &res, true);
174
175 } else {
176 int * idistances = new int [n * k];
177 ScopeDeleter<int> del (idistances);
178
179 int_maxheap_array_t res = {
180 size_t (n), size_t (k), labels, idistances};
181
182 if (search_type == ST_HE) {
183
184 hammings_knn (&res, q_codes, codes.data(),
185 ntotal, pq.code_size, true);
186
187 } else if (search_type == ST_generalized_HE) {
188
189 generalized_hammings_knn (&res, q_codes, codes.data(),
190 ntotal, pq.code_size, true);
191 }
192
193 // convert distances to floats
194 for (int i = 0; i < k * n; i++)
195 distances[i] = idistances[i];
196
197 }
198
199
200 indexPQ_stats.nq += n;
201 indexPQ_stats.ncode += n * ntotal;
202 }
203 }
204
205
206
207
208
209 void IndexPQStats::reset()
210 {
211 nq = ncode = n_hamming_pass = 0;
212 }
213
214 IndexPQStats indexPQ_stats;
215
216
217 template <class HammingComputer>
218 static size_t polysemous_inner_loop (
219 const IndexPQ & index,
220 const float *dis_table_qi, const uint8_t *q_code,
221 size_t k, float *heap_dis, long *heap_ids)
222 {
223
224 int M = index.pq.M;
225 int code_size = index.pq.code_size;
226 int ksub = index.pq.ksub;
227 size_t ntotal = index.ntotal;
228 int ht = index.polysemous_ht;
229
230 const uint8_t *b_code = index.codes.data();
231
232 size_t n_pass_i = 0;
233
234 HammingComputer hc (q_code, code_size);
235
236 for (long bi = 0; bi < ntotal; bi++) {
237 int hd = hc.hamming (b_code);
238
239 if (hd < ht) {
240 n_pass_i ++;
241
242 float dis = 0;
243 const float * dis_table = dis_table_qi;
244 for (int m = 0; m < M; m++) {
245 dis += dis_table [b_code[m]];
246 dis_table += ksub;
247 }
248
249 if (dis < heap_dis[0]) {
250 maxheap_pop (k, heap_dis, heap_ids);
251 maxheap_push (k, heap_dis, heap_ids, dis, bi);
252 }
253 }
254 b_code += code_size;
255 }
256 return n_pass_i;
257 }
258
259
260 void IndexPQ::search_core_polysemous (idx_t n, const float *x, idx_t k,
261 float *distances, idx_t *labels) const
262 {
263 FAISS_THROW_IF_NOT (pq.byte_per_idx == 1);
264
265 // PQ distance tables
266 float * dis_tables = new float [n * pq.ksub * pq.M];
267 ScopeDeleter<float> del (dis_tables);
268 pq.compute_distance_tables (n, x, dis_tables);
269
270 // Hamming embedding queries
271 uint8_t * q_codes = new uint8_t [n * pq.code_size];
272 ScopeDeleter<uint8_t> del2 (q_codes);
273
274 if (false) {
275 pq.compute_codes (x, q_codes, n);
276 } else {
277 #pragma omp parallel for
278 for (idx_t qi = 0; qi < n; qi++) {
279 pq.compute_code_from_distance_table
280 (dis_tables + qi * pq.M * pq.ksub,
281 q_codes + qi * pq.code_size);
282 }
283 }
284
285 size_t n_pass = 0;
286
287 #pragma omp parallel for reduction (+: n_pass)
288 for (idx_t qi = 0; qi < n; qi++) {
289 const uint8_t * q_code = q_codes + qi * pq.code_size;
290
291 const float * dis_table_qi = dis_tables + qi * pq.M * pq.ksub;
292
293 long * heap_ids = labels + qi * k;
294 float *heap_dis = distances + qi * k;
295 maxheap_heapify (k, heap_dis, heap_ids);
296
297 if (search_type == ST_polysemous) {
298
299 switch (pq.code_size) {
300 case 4:
301 n_pass += polysemous_inner_loop<HammingComputer4>
302 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
303 break;
304 case 8:
305 n_pass += polysemous_inner_loop<HammingComputer8>
306 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
307 break;
308 case 16:
309 n_pass += polysemous_inner_loop<HammingComputer16>
310 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
311 break;
312 case 32:
313 n_pass += polysemous_inner_loop<HammingComputer32>
314 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
315 break;
316 case 20:
317 n_pass += polysemous_inner_loop<HammingComputer20>
318 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
319 break;
320 default:
321 if (pq.code_size % 8 == 0) {
322 n_pass += polysemous_inner_loop<HammingComputerM8>
323 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
324 } else if (pq.code_size % 4 == 0) {
325 n_pass += polysemous_inner_loop<HammingComputerM4>
326 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
327 } else {
328 FAISS_THROW_FMT(
329 "code size %zd not supported for polysemous",
330 pq.code_size);
331 }
332 break;
333 }
334 } else {
335 switch (pq.code_size) {
336 case 8:
337 n_pass += polysemous_inner_loop<GenHammingComputer8>
338 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
339 break;
340 case 16:
341 n_pass += polysemous_inner_loop<GenHammingComputer16>
342 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
343 break;
344 case 32:
345 n_pass += polysemous_inner_loop<GenHammingComputer32>
346 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
347 break;
348 default:
349 if (pq.code_size % 8 == 0) {
350 n_pass += polysemous_inner_loop<GenHammingComputerM8>
351 (*this, dis_table_qi, q_code, k, heap_dis, heap_ids);
352 } else {
353 FAISS_THROW_FMT(
354 "code size %zd not supported for polysemous",
355 pq.code_size);
356 }
357 break;
358 }
359 }
360 maxheap_reorder (k, heap_dis, heap_ids);
361 }
362
363 indexPQ_stats.nq += n;
364 indexPQ_stats.ncode += n * ntotal;
365 indexPQ_stats.n_hamming_pass += n_pass;
366
367
368 }
369
370
371
372
373 /*****************************************
374 * Stats of IndexPQ codes
375 ******************************************/
376
377
378
379
380 void IndexPQ::hamming_distance_table (idx_t n, const float *x,
381 int32_t *dis) const
382 {
383 uint8_t * q_codes = new uint8_t [n * pq.code_size];
384 ScopeDeleter<uint8_t> del (q_codes);
385
386 pq.compute_codes (x, q_codes, n);
387
388 hammings (q_codes, codes.data(), n, ntotal, pq.code_size, dis);
389 }
390
391
392 void IndexPQ::hamming_distance_histogram (idx_t n, const float *x,
393 idx_t nb, const float *xb,
394 long *hist)
395 {
396 FAISS_THROW_IF_NOT (metric_type == METRIC_L2);
397 FAISS_THROW_IF_NOT (pq.code_size % 8 == 0);
398 FAISS_THROW_IF_NOT (pq.byte_per_idx == 1);
399
400 // Hamming embedding queries
401 uint8_t * q_codes = new uint8_t [n * pq.code_size];
402 ScopeDeleter <uint8_t> del (q_codes);
403 pq.compute_codes (x, q_codes, n);
404
405 uint8_t * b_codes ;
406 ScopeDeleter <uint8_t> del_b_codes;
407
408 if (xb) {
409 b_codes = new uint8_t [nb * pq.code_size];
410 del_b_codes.set (b_codes);
411 pq.compute_codes (xb, b_codes, nb);
412 } else {
413 nb = ntotal;
414 b_codes = codes.data();
415 }
416 int nbits = pq.M * pq.nbits;
417 memset (hist, 0, sizeof(*hist) * (nbits + 1));
418 size_t bs = 256;
419
420 #pragma omp parallel
421 {
422 std::vector<long> histi (nbits + 1);
423 hamdis_t *distances = new hamdis_t [nb * bs];
424 ScopeDeleter<hamdis_t> del (distances);
425 #pragma omp for
426 for (size_t q0 = 0; q0 < n; q0 += bs) {
427 // printf ("dis stats: %ld/%ld\n", q0, n);
428 size_t q1 = q0 + bs;
429 if (q1 > n) q1 = n;
430
431 hammings (q_codes + q0 * pq.code_size, b_codes,
432 q1 - q0, nb,
433 pq.code_size, distances);
434
435 for (size_t i = 0; i < nb * (q1 - q0); i++)
436 histi [distances [i]]++;
437 }
438 #pragma omp critical
439 {
440 for (int i = 0; i <= nbits; i++)
441 hist[i] += histi[i];
442 }
443 }
444
445 }
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466 /*****************************************
467 * MultiIndexQuantizer
468 ******************************************/
469
470
471
472 template <typename T>
473 struct ArgSort {
474 const T * x;
475 bool operator() (size_t i, size_t j) {
476 return x[i] < x[j];
477 }
478 };
479
480
481 /** Array that maintains a permutation of its elements so that the
482 * array's elements are sorted
483 */
484 template <typename T>
485 struct SortedArray {
486 const T * x;
487 int N;
488 std::vector<int> perm;
489
490 explicit SortedArray (int N) {
491 this->N = N;
492 perm.resize (N);
493 }
494
495 void init (const T*x) {
496 this->x = x;
497 for (int n = 0; n < N; n++)
498 perm[n] = n;
499 ArgSort<T> cmp = {x };
500 std::sort (perm.begin(), perm.end(), cmp);
501 }
502
503 // get smallest value
504 T get_0 () {
505 return x[perm[0]];
506 }
507
508 // get delta between n-smallest and n-1 -smallest
509 T get_diff (int n) {
510 return x[perm[n]] - x[perm[n - 1]];
511 }
512
513 // remap orders counted from smallest to indices in array
514 int get_ord (int n) {
515 return perm[n];
516 }
517 };
518
519
520
521 /** Array has n values. Sort the k first ones and copy the other ones
522 * into elements k..n-1
523 */
524 template <class C>
525 void partial_sort (int k, int n,
526 const typename C::T * vals, typename C::TI * perm) {
527 // insert first k elts in heap
528 for (int i = 1; i < k; i++) {
529 indirect_heap_push<C> (i + 1, vals, perm, perm[i]);
530 }
531
532 // insert next n - k elts in heap
533 for (int i = k; i < n; i++) {
534 typename C::TI id = perm[i];
535 typename C::TI top = perm[0];
536
537 if (C::cmp(vals[top], vals[id])) {
538 indirect_heap_pop<C> (k, vals, perm);
539 indirect_heap_push<C> (k, vals, perm, id);
540 perm[i] = top;
541 } else {
542 // nothing, elt at i is good where it is.
543 }
544 }
545
546 // order the k first elements in heap
547 for (int i = k - 1; i > 0; i--) {
548 typename C::TI top = perm[0];
549 indirect_heap_pop<C> (i + 1, vals, perm);
550 perm[i] = top;
551 }
552 }
553
554 /** same as SortedArray, but only the k first elements are sorted */
555 template <typename T>
556 struct SemiSortedArray {
557 const T * x;
558 int N;
559
560 // type of the heap: CMax = sort ascending
561 typedef CMax<T, int> HC;
562 std::vector<int> perm;
563
564 int k; // k elements are sorted
565
566 int initial_k, k_factor;
567
568 explicit SemiSortedArray (int N) {
569 this->N = N;
570 perm.resize (N);
571 perm.resize (N);
572 initial_k = 3;
573 k_factor = 4;
574 }
575
576 void init (const T*x) {
577 this->x = x;
578 for (int n = 0; n < N; n++)
579 perm[n] = n;
580 k = 0;
581 grow (initial_k);
582 }
583
584 /// grow the sorted part of the array to size next_k
585 void grow (int next_k) {
586 if (next_k < N) {
587 partial_sort<HC> (next_k - k, N - k, x, &perm[k]);
588 k = next_k;
589 } else { // full sort of remainder of array
590 ArgSort<T> cmp = {x };
591 std::sort (perm.begin() + k, perm.end(), cmp);
592 k = N;
593 }
594 }
595
596 // get smallest value
597 T get_0 () {
598 return x[perm[0]];
599 }
600
601 // get delta between n-smallest and n-1 -smallest
602 T get_diff (int n) {
603 if (n >= k) {
604 // want to keep powers of 2 - 1
605 int next_k = (k + 1) * k_factor - 1;
606 grow (next_k);
607 }
608 return x[perm[n]] - x[perm[n - 1]];
609 }
610
611 // remap orders counted from smallest to indices in array
612 int get_ord (int n) {
613 assert (n < k);
614 return perm[n];
615 }
616 };
617
618
619
620 /*****************************************
621 * Find the k smallest sums of M terms, where each term is taken in a
622 * table x of n values.
623 *
624 * A combination of terms is encoded as a scalar 0 <= t < n^M. The
625 * combination t0 ... t(M-1) that correspond to the sum
626 *
627 * sum = x[0, t0] + x[1, t1] + .... + x[M-1, t(M-1)]
628 *
629 * is encoded as
630 *
631 * t = t0 + t1 * n + t2 * n^2 + ... + t(M-1) * n^(M-1)
632 *
633 * MinSumK is an object rather than a function, so that storage can be
634 * re-used over several computations with the same sizes. use_seen is
635 * good when there may be ties in the x array and it is a concern if
636 * occasionally several t's are returned.
637 *
638 * @param x size M * n, values to add up
639 * @parms k nb of results to retrieve
640 * @param M nb of terms
641 * @param n nb of distinct values
642 * @param sums output, size k, sorted
643 * @prarm terms output, size k, with encoding as above
644 *
645 ******************************************/
646 template <typename T, class SSA, bool use_seen>
647 struct MinSumK {
648 int K; ///< nb of sums to return
649 int M; ///< nb of elements to sum up
650 int N; ///< nb of possible elements for each of the M terms
651
652 /** the heap.
653 * We use a heap to maintain a queue of sums, with the associated
654 * terms involved in the sum.
655 */
656 typedef CMin<T, long> HC;
657 size_t heap_capacity, heap_size;
658 T *bh_val;
659 long *bh_ids;
660
661 std::vector <SSA> ssx;
662 std::vector <long> weights;
663
664 // all results get pushed several times. When there are ties, they
665 // are popped interleaved with others, so it is not easy to
666 // identify them. Therefore, this bit array just marks elements
667 // that were seen before.
668 std::vector <uint8_t> seen;
669
670 MinSumK (int K, int M, int N): K(K), M(M), N(N) {
671 heap_capacity = K * M;
672 // we'll do k steps, each step pushes at most M vals
673 bh_val = new T[heap_capacity];
674 bh_ids = new long[heap_capacity];
675
676 weights.push_back (1);
677 for (int m = 1; m < M; m++)
678 weights.push_back(weights[m - 1] * N);
679
680 if (use_seen) {
681 long n_ids = weights.back() * N;
682 seen.resize ((n_ids + 7) / 8);
683 }
684
685 for (int m = 0; m < M; m++)
686 ssx.push_back (SSA(N));
687
688 }
689
690 bool is_seen (long i) {
691 return (seen[i >> 3] >> (i & 7)) & 1;
692 }
693
694 void mark_seen (long i) {
695 if (use_seen)
696 seen [i >> 3] |= 1 << (i & 7);
697 }
698
699 void run (const T *x, T * sums, long * terms) {
700 heap_size = 0;
701
702 for (int m = 0; m < M; m++)
703 ssx[m].init(x + N * m);
704
705 { // intial result: take min for all elements
706 T sum = 0;
707 terms[0] = 0;
708 mark_seen (0);
709 for (int m = 0; m < M; m++) {
710 sum += ssx[m].get_0();
711 }
712 sums[0] = sum;
713 for (int m = 0; m < M; m++) {
714 heap_push<HC> (++heap_size, bh_val, bh_ids,
715 sum + ssx[m].get_diff(1),
716 weights[m]);
717 }
718 }
719
720 for (int k = 1; k < K; k++) {
721 // pop smallest value from heap
722 if (use_seen) {// skip already seen elements
723 while (is_seen (bh_ids[0])) {
724 assert (heap_size > 0);
725 heap_pop<HC> (heap_size--, bh_val, bh_ids);
726 }
727 }
728 assert (heap_size > 0);
729
730 T sum = sums[k] = bh_val[0];
731 long ti = terms[k] = bh_ids[0];
732
733 if (use_seen) {
734 mark_seen (ti);
735 heap_pop<HC> (heap_size--, bh_val, bh_ids);
736 } else {
737 do {
738 heap_pop<HC> (heap_size--, bh_val, bh_ids);
739 } while (heap_size > 0 && bh_ids[0] == ti);
740 }
741
742 // enqueue followers
743 long ii = ti;
744 for (int m = 0; m < M; m++) {
745 long n = ii % N;
746 ii /= N;
747 if (n + 1 >= N) continue;
748
749 enqueue_follower (ti, m, n, sum);
750 }
751 }
752
753 /*
754 for (int k = 0; k < K; k++)
755 for (int l = k + 1; l < K; l++)
756 assert (terms[k] != terms[l]);
757 */
758
759 // convert indices by applying permutation
760 for (int k = 0; k < K; k++) {
761 long ii = terms[k];
762 if (use_seen) {
763 // clear seen for reuse at next loop
764 seen[ii >> 3] = 0;
765 }
766 long ti = 0;
767 for (int m = 0; m < M; m++) {
768 long n = ii % N;
769 ti += weights[m] * ssx[m].get_ord(n);
770 ii /= N;
771 }
772 terms[k] = ti;
773 }
774 }
775
776
777 void enqueue_follower (long ti, int m, int n, T sum) {
778 T next_sum = sum + ssx[m].get_diff(n + 1);
779 long next_ti = ti + weights[m];
780 heap_push<HC> (++heap_size, bh_val, bh_ids, next_sum, next_ti);
781 }
782
783
784 ~MinSumK () {
785 delete [] bh_ids;
786 delete [] bh_val;
787 }
788 };
789
790
791
792
793 MultiIndexQuantizer::MultiIndexQuantizer (int d,
794 size_t M,
795 size_t nbits):
796 Index(d, METRIC_L2), pq(d, M, nbits)
797 {
798 is_trained = false;
799 pq.verbose = verbose;
800 }
801
802
803
804 void MultiIndexQuantizer::train(idx_t n, const float *x)
805 {
806 pq.verbose = verbose;
807 pq.train (n, x);
808 is_trained = true;
809 // count virtual elements in index
810 ntotal = 1;
811 for (int m = 0; m < pq.M; m++)
812 ntotal *= pq.ksub;
813 }
814
815
816 void MultiIndexQuantizer::search (idx_t n, const float *x, idx_t k,
817 float *distances, idx_t *labels) const {
818 if (n == 0) return;
819
820 float * dis_tables = new float [n * pq.ksub * pq.M];
821 ScopeDeleter<float> del (dis_tables);
822
823 pq.compute_distance_tables (n, x, dis_tables);
824
825 if (k == 1) {
826 // simple version that just finds the min in each table
827
828 #pragma omp parallel for
829 for (int i = 0; i < n; i++) {
830 const float * dis_table = dis_tables + i * pq.ksub * pq.M;
831 float dis = 0;
832 idx_t label = 0;
833
834 for (int s = 0; s < pq.M; s++) {
835 float vmin = HUGE_VALF;
836 idx_t lmin = -1;
837
838 for (idx_t j = 0; j < pq.ksub; j++) {
839 if (dis_table[j] < vmin) {
840 vmin = dis_table[j];
841 lmin = j;
842 }
843 }
844 dis += vmin;
845 label |= lmin << (s * pq.nbits);
846 dis_table += pq.ksub;
847 }
848
849 distances [i] = dis;
850 labels [i] = label;
851 }
852
853
854 } else {
855
856 #pragma omp parallel if(n > 1)
857 {
858 MinSumK <float, SemiSortedArray<float>, false>
859 msk(k, pq.M, pq.ksub);
860 #pragma omp for
861 for (int i = 0; i < n; i++) {
862 msk.run (dis_tables + i * pq.ksub * pq.M,
863 distances + i * k, labels + i * k);
864
865 }
866 }
867 }
868
869 }
870
871
872 void MultiIndexQuantizer::reconstruct (idx_t key, float * recons) const
873 {
874 if (pq.byte_per_idx == 1) {
875 uint8_t code[pq.M];
876 long jj = key;
877 for (int m = 0; m < pq.M; m++) {
878 long n = jj % pq.ksub;
879 jj /= pq.ksub;
880 code[m] = n;
881 }
882 pq.decode (code, recons);
883 } else if (pq.byte_per_idx == 2) {
884 uint16_t code[pq.M];
885 long jj = key;
886 for (int m = 0; m < pq.M; m++) {
887 long n = jj % pq.ksub;
888 jj /= pq.ksub;
889 code[m] = n;
890 }
891 pq.decode ((uint8_t*)code, recons);
892 } else FAISS_THROW_MSG( "only 1 or 2 bytes per index supported");
893 }
894
895 void MultiIndexQuantizer::add(idx_t /*n*/, const float* /*x*/) {
896 FAISS_THROW_MSG(
897 "This index has virtual elements, "
898 "it does not support add");
899 }
900
901 void MultiIndexQuantizer::reset ()
902 {
903 FAISS_THROW_MSG ( "This index has virtual elements, "
904 "it does not support reset");
905 }
906
907
908
909
910 } // END namespace faiss
faiss::MinSumK::M
int M
nb of elements to sum up
Definition: IndexPQ.cpp:649
faiss::IndexPQ::codes
std::vector< uint8_t > codes
Codes. Size ntotal * pq.code_size.
Definition: IndexPQ.h:34
faiss::ProductQuantizer::nbits
size_t nbits
number of bits per quantization index
Definition: ProductQuantizer.h:29
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:338
faiss::IndexPQ::ST_HE
Hamming distance on codes.
Definition: IndexPQ.h:77
faiss::IndexPQ::do_polysemous_training
bool do_polysemous_training
false = standard PQ
Definition: IndexPQ.h:69
faiss::IndexPQ::train
void train(idx_t n, const float *x) override
Definition: IndexPQ.cpp:54
faiss::ProductQuantizer::byte_per_idx
size_t byte_per_idx
nb bytes per code component (1 or 2)
Definition: ProductQuantizer.h:33
faiss::MultiIndexQuantizer::reset
void reset() override
removes all elements from the database.
Definition: IndexPQ.cpp:901
faiss::CMax
Definition: Heap.h:48
faiss::partial_sort
void partial_sort(int k, int n, const typename C::T *vals, typename C::TI *perm)
Definition: IndexPQ.cpp:525
faiss::MultiIndexQuantizer::train
void train(idx_t n, const float *x) override
Definition: IndexPQ.cpp:804
faiss::MinSumK
Definition: IndexPQ.cpp:647
faiss::MinSumK::HC
CMin< T, long > HC
Definition: IndexPQ.cpp:656
faiss::ArgSort
Definition: IndexPQ.cpp:473
faiss::SemiSortedArray::grow
void grow(int next_k)
grow the sorted part of the array to size next_k
Definition: IndexPQ.cpp:585
faiss::SemiSortedArray
Definition: IndexPQ.cpp:556
faiss::ScopeDeleter
Definition: FaissException.h:39
faiss::ProductQuantizer::compute_distance_tables
void compute_distance_tables(size_t nx, const float *x, float *dis_tables) const
Definition: ProductQuantizer.cpp:439
faiss::SortedArray
Definition: IndexPQ.cpp:485
faiss::generalized_hammings_knn
void generalized_hammings_knn(int_maxheap_array_t *ha, const uint8_t *a, const uint8_t *b, size_t nb, size_t code_size, int ordered)
Definition: hamming.cpp:639
faiss::ProductQuantizer::compute_code_from_distance_table
void compute_code_from_distance_table(const float *tab, uint8_t *code) const
Definition: ProductQuantizer.cpp:363
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:385
faiss::Index::d
int d
vector dimension
Definition: Index.h:64
faiss::IndexPQ::hamming_distance_histogram
void hamming_distance_histogram(idx_t n, const float *x, idx_t nb, const float *xb, long *dist_histogram)
Definition: IndexPQ.cpp:392
faiss::ProductQuantizer::search
void search(const float *x, size_t nx, const uint8_t *codes, const size_t ncodes, float_maxheap_array_t *res, bool init_finalize_heap=true) const
Definition: ProductQuantizer.cpp:540
faiss::ProductQuantizer::code_size
size_t code_size
byte per indexed vector
Definition: ProductQuantizer.h:34
faiss::IndexPQ::ST_polysemous_generalize
Filter on generalized Hamming.
Definition: IndexPQ.h:81
faiss::ProductQuantizer::ksub
size_t ksub
number of centroids for each subquantizer
Definition: ProductQuantizer.h:35
faiss::HeapArray
Definition: Heap.h:350
faiss::ProductQuantizer::search_ip
void search_ip(const float *x, size_t nx, const uint8_t *codes, const size_t ncodes, float_minheap_array_t *res, bool init_finalize_heap=true) const
Definition: ProductQuantizer.cpp:565
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::hammings_knn
void hammings_knn(int_maxheap_array_t *ha, const uint8_t *a, const uint8_t *b, size_t nb, size_t ncodes, int order)
Definition: hamming.cpp:474
faiss::IndexPQ::pq
ProductQuantizer pq
The product quantizer used to encode the vectors.
Definition: IndexPQ.h:31
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:65
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:66
faiss::IndexPQ::add
void add(idx_t n, const float *x) override
Definition: IndexPQ.cpp:78
faiss::IndexPQ::ST_generalized_HE
nb of same codes
Definition: IndexPQ.h:78
faiss::MinSumK::K
int K
nb of sums to return
Definition: IndexPQ.cpp:648
faiss::IndexPQ::hamming_distance_table
void hamming_distance_table(idx_t n, const float *x, int32_t *dis) const
Definition: IndexPQ.cpp:380
faiss::MultiIndexQuantizer::search
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexPQ.cpp:816
faiss::IndexPQ::reconstruct
void reconstruct(idx_t key, float *recons) const override
Definition: IndexPQ.cpp:104
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:72
faiss::ProductQuantizer::M
size_t M
number of subquantizers
Definition: ProductQuantizer.h:28
faiss::MinSumK::N
int N
nb of possible elements for each of the M terms
Definition: IndexPQ.cpp:650
faiss::Index
Definition: Index.h:60
faiss::IndexPQ::reconstruct_n
void reconstruct_n(idx_t i0, idx_t ni, float *recons) const override
Definition: IndexPQ.cpp:94
faiss::IndexPQ::ST_PQ
asymmetric product quantizer (default)
Definition: IndexPQ.h:76
faiss::MultiIndexQuantizer::reconstruct
void reconstruct(idx_t key, float *recons) const override
Definition: IndexPQ.cpp:872
faiss::IndexPQ::ST_polysemous
HE filter (using ht) + PQ combination.
Definition: IndexPQ.h:80
faiss::MultiIndexQuantizer::add
void add(idx_t n, const float *x) override
add and reset will crash at runtime
Definition: IndexPQ.cpp:895
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:69
faiss::IndexPQ::reset
void reset() override
removes all elements from the database.
Definition: IndexPQ.cpp:88
faiss::PolysemousTraining::optimize_pq_for_hamming
void optimize_pq_for_hamming(ProductQuantizer &pq, size_t n, const float *x) const
Definition: PolysemousTraining.cpp:933
faiss::ProductQuantizer::verbose
bool verbose
verbose during training?
Definition: ProductQuantizer.h:36
faiss::IndexPQ::search
void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const override
Definition: IndexPQ.cpp:124
faiss::IndexPQ::ST_SDC
symmetric product quantizer (SDC)
Definition: IndexPQ.h:79
faiss::IndexPQ::polysemous_ht
int polysemous_ht
Hamming threshold used for polysemy.
Definition: IndexPQ.h:91
faiss::IndexPQ::polysemous_training
PolysemousTraining polysemous_training
parameters used for the polysemous training
Definition: IndexPQ.h:72
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:43
faiss::CMin
Definition: Heap.h:53