Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PolysemousTraining.cpp
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 #include "PolysemousTraining.h"
11
12 #include <cstdlib>
13 #include <cmath>
14 #include <cstring>
15
16 #include <algorithm>
17
18 #include "utils.h"
19 #include "hamming.h"
20
21 #include "FaissAssert.h"
22
23 /*****************************************
24 * Mixed PQ / Hamming
25 ******************************************/
26
27 namespace faiss {
28
29 /****************************************************
30 * Optimization code
31 ****************************************************/
32
33
34 SimulatedAnnealingParameters::SimulatedAnnealingParameters ()
35 {
36 // set some reasonable defaults for the optimization
37 init_temperature = 0.7;
38 temperature_decay = pow (0.9, 1/500.);
39 // reduce by a factor 0.9 every 500 it
40 n_iter = 500000;
41 n_redo = 2;
42 seed = 123;
43 verbose = 0;
44 only_bit_flips = false;
45 init_random = false;
46 }
47
48 // what would the cost update be if iw and jw were swapped?
49 // default implementation just computes both and computes the difference
50 double PermutationObjective::cost_update (
51 const int *perm, int iw, int jw) const
52 {
53 double orig_cost = compute_cost (perm);
54
55 std::vector<int> perm2 (n);
56 for (int i = 0; i < n; i++)
57 perm2[i] = perm[i];
58 perm2[iw] = perm[jw];
59 perm2[jw] = perm[iw];
60
61 double new_cost = compute_cost (perm2.data());
62 return new_cost - orig_cost;
63 }
64
65
66
67
68 SimulatedAnnealingOptimizer::SimulatedAnnealingOptimizer (
69 PermutationObjective *obj,
70 const SimulatedAnnealingParameters &p):
71 SimulatedAnnealingParameters (p),
72 obj (obj),
73 n(obj->n),
74 logfile (nullptr)
75 {
76 rnd = new RandomGenerator (p.seed);
77 FAISS_ASSERT (n < 100000 && n >=0 );
78 }
79
80 SimulatedAnnealingOptimizer::~SimulatedAnnealingOptimizer ()
81 {
82 delete rnd;
83 }
84
85 // run the optimization and return the best result in best_perm
86 double SimulatedAnnealingOptimizer::run_optimization (int * best_perm)
87 {
88 double min_cost = 1e30;
89
90 // just do a few runs of the annealing and keep the lowest output cost
91 for (int it = 0; it < n_redo; it++) {
92 std::vector<int> perm(n);
93 for (int i = 0; i < n; i++)
94 perm[i] = i;
95 if (init_random) {
96 for (int i = 0; i < n; i++) {
97 int j = i + rnd->rand_int (n - i);
98 std::swap (perm[i], perm[j]);
99 }
100 }
101 float cost = optimize (perm.data());
102 if (logfile) fprintf (logfile, "\n");
103 if(verbose > 1) {
104 printf (" optimization run %d: cost=%g %s\n",
105 it, cost, cost < min_cost ? "keep" : "");
106 }
107 if (cost < min_cost) {
108 memcpy (best_perm, perm.data(), sizeof(perm[0]) * n);
109 min_cost = cost;
110 }
111 }
112 return min_cost;
113 }
114
115 // perform the optimization loop, starting from and modifying
116 // permutation in-place
117 double SimulatedAnnealingOptimizer::optimize (int *perm)
118 {
119 double cost = init_cost = obj->compute_cost (perm);
120 int log2n = 0;
121 while (!(n <= (1 << log2n))) log2n++;
122 double temperature = init_temperature;
123 int n_swap = 0, n_hot = 0;
124 for (int it = 0; it < n_iter; it++) {
125 temperature = temperature * temperature_decay;
126 int iw, jw;
127 if (only_bit_flips) {
128 iw = rnd->rand_int (n);
129 jw = iw ^ (1 << rnd->rand_int (log2n));
130 } else {
131 iw = rnd->rand_int (n);
132 jw = rnd->rand_int (n - 1);
133 if (jw == iw) jw++;
134 }
135 double delta_cost = obj->cost_update (perm, iw, jw);
136 if (delta_cost < 0 || rnd->rand_float () < temperature) {
137 std::swap (perm[iw], perm[jw]);
138 cost += delta_cost;
139 n_swap++;
140 if (delta_cost >= 0) n_hot++;
141 }
142 if (verbose > 2 || (verbose > 1 && it % 10000 == 0)) {
143 printf (" iteration %d cost %g temp %g n_swap %d "
144 "(%d hot) \r",
145 it, cost, temperature, n_swap, n_hot);
146 fflush(stdout);
147 }
148 if (logfile) {
149 fprintf (logfile, "%d %g %g %d %d\n",
150 it, cost, temperature, n_swap, n_hot);
151 }
152 }
153 if (verbose > 1) printf("\n");
154 return cost;
155 }
156
157
158
159
160
161 /****************************************************
162 * Cost functions: ReproduceDistanceTable
163 ****************************************************/
164
165
166
167
168
169
170 static inline int hamming_dis (uint64_t a, uint64_t b)
171 {
172 return __builtin_popcountl (a ^ b);
173 }
174
175 namespace {
176
177 /// optimize permutation to reproduce a distance table with Hamming distances
178 struct ReproduceWithHammingObjective : PermutationObjective {
179 int nbits;
180 double dis_weight_factor;
181
182 static double sqr (double x) { return x * x; }
183
184
185 // weihgting of distances: it is more important to reproduce small
186 // distances well
187 double dis_weight (double x) const
188 {
189 return exp (-dis_weight_factor * x);
190 }
191
192 std::vector<double> target_dis; // wanted distances (size n^2)
193 std::vector<double> weights; // weights for each distance (size n^2)
194
195 // cost = quadratic difference between actual distance and Hamming distance
196 virtual double compute_cost (const int *perm) const
197 {
198 double cost = 0;
199 for (int i = 0; i < n; i++) {
200 for (int j = 0; j < n; j++) {
201 double wanted = target_dis [i * n + j];
202 double w = weights [i * n + j];
203 double actual = hamming_dis (perm[i], perm[j]);
204 cost += w * sqr (wanted - actual);
205 }
206 }
207 return cost;
208 }
209
210
211 // what would the cost update be if iw and jw were swapped?
212 // computed in O(n) instead of O(n^2) for the full re-computation
213 double cost_update (const int *perm, int iw, int jw) const
214 {
215 double delta_cost = 0;
216
217 for (int i = 0; i < n; i++) {
218 if (i == iw) {
219 for (int j = 0; j < n; j++) {
220 double wanted = target_dis [i * n + j],
221 w = weights [i * n + j];
222 double actual = hamming_dis (perm[i], perm[j]);
223 delta_cost -= w * sqr (wanted - actual);
224 double new_actual = hamming_dis (
225 perm[jw],
226 perm[j == iw ? jw : j == jw ? iw : j]);
227 delta_cost += w * sqr (wanted - new_actual);
228 }
229 } else if (i == jw) {
230 for (int j = 0; j < n; j++) {
231 double wanted = target_dis [i * n + j],
232 w = weights [i * n + j];
233 double actual = hamming_dis (perm[i], perm[j]);
234 delta_cost -= w * sqr (wanted - actual);
235 double new_actual = hamming_dis (
236 perm[iw],
237 perm[j == iw ? jw : j == jw ? iw : j]);
238 delta_cost += w * sqr (wanted - new_actual);
239 }
240 } else {
241 int j = iw;
242 {
243 double wanted = target_dis [i * n + j],
244 w = weights [i * n + j];
245 double actual = hamming_dis (perm[i], perm[j]);
246 delta_cost -= w * sqr (wanted - actual);
247 double new_actual = hamming_dis (perm[i], perm[jw]);
248 delta_cost += w * sqr (wanted - new_actual);
249 }
250 j = jw;
251 {
252 double wanted = target_dis [i * n + j],
253 w = weights [i * n + j];
254 double actual = hamming_dis (perm[i], perm[j]);
255 delta_cost -= w * sqr (wanted - actual);
256 double new_actual = hamming_dis (perm[i], perm[iw]);
257 delta_cost += w * sqr (wanted - new_actual);
258 }
259 }
260 }
261
262 return delta_cost;
263 }
264
265
266
267 ReproduceWithHammingObjective (
268 int nbits,
269 const std::vector<double> & dis_table,
270 double dis_weight_factor):
271 nbits (nbits), dis_weight_factor (dis_weight_factor)
272 {
273 n = 1 << nbits;
274 FAISS_ASSERT (dis_table.size() == n * n);
275 set_affine_target_dis (dis_table);
276 }
277
278 void set_affine_target_dis (const std::vector<double> & dis_table)
279 {
280 double sum = 0, sum2 = 0;
281 int n2 = n * n;
282 for (int i = 0; i < n2; i++) {
283 sum += dis_table [i];
284 sum2 += dis_table [i] * dis_table [i];
285 }
286 double mean = sum / n2;
287 double stddev = sqrt(sum2 / n2 - (sum / n2) * (sum / n2));
288
289 target_dis.resize (n2);
290
291 for (int i = 0; i < n2; i++) {
292 // the mapping function
293 double td = (dis_table [i] - mean) / stddev * sqrt(nbits / 4) +
294 nbits / 2;
295 target_dis[i] = td;
296 // compute a weight
297 weights.push_back (dis_weight (td));
298 }
299
300 }
301
302 virtual ~ReproduceWithHammingObjective () {}
303 };
304
305 } // anonymous namespace
306
307 // weihgting of distances: it is more important to reproduce small
308 // distances well
309 double ReproduceDistancesObjective::dis_weight (double x) const
310 {
311 return exp (-dis_weight_factor * x);
312 }
313
314
315 double ReproduceDistancesObjective::get_source_dis (int i, int j) const
316 {
317 return source_dis [i * n + j];
318 }
319
320 // cost = quadratic difference between actual distance and Hamming distance
321 double ReproduceDistancesObjective::compute_cost (const int *perm) const
322 {
323 double cost = 0;
324 for (int i = 0; i < n; i++) {
325 for (int j = 0; j < n; j++) {
326 double wanted = target_dis [i * n + j];
327 double w = weights [i * n + j];
328 double actual = get_source_dis (perm[i], perm[j]);
329 cost += w * sqr (wanted - actual);
330 }
331 }
332 return cost;
333 }
334
335 // what would the cost update be if iw and jw were swapped?
336 // computed in O(n) instead of O(n^2) for the full re-computation
337 double ReproduceDistancesObjective::cost_update(
338 const int *perm, int iw, int jw) const
339 {
340 double delta_cost = 0;
341 for (int i = 0; i < n; i++) {
342 if (i == iw) {
343 for (int j = 0; j < n; j++) {
344 double wanted = target_dis [i * n + j],
345 w = weights [i * n + j];
346 double actual = get_source_dis (perm[i], perm[j]);
347 delta_cost -= w * sqr (wanted - actual);
348 double new_actual = get_source_dis (
349 perm[jw],
350 perm[j == iw ? jw : j == jw ? iw : j]);
351 delta_cost += w * sqr (wanted - new_actual);
352 }
353 } else if (i == jw) {
354 for (int j = 0; j < n; j++) {
355 double wanted = target_dis [i * n + j],
356 w = weights [i * n + j];
357 double actual = get_source_dis (perm[i], perm[j]);
358 delta_cost -= w * sqr (wanted - actual);
359 double new_actual = get_source_dis (
360 perm[iw],
361 perm[j == iw ? jw : j == jw ? iw : j]);
362 delta_cost += w * sqr (wanted - new_actual);
363 }
364 } else {
365 int j = iw;
366 {
367 double wanted = target_dis [i * n + j],
368 w = weights [i * n + j];
369 double actual = get_source_dis (perm[i], perm[j]);
370 delta_cost -= w * sqr (wanted - actual);
371 double new_actual = get_source_dis (perm[i], perm[jw]);
372 delta_cost += w * sqr (wanted - new_actual);
373 }
374 j = jw;
375 {
376 double wanted = target_dis [i * n + j],
377 w = weights [i * n + j];
378 double actual = get_source_dis (perm[i], perm[j]);
379 delta_cost -= w * sqr (wanted - actual);
380 double new_actual = get_source_dis (perm[i], perm[iw]);
381 delta_cost += w * sqr (wanted - new_actual);
382 }
383 }
384 }
385 return delta_cost;
386 }
387
388
389
390 ReproduceDistancesObjective::ReproduceDistancesObjective (
391 int n,
392 const double *source_dis_in,
393 const double *target_dis_in,
394 double dis_weight_factor):
395 dis_weight_factor (dis_weight_factor),
396 target_dis (target_dis_in)
397 {
398 this->n = n;
399 set_affine_target_dis (source_dis_in);
400 }
401
402 void ReproduceDistancesObjective::compute_mean_stdev (
403 const double *tab, size_t n2,
404 double *mean_out, double *stddev_out)
405 {
406 double sum = 0, sum2 = 0;
407 for (int i = 0; i < n2; i++) {
408 sum += tab [i];
409 sum2 += tab [i] * tab [i];
410 }
411 double mean = sum / n2;
412 double stddev = sqrt(sum2 / n2 - (sum / n2) * (sum / n2));
413 *mean_out = mean;
414 *stddev_out = stddev;
415 }
416
417 void ReproduceDistancesObjective::set_affine_target_dis (
418 const double *source_dis_in)
419 {
420 int n2 = n * n;
421
422 double mean_src, stddev_src;
423 compute_mean_stdev (source_dis_in, n2, &mean_src, &stddev_src);
424
425 double mean_target, stddev_target;
426 compute_mean_stdev (target_dis, n2, &mean_target, &stddev_target);
427
428 printf ("map mean %g std %g -> mean %g std %g\n",
429 mean_src, stddev_src, mean_target, stddev_target);
430
431 source_dis.resize (n2);
432 weights.resize (n2);
433
434 for (int i = 0; i < n2; i++) {
435 // the mapping function
436 source_dis[i] = (source_dis_in[i] - mean_src) / stddev_src
437 * stddev_target + mean_target;
438
439 // compute a weight
440 weights [i] = dis_weight (target_dis[i]);
441 }
442
443 }
444
445 /****************************************************
446 * Cost functions: RankingScore
447 ****************************************************/
448
449 /// Maintains a 3D table of elementary costs.
450 /// Accumulates elements based on Hamming distance comparisons
451 template <typename Ttab, typename Taccu>
452 struct Score3Computer: PermutationObjective {
453
454 int nc;
455
456 // cost matrix of size nc * nc *nc
457 // n_gt (i,j,k) = count of d_gt(x, y-) < d_gt(x, y+)
458 // where x has PQ code i, y- PQ code j and y+ PQ code k
459 std::vector<Ttab> n_gt;
460
461
462 /// the cost is a triple loop on the nc * nc * nc matrix of entries.
463 ///
464 Taccu compute (const int * perm) const
465 {
466 Taccu accu = 0;
467 const Ttab *p = n_gt.data();
468 for (int i = 0; i < nc; i++) {
469 int ip = perm [i];
470 for (int j = 0; j < nc; j++) {
471 int jp = perm [j];
472 for (int k = 0; k < nc; k++) {
473 int kp = perm [k];
474 if (hamming_dis (ip, jp) <
475 hamming_dis (ip, kp)) {
476 accu += *p; // n_gt [ ( i * nc + j) * nc + k];
477 }
478 p++;
479 }
480 }
481 }
482 return accu;
483 }
484
485
486 /** cost update if entries iw and jw of the permutation would be
487 * swapped.
488 *
489 * The computation is optimized by avoiding elements in the
490 * nc*nc*nc cube that are known not to change. For nc=256, this
491 * reduces the nb of cells to visit to about 6/256 th of the
492 * cells. Practical speedup is about 8x, and the code is quite
493 * complex :-/
494 */
495 Taccu compute_update (const int *perm, int iw, int jw) const
496 {
497 assert (iw != jw);
498 if (iw > jw) std::swap (iw, jw);
499
500 Taccu accu = 0;
501 const Ttab * n_gt_i = n_gt.data();
502 for (int i = 0; i < nc; i++) {
503 int ip0 = perm [i];
504 int ip = perm [i == iw ? jw : i == jw ? iw : i];
505
506 //accu += update_i (perm, iw, jw, ip0, ip, n_gt_i);
507
508 accu += update_i_cross (perm, iw, jw,
509 ip0, ip, n_gt_i);
510
511 if (ip != ip0)
512 accu += update_i_plane (perm, iw, jw,
513 ip0, ip, n_gt_i);
514
515 n_gt_i += nc * nc;
516 }
517
518 return accu;
519 }
520
521
522 Taccu update_i (const int *perm, int iw, int jw,
523 int ip0, int ip, const Ttab * n_gt_i) const
524 {
525 Taccu accu = 0;
526 const Ttab *n_gt_ij = n_gt_i;
527 for (int j = 0; j < nc; j++) {
528 int jp0 = perm[j];
529 int jp = perm [j == iw ? jw : j == jw ? iw : j];
530 for (int k = 0; k < nc; k++) {
531 int kp0 = perm [k];
532 int kp = perm [k == iw ? jw : k == jw ? iw : k];
533 int ng = n_gt_ij [k];
534 if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
535 accu += ng;
536 }
537 if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp0)) {
538 accu -= ng;
539 }
540 }
541 n_gt_ij += nc;
542 }
543 return accu;
544 }
545
546 // 2 inner loops for the case ip0 != ip
547 Taccu update_i_plane (const int *perm, int iw, int jw,
548 int ip0, int ip, const Ttab * n_gt_i) const
549 {
550 Taccu accu = 0;
551 const Ttab *n_gt_ij = n_gt_i;
552
553 for (int j = 0; j < nc; j++) {
554 if (j != iw && j != jw) {
555 int jp = perm[j];
556 for (int k = 0; k < nc; k++) {
557 if (k != iw && k != jw) {
558 int kp = perm [k];
559 Ttab ng = n_gt_ij [k];
560 if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
561 accu += ng;
562 }
563 if (hamming_dis (ip0, jp) < hamming_dis (ip0, kp)) {
564 accu -= ng;
565 }
566 }
567 }
568 }
569 n_gt_ij += nc;
570 }
571 return accu;
572 }
573
574 /// used for the 8 cells were the 3 indices are swapped
575 inline Taccu update_k (const int *perm, int iw, int jw,
576 int ip0, int ip, int jp0, int jp,
577 int k,
578 const Ttab * n_gt_ij) const
579 {
580 Taccu accu = 0;
581 int kp0 = perm [k];
582 int kp = perm [k == iw ? jw : k == jw ? iw : k];
583 Ttab ng = n_gt_ij [k];
584 if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
585 accu += ng;
586 }
587 if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp0)) {
588 accu -= ng;
589 }
590 return accu;
591 }
592
593 /// compute update on a line of k's, where i and j are swapped
594 Taccu update_j_line (const int *perm, int iw, int jw,
595 int ip0, int ip, int jp0, int jp,
596 const Ttab * n_gt_ij) const
597 {
598 Taccu accu = 0;
599 for (int k = 0; k < nc; k++) {
600 if (k == iw || k == jw) continue;
601 int kp = perm [k];
602 Ttab ng = n_gt_ij [k];
603 if (hamming_dis (ip, jp) < hamming_dis (ip, kp)) {
604 accu += ng;
605 }
606 if (hamming_dis (ip0, jp0) < hamming_dis (ip0, kp)) {
607 accu -= ng;
608 }
609 }
610 return accu;
611 }
612
613
614 /// considers the 2 pairs of crossing lines j=iw or jw and k = iw or kw
615 Taccu update_i_cross (const int *perm, int iw, int jw,
616 int ip0, int ip, const Ttab * n_gt_i) const
617 {
618 Taccu accu = 0;
619 const Ttab *n_gt_ij = n_gt_i;
620
621 for (int j = 0; j < nc; j++) {
622 int jp0 = perm[j];
623 int jp = perm [j == iw ? jw : j == jw ? iw : j];
624
625 accu += update_k (perm, iw, jw, ip0, ip, jp0, jp, iw, n_gt_ij);
626 accu += update_k (perm, iw, jw, ip0, ip, jp0, jp, jw, n_gt_ij);
627
628 if (jp != jp0)
629 accu += update_j_line (perm, iw, jw, ip0, ip, jp0, jp, n_gt_ij);
630
631 n_gt_ij += nc;
632 }
633 return accu;
634 }
635
636
637 /// PermutationObjective implementeation (just negates the scores
638 /// for minimization)
639
640 virtual double compute_cost (const int *perm) const {
641 return -compute (perm);
642 }
643
644
645
646 virtual double cost_update (const int *perm, int iw, int jw) const
647 {
648 double ret = -compute_update (perm, iw, jw);
649 return ret;
650 }
651
652 virtual ~Score3Computer () {}
653
654 };
655
656
657
658
659
660 struct IndirectSort {
661 const float *tab;
662 bool operator () (int a, int b) {return tab[a] < tab[b]; }
663 };
664
665
666
667 struct RankingScore2: Score3Computer<float, double> {
668 int nbits;
669 int nq, nb;
670 const uint32_t *qcodes, *bcodes;
671 const float *gt_distances;
672
673 RankingScore2 (int nbits, int nq, int nb,
674 const uint32_t *qcodes, const uint32_t *bcodes,
675 const float *gt_distances):
676 nbits(nbits), nq(nq), nb(nb), qcodes(qcodes),
677 bcodes(bcodes), gt_distances(gt_distances)
678 {
679 n = nc = 1 << nbits;
680 n_gt.resize (nc * nc * nc);
681 init_n_gt ();
682 }
683
684
685 double rank_weight (int r)
686 {
687 return 1.0 / (r + 1);
688 }
689
690 /// count nb of i, j in a x b st. i < j
691 /// a and b should be sorted on input
692 /// they are the ranks of j and k respectively.
693 /// specific version for diff-of-rank weighting, cannot optimized
694 /// with a cumulative table
695 double accum_gt_weight_diff (const std::vector<int> & a,
696 const std::vector<int> & b)
697 {
698 int nb = b.size(), na = a.size();
699
700 double accu = 0;
701 int j = 0;
702 for (int i = 0; i < na; i++) {
703 int ai = a[i];
704 while (j < nb && ai >= b[j]) j++;
705
706 double accu_i = 0;
707 for (int k = j; k < b.size(); k++)
708 accu_i += rank_weight (b[k] - ai);
709
710 accu += rank_weight (ai) * accu_i;
711
712 }
713 return accu;
714 }
715
716 void init_n_gt ()
717 {
718 for (int q = 0; q < nq; q++) {
719 const float *gtd = gt_distances + q * nb;
720 const uint32_t *cb = bcodes;// all same codes
721 float * n_gt_q = & n_gt [qcodes[q] * nc * nc];
722
723 printf("init gt for q=%d/%d \r", q, nq); fflush(stdout);
724
725 std::vector<int> rankv (nb);
726 int * ranks = rankv.data();
727
728 // elements in each code bin, ordered by rank within each bin
729 std::vector<std::vector<int> > tab (nc);
730
731 { // build rank table
732 IndirectSort s = {gtd};
733 for (int j = 0; j < nb; j++) ranks[j] = j;
734 std::sort (ranks, ranks + nb, s);
735 }
736
737 for (int rank = 0; rank < nb; rank++) {
738 int i = ranks [rank];
739 tab [cb[i]].push_back (rank);
740 }
741
742
743 // this is very expensive. Any suggestion for improvement
744 // welcome.
745 for (int i = 0; i < nc; i++) {
746 std::vector<int> & di = tab[i];
747 for (int j = 0; j < nc; j++) {
748 std::vector<int> & dj = tab[j];
749 n_gt_q [i * nc + j] += accum_gt_weight_diff (di, dj);
750
751 }
752 }
753
754 }
755
756 }
757
758 };
759
760
761 /*****************************************
762 * PolysemousTraining
763 ******************************************/
764
765
766
767 PolysemousTraining::PolysemousTraining ()
768 {
769 optimization_type = OT_ReproduceDistances_affine;
770 ntrain_permutation = 0;
771 dis_weight_factor = log(2);
772 }
773
774
775
776 void PolysemousTraining::optimize_reproduce_distances (
777 ProductQuantizer &pq) const
778 {
779
780 int dsub = pq.dsub;
781
782 int n = pq.ksub;
783 int nbits = pq.nbits;
784
785 #pragma omp parallel for
786 for (int m = 0; m < pq.M; m++) {
787 std::vector<double> dis_table;
788
789 // printf ("Optimizing quantizer %d\n", m);
790
791 float * centroids = pq.get_centroids (m, 0);
792
793 for (int i = 0; i < n; i++) {
794 for (int j = 0; j < n; j++) {
795 dis_table.push_back (fvec_L2sqr (centroids + i * dsub,
796 centroids + j * dsub,
797 dsub));
798 }
799 }
800
801 std::vector<int> perm (n);
802 ReproduceWithHammingObjective obj (
803 nbits, dis_table,
804 dis_weight_factor);
805
806
807 SimulatedAnnealingOptimizer optim (&obj, *this);
808
809 if (log_pattern.size()) {
810 char fname[256];
811 snprintf (fname, 256, log_pattern.c_str(), m);
812 printf ("opening log file %s\n", fname);
813 optim.logfile = fopen (fname, "w");
814 FAISS_ASSERT (optim.logfile || !"could not open logfile");
815 }
816 double final_cost = optim.run_optimization (perm.data());
817
818 if (verbose > 0) {
819 printf ("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
820 m, optim.init_cost, final_cost);
821 }
822
823 if (log_pattern.size()) fclose (optim.logfile);
824
825 std::vector<float> centroids_copy;
826 for (int i = 0; i < dsub * n; i++)
827 centroids_copy.push_back (centroids[i]);
828
829 for (int i = 0; i < n; i++)
830 memcpy (centroids + perm[i] * dsub,
831 centroids_copy.data() + i * dsub,
832 dsub * sizeof(centroids[0]));
833
834 }
835
836 }
837
838
839 void PolysemousTraining::optimize_ranking (
840 ProductQuantizer &pq, size_t n, const float *x) const
841 {
842
843 int dsub = pq.dsub;
844
845 int nbits = pq.nbits;
846
847 std::vector<uint8_t> all_codes (pq.code_size * n);
848
849 pq.compute_codes (x, all_codes.data(), n);
850
851 FAISS_ASSERT (pq.byte_per_idx == 1);
852
853 if (n == 0)
854 pq.compute_sdc_table ();
855
856 #pragma omp parallel for
857 for (int m = 0; m < pq.M; m++) {
858 size_t nq, nb;
859 std::vector <uint32_t> codes; // query codes, then db codes
860 std::vector <float> gt_distances; // nq * nb matrix of distances
861
862 if (n > 0) {
863 std::vector<float> xtrain (n * dsub);
864 for (int i = 0; i < n; i++)
865 memcpy (xtrain.data() + i * dsub,
866 x + i * pq.d + m * dsub,
867 sizeof(float) * dsub);
868
869 codes.resize (n);
870 for (int i = 0; i < n; i++)
871 codes [i] = all_codes [i * pq.code_size + m];
872
873 nq = n / 4; nb = n - nq;
874 const float *xq = xtrain.data();
875 const float *xb = xq + nq * dsub;
876
877 gt_distances.resize (nq * nb);
878
879 pairwise_L2sqr (dsub,
880 nq, xq,
881 nb, xb,
882 gt_distances.data());
883 } else {
884 nq = nb = pq.ksub;
885 codes.resize (2 * nq);
886 for (int i = 0; i < nq; i++)
887 codes[i] = codes [i + nq] = i;
888
889 gt_distances.resize (nq * nb);
890
891 memcpy (gt_distances.data (),
892 pq.sdc_table.data () + m * nq * nb,
893 sizeof (float) * nq * nb);
894 }
895
896 double t0 = getmillisecs ();
897
898 PermutationObjective *obj = new RankingScore2 (
899 nbits, nq, nb,
900 codes.data(), codes.data() + nq,
901 gt_distances.data ());
902 if (verbose > 0) {
903 printf(" m=%d, nq=%ld, nb=%ld, intialize RankingScore "
904 "in %.3f ms\n",
905 m, nq, nb, getmillisecs () - t0);
906 }
907
908 SimulatedAnnealingOptimizer optim (obj, *this);
909
910 if (log_pattern.size()) {
911 char fname[256];
912 snprintf (fname, 256, log_pattern.c_str(), m);
913 printf ("opening log file %s\n", fname);
914 optim.logfile = fopen (fname, "w");
915 FAISS_ASSERT (optim.logfile || !"could not open logfile");
916 }
917
918 std::vector<int> perm (pq.ksub);
919
920 double final_cost = optim.run_optimization (perm.data());
921 printf ("SimulatedAnnealingOptimizer for m=%d: %g -> %g\n",
922 m, optim.init_cost, final_cost);
923
924 if (log_pattern.size()) fclose (optim.logfile);
925
926 delete obj;
927
928 float * centroids = pq.get_centroids (m, 0);
929
930 std::vector<float> centroids_copy;
931 for (int i = 0; i < dsub * pq.ksub; i++)
932 centroids_copy.push_back (centroids[i]);
933
934 for (int i = 0; i < pq.ksub; i++)
935 memcpy (centroids + perm[i] * dsub,
936 centroids_copy.data() + i * dsub,
937 dsub * sizeof(centroids[0]));
938
939 }
940
941 }
942
943
944
945 void PolysemousTraining::optimize_pq_for_hamming (ProductQuantizer &pq,
946 size_t n, const float *x) const
947 {
948 if (optimization_type == OT_None) {
949
950 } else if (optimization_type == OT_ReproduceDistances_affine) {
951 optimize_reproduce_distances (pq);
952 } else {
953 optimize_ranking (pq, n, x);
954 }
955
956 pq.compute_sdc_table ();
957
958 }
959
960
961
962 } // namespace faiss
faiss::RandomGenerator
random generator that can be used in multithreaded contexts
Definition: utils.h:49
faiss::ProductQuantizer::nbits
size_t nbits
number of bits per quantization index
Definition: ProductQuantizer.h:30
faiss::fvec_L2sqr
float fvec_L2sqr(const float *x, const float *y, size_t d)
Squared L2 distance between two vectors.
Definition: utils.cpp:430
faiss::ProductQuantizer::byte_per_idx
size_t byte_per_idx
nb bytes per code component (1 or 2)
Definition: ProductQuantizer.h:34
faiss::Score3Computer::compute_update
Taccu compute_update(const int *perm, int iw, int jw) const
Definition: PolysemousTraining.cpp:495
faiss::ProductQuantizer::sdc_table
std::vector< float > sdc_table
Symmetric Distance Table.
Definition: ProductQuantizer.h:156
faiss::SimulatedAnnealingOptimizer::SimulatedAnnealingOptimizer
SimulatedAnnealingOptimizer(PermutationObjective *obj, const SimulatedAnnealingParameters &p)
logs values of the cost function
Definition: PolysemousTraining.cpp:68
faiss::SimulatedAnnealingOptimizer::n
int n
size of the permutation
Definition: PolysemousTraining.h:103
faiss::Score3Computer::compute
Taccu compute(const int *perm) const
Definition: PolysemousTraining.cpp:464
faiss::ProductQuantizer::dsub
size_t dsub
dimensionality of each subvector
Definition: ProductQuantizer.h:33
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::Score3Computer::update_j_line
Taccu update_j_line(const int *perm, int iw, int jw, int ip0, int ip, int jp0, int jp, const Ttab *n_gt_ij) const
compute update on a line of k&#39;s, where i and j are swapped
Definition: PolysemousTraining.cpp:594
faiss::ReproduceDistancesObjective::target_dis
const double * target_dis
wanted distances (size n^2)
Definition: PolysemousTraining.h:70
faiss::ProductQuantizer::code_size
size_t code_size
byte per indexed vector
Definition: ProductQuantizer.h:35
faiss::SimulatedAnnealingOptimizer::init_cost
double init_cost
remember intial cost of optimization
Definition: PolysemousTraining.h:111
faiss::RandomGenerator::rand_int
int rand_int()
random 31-bit positive integer
faiss::ProductQuantizer::ksub
size_t ksub
number of centroids for each subquantizer
Definition: ProductQuantizer.h:36
faiss::PolysemousTraining::optimize_ranking
void optimize_ranking(ProductQuantizer &pq, size_t n, const float *x) const
called by optimize_pq_for_hamming
Definition: PolysemousTraining.cpp:839
faiss::pairwise_L2sqr
void pairwise_L2sqr(long d, long nq, const float *xq, long nb, const float *xb, float *dis, long ldq, long ldb, long ldd)
Definition: utils.cpp:1228
faiss::Score3Computer
Definition: PolysemousTraining.cpp:452
faiss::getmillisecs
double getmillisecs()
ms elapsed since some arbitrary epoch
Definition: utils.cpp:71
faiss::ReproduceDistancesObjective::weights
std::vector< double > weights
weights for each distance (size n^2)
Definition: PolysemousTraining.h:71
faiss::RankingScore2::accum_gt_weight_diff
double accum_gt_weight_diff(const std::vector< int > &a, const std::vector< int > &b)
Definition: PolysemousTraining.cpp:695
faiss::SimulatedAnnealingParameters
parameters used for the simulated annealing method
Definition: PolysemousTraining.h:26
faiss::IndirectSort
Definition: PolysemousTraining.cpp:660
faiss::Score3Computer::update_i_cross
Taccu update_i_cross(const int *perm, int iw, int jw, int ip0, int ip, const Ttab *n_gt_i) const
considers the 2 pairs of crossing lines j=iw or jw and k = iw or kw
Definition: PolysemousTraining.cpp:615
faiss::ProductQuantizer::M
size_t M
number of subquantizers
Definition: ProductQuantizer.h:29
faiss::PolysemousTraining::OT_ReproduceDistances_affine
default
Definition: PolysemousTraining.h:131
faiss::PermutationObjective
abstract class for the loss function
Definition: PolysemousTraining.h:45
faiss::ProductQuantizer
Definition: ProductQuantizer.h:26
faiss::Score3Computer::update_k
Taccu update_k(const int *perm, int iw, int jw, int ip0, int ip, int jp0, int jp, int k, const Ttab *n_gt_ij) const
used for the 8 cells were the 3 indices are swapped
Definition: PolysemousTraining.cpp:575
faiss::ReproduceDistancesObjective::source_dis
std::vector< double > source_dis
&quot;real&quot; corrected distances (size n^2)
Definition: PolysemousTraining.h:69
faiss::ProductQuantizer::get_centroids
float * get_centroids(size_t m, size_t i)
return the centroids associated with subvector m
Definition: ProductQuantizer.h:56
faiss::PolysemousTraining::optimize_reproduce_distances
void optimize_reproduce_distances(ProductQuantizer &pq) const
called by optimize_pq_for_hamming
Definition: PolysemousTraining.cpp:776
faiss::PolysemousTraining::optimize_pq_for_hamming
void optimize_pq_for_hamming(ProductQuantizer &pq, size_t n, const float *x) const
Definition: PolysemousTraining.cpp:945
faiss::ProductQuantizer::d
size_t d
size of the input vectors
Definition: ProductQuantizer.h:28
faiss::RankingScore2
Definition: PolysemousTraining.cpp:667
faiss::SimulatedAnnealingOptimizer
Simulated annealing optimization algorithm for permutations.
Definition: PolysemousTraining.h:100
faiss::Score3Computer::compute_cost
virtual double compute_cost(const int *perm) const
Definition: PolysemousTraining.cpp:640