Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexIVFPQ.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 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #include "../../IndexFlat.h"
13 #include "../../IndexIVFPQ.h"
14 #include "../GpuIndexIVFPQ.h"
15 #include "../StandardGpuResources.h"
16 #include "../utils/DeviceUtils.h"
17 #include "../test/TestUtils.h"
18 #include <cmath>
19 #include <gtest/gtest.h>
20 #include <sstream>
21 #include <vector>
22
23
24 void pickEncoding(int& codes, int& dim) {
25 std::vector<int> codeSizes{3, 4, 8, 12, 16, 20, 24, 28, 32, 40, 48, 56, 64};
26 std::vector<int> dimSizes{4, 8, 16, 32};
27
28 codes = codeSizes[faiss::gpu::randVal(0, codeSizes.size() - 1)];
29
30 while (true) {
31 dim = codes * dimSizes[faiss::gpu::randVal(0, dimSizes.size() - 1)];
32
33 // for such a small test, super-low or high dim is more likely to
34 // generate comparison errors
35 if (dim < 512 && dim >= 64) {
36 return;
37 }
38 }
39 }
40
41 struct Options {
42 Options() {
43 numAdd = faiss::gpu::randVal(10000, 30000);
44 numCentroids = std::sqrt((float) numAdd);
45 numTrain = numCentroids * 40;
46
47 pickEncoding(codes, dim);
48
49 bitsPerCode = faiss::gpu::randVal(3, 8);
50 nprobe = std::min(faiss::gpu::randVal(40, 1000), numCentroids);
51 numQuery = faiss::gpu::randVal(32, 256);
52 k = std::min(faiss::gpu::randVal(10, 50), numAdd / 40);
53 usePrecomputed = faiss::gpu::randBool();
54 indicesOpt = faiss::gpu::randSelect({
55 faiss::gpu::INDICES_CPU,
56 faiss::gpu::INDICES_32_BIT,
57 faiss::gpu::INDICES_64_BIT});
58 if (codes == 64) {
59 // 64x8 lookup can only fit using float16
60 useFloat16 = true;
61 } else {
62 useFloat16 = faiss::gpu::randBool();
63 }
64
65 device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
66 }
67
68 std::string toString() const {
69 std::stringstream str;
70 str << "IVFPQ device " << device
71 << " numVecs " << numAdd
72 << " dim " << dim
73 << " numCentroids " << numCentroids
74 << " codes " << codes
75 << " bitsPerCode " << bitsPerCode
76 << " nprobe " << nprobe
77 << " numQuery " << numQuery
78 << " k " << k
79 << " usePrecomputed " << usePrecomputed
80 << " indicesOpt " << indicesOpt
81 << " useFloat16 " << useFloat16;
82
83 return str.str();
84 }
85
86 float getCompareEpsilon() const {
87 return 0.015f;
88 }
89
90 float getPctMaxDiff1() const {
91 return useFloat16 ? 0.30f : 0.10f;
92 }
93
94 float getPctMaxDiffN() const {
95 return useFloat16 ? 0.05f : 0.015f;
96 }
97
98 int numAdd;
99 int numCentroids;
100 int numTrain;
101 int codes;
102 int dim;
103 int bitsPerCode;
104 int nprobe;
105 int numQuery;
106 int k;
107 bool usePrecomputed;
108 faiss::gpu::IndicesOptions indicesOpt;
109 bool useFloat16;
110 int device;
111 };
112
113 TEST(TestGpuIndexIVFPQ, Query) {
114 for (int tries = 0; tries < 5; ++tries) {
115 faiss::gpu::newTestSeed();
116
117 Options opt;
118
119 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
120 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
121
122 faiss::IndexFlatL2 coarseQuantizer(opt.dim);
123 faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
124 opt.codes, opt.bitsPerCode);
125 cpuIndex.nprobe = opt.nprobe;
126 cpuIndex.train(opt.numTrain, trainVecs.data());
127 cpuIndex.add(opt.numAdd, addVecs.data());
128
129 faiss::gpu::StandardGpuResources res;
130 res.noTempMemory();
131
132 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
133 opt.indicesOpt,
134 opt.useFloat16,
135 &cpuIndex);
136 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
137 gpuIndex.setNumProbes(opt.nprobe);
138
139 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
140 opt.numQuery, opt.dim, opt.k, opt.toString(),
141 opt.getCompareEpsilon(),
142 opt.getPctMaxDiff1(),
143 opt.getPctMaxDiffN());
144 }
145 }
146
147 TEST(TestGpuIndexIVFPQ, Add) {
148 for (int tries = 0; tries < 5; ++tries) {
149 faiss::gpu::newTestSeed();
150
151 Options opt;
152
153 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
154 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
155
156 faiss::IndexFlatL2 coarseQuantizer(opt.dim);
157 faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
158 opt.codes, opt.bitsPerCode);
159 cpuIndex.nprobe = opt.nprobe;
160 cpuIndex.train(opt.numTrain, trainVecs.data());
161
162 faiss::gpu::StandardGpuResources res;
163 res.noTempMemory();
164
165 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
166 opt.indicesOpt,
167 opt.useFloat16,
168 &cpuIndex);
169 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
170 gpuIndex.setNumProbes(opt.nprobe);
171
172 gpuIndex.add(opt.numAdd, addVecs.data());
173 cpuIndex.add(opt.numAdd, addVecs.data());
174
175 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
176 opt.numQuery, opt.dim, opt.k, opt.toString(),
177 opt.getCompareEpsilon(),
178 opt.getPctMaxDiff1(),
179 opt.getPctMaxDiffN());
180 }
181 }
182
183 TEST(TestGpuIndexIVFPQ, CopyTo) {
184 faiss::gpu::newTestSeed();
185
186 Options opt;
187 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
188 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
189
190 faiss::gpu::StandardGpuResources res;
191 res.noTempMemory();
192
193 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
194 opt.dim,
195 opt.numCentroids,
196 opt.codes,
197 opt.bitsPerCode,
198 opt.usePrecomputed,
199 opt.indicesOpt,
200 opt.useFloat16,
201 faiss::METRIC_L2);
202 gpuIndex.setNumProbes(opt.nprobe);
203 gpuIndex.train(opt.numTrain, trainVecs.data());
204 gpuIndex.add(opt.numAdd, addVecs.data());
205
206 // Use garbage values to see if we overwrite them
207 faiss::IndexFlatL2 cpuQuantizer(1);
208 faiss::IndexIVFPQ cpuIndex(&cpuQuantizer, 1, 1, 1, 1);
209
210 gpuIndex.copyTo(&cpuIndex);
211
212 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
213 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
214
215 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
216 EXPECT_EQ(cpuIndex.d, opt.dim);
217 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
218 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
219 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
220 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
221 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
222 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
223
224 // Query both objects; results should be equivalent
225 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
226 opt.numQuery, opt.dim, opt.k, opt.toString(),
227 opt.getCompareEpsilon(),
228 opt.getPctMaxDiff1(),
229 opt.getPctMaxDiffN());
230 }
231
232 TEST(TestGpuIndexIVFPQ, CopyFrom) {
233 faiss::gpu::newTestSeed();
234
235 Options opt;
236 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
237 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
238
239 faiss::IndexFlatL2 coarseQuantizer(opt.dim);
240 faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
241 opt.codes, opt.bitsPerCode);
242 cpuIndex.nprobe = opt.nprobe;
243 cpuIndex.train(opt.numTrain, trainVecs.data());
244 cpuIndex.add(opt.numAdd, addVecs.data());
245
246 // Use garbage values to see if we overwrite them
247 faiss::gpu::StandardGpuResources res;
248 res.noTempMemory();
249
250 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
251 1, 1, 1, 1, false, opt.indicesOpt,
252 opt.useFloat16, faiss::METRIC_L2);
253 gpuIndex.setNumProbes(1);
254
255 gpuIndex.copyFrom(&cpuIndex);
256
257 // Make sure we are equivalent
258 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
259 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
260
261 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
262 EXPECT_EQ(cpuIndex.d, opt.dim);
263 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
264 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
265 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
266 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
267 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
268 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
269
270 // Query both objects; results should be equivalent
271 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
272 opt.numQuery, opt.dim, opt.k, opt.toString(),
273 opt.getCompareEpsilon(),
274 opt.getPctMaxDiff1(),
275 opt.getPctMaxDiffN());
276 }
277
278 TEST(TestGpuIndexIVFPQ, QueryNaN) {
279 faiss::gpu::newTestSeed();
280
281 Options opt;
282
283 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
284 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
285
286 faiss::gpu::StandardGpuResources res;
287 res.noTempMemory();
288
289 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
290 opt.dim,
291 opt.numCentroids,
292 opt.codes,
293 opt.bitsPerCode,
294 opt.usePrecomputed,
295 opt.indicesOpt,
296 opt.useFloat16,
297 faiss::METRIC_L2);
298
299 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
300 gpuIndex.setNumProbes(opt.nprobe);
301
302 gpuIndex.train(opt.numTrain, trainVecs.data());
303 gpuIndex.add(opt.numAdd, addVecs.data());
304
305 int numQuery = 10;
306 std::vector<float> nans(numQuery * opt.dim,
307 std::numeric_limits<float>::quiet_NaN());
308
309 std::vector<float> distances(numQuery * opt.k, 0);
310 std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
311
312 gpuIndex.search(numQuery,
313 nans.data(),
314 opt.k,
315 distances.data(),
316 indices.data());
317
318 for (int q = 0; q < numQuery; ++q) {
319 for (int k = 0; k < opt.k; ++k) {
320 EXPECT_EQ(indices[q * opt.k + k], -1);
321 EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
322 }
323 }
324 }
325
326 TEST(TestGpuIndexIVFPQ, AddNaN) {
327 faiss::gpu::newTestSeed();
328
329 Options opt;
330
331 faiss::gpu::StandardGpuResources res;
332 res.noTempMemory();
333
334 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
335 opt.dim,
336 opt.numCentroids,
337 opt.codes,
338 opt.bitsPerCode,
339 opt.usePrecomputed,
340 opt.indicesOpt,
341 opt.useFloat16,
342 faiss::METRIC_L2);
343
344 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
345 gpuIndex.setNumProbes(opt.nprobe);
346
347 int numNans = 10;
348 std::vector<float> nans(numNans * opt.dim,
349 std::numeric_limits<float>::quiet_NaN());
350
351 // Make one vector valid, which should actually add
352 for (int i = 0; i < opt.dim; ++i) {
353 nans[i] = 0.0f;
354 }
355
356 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
357 gpuIndex.train(opt.numTrain, trainVecs.data());
358
359 // should not crash
360 EXPECT_EQ(gpuIndex.ntotal, 0);
361 gpuIndex.add(numNans, nans.data());
362
363 // Only the single valid vector should have added
364 EXPECT_EQ(gpuIndex.ntotal, 1);
365
366 std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
367 std::vector<float> distance(opt.numQuery * opt.k, 0);
368 std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
369
370 // should not crash
371 gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
372 distance.data(), indices.data());
373
374 }
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:72
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:30
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:24
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:25
Options
Definition: TestGpuIndexIVFFlat.cpp:27