1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
/*
* Copyright (c) Meta Platforms, Inc. and affiliates.
*
* This source code is licensed under the MIT license found in the
* LICENSE file in the root directory of this source tree.
*/
#include <faiss/IndexAdditiveQuantizer.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/utils/random.h>
#include <gtest/gtest.h>
/* This test creates a 3-level RCQ and performs a search on it.
* Then it crops the RCQ to just the 2 first levels and verifies that
* the 3-level vectors are in a subtree that was visited in the 2-level RCQ. */
TEST(RCQCropping, test_cropping) {
size_t nq = 10, nt = 2000, nb = 1000, d = 32;
using idx_t = faiss::idx_t;
std::vector<float> buf((nq + nb + nt) * d);
faiss::rand_smooth_vectors(nq + nb + nt, d, buf.data(), 1234);
const float* xt = buf.data();
const float* xb = xt + nt * d;
const float* xq = xb + nb * d;
std::vector<size_t> nbits = {5, 4, 4};
faiss::ResidualCoarseQuantizer rcq(d, nbits);
rcq.train(nt, xt);
// the test below works only for beam size == nprobe
rcq.set_beam_factor(1.0);
// perform search
int nprobe = 15;
std::vector<faiss::idx_t> Iref(nq * nprobe);
std::vector<float> Dref(nq * nprobe);
rcq.search(nq, xq, nprobe, Dref.data(), Iref.data());
// crop to the first 2 quantization levels
int last_nbits = nbits.back();
nbits.pop_back();
faiss::ResidualCoarseQuantizer rcq_cropped(d, nbits);
rcq_cropped.initialize_from(rcq);
EXPECT_EQ(rcq_cropped.ntotal, rcq.ntotal >> last_nbits);
// perform search
std::vector<faiss::idx_t> Inew(nq * nprobe);
std::vector<float> Dnew(nq * nprobe);
rcq_cropped.search(nq, xq, nprobe, Dnew.data(), Inew.data());
// these bits are in common between the two RCQs
idx_t mask = ((idx_t)1 << rcq_cropped.rq.tot_bits) - 1;
for (int q = 0; q < nq; q++) {
for (int i = 0; i < nprobe; i++) {
idx_t fine = Iref[q * nprobe + i];
EXPECT_GE(fine, 0);
bool found = false;
// fine should be generated from a path that passes through coarse
for (int j = 0; j < nprobe; j++) {
idx_t coarse = Inew[q * nprobe + j];
if ((fine & mask) == coarse) {
found = true;
break;
}
}
EXPECT_TRUE(found);
}
}
}
TEST(RCQCropping, search_params) {
size_t nq = 10, nt = 2000, nb = 1000, d = 32;
using idx_t = faiss::idx_t;
std::vector<float> buf((nq + nb + nt) * d);
faiss::rand_smooth_vectors(nq + nb + nt, d, buf.data(), 1234);
const float* xt = buf.data();
const float* xb = xt + nt * d;
const float* xq = xb + nb * d;
std::vector<size_t> nbits = {3, 6, 3};
faiss::ResidualCoarseQuantizer quantizer(d, nbits);
size_t ntotal = (size_t)1 << quantizer.rq.tot_bits;
faiss::IndexIVFScalarQuantizer index(
&quantizer, d, ntotal, faiss::ScalarQuantizer::QT_8bit);
index.quantizer_trains_alone = true;
index.train(nt, xt);
index.add(nb, xb);
index.nprobe = 10;
int k = 4;
float beam_factor_1 = 8.0;
quantizer.set_beam_factor(beam_factor_1);
std::vector<idx_t> I1(nq * k);
std::vector<float> D1(nq * k);
index.search(nq, xq, k, D1.data(), I1.data());
// change from 8 to 1
quantizer.set_beam_factor(1.0f);
std::vector<idx_t> I2(nq * k);
std::vector<float> D2(nq * k);
index.search(nq, xq, k, D2.data(), I2.data());
// make sure it changes the result
EXPECT_NE(I1, I2);
EXPECT_NE(D1, D2);
// override the class level beam factor
faiss::SearchParametersResidualCoarseQuantizer params1;
params1.beam_factor = beam_factor_1;
faiss::SearchParametersIVF params;
params.nprobe = index.nprobe;
params.quantizer_params = ¶ms1;
std::vector<idx_t> I3(nq * k);
std::vector<float> D3(nq * k);
index.search(nq, xq, k, D3.data(), I3.data(), ¶ms);
// make sure we find back the original results
EXPECT_EQ(I1, I3);
EXPECT_EQ(D1, D3);
}
|
