Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexIVFPQ.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
10 #include "../../IndexFlat.h"
11 #include "../../IndexIVFPQ.h"
12 #include "../GpuIndexIVFPQ.h"
13 #include "../StandardGpuResources.h"
14 #include "../utils/DeviceUtils.h"
15 #include "../test/TestUtils.h"
16 #include <cmath>
17 #include <gtest/gtest.h>
18 #include <sstream>
19 #include <vector>
20
21
22 void pickEncoding(int& codes, int& dim) {
23 std::vector<int> codeSizes{
24 3, 4, 8, 12, 16, 20, 24,
25 28, 32, 40, 48, 56, 64, 96
26 };
27
28 // Above 32 doesn't work with no precomputed codes
29 std::vector<int> dimSizes{4, 8, 10, 12, 16, 20, 24, 28, 32};
30
31 while (true) {
32 codes = codeSizes[faiss::gpu::randVal(0, codeSizes.size() - 1)];
33 dim = codes * dimSizes[faiss::gpu::randVal(0, dimSizes.size() - 1)];
34
35 // for such a small test, super-low or high dim is more likely to
36 // generate comparison errors
37 if (dim < 512 && dim >= 64) {
38 return;
39 }
40 }
41 }
42
43 struct Options {
44 Options() {
45 numAdd = faiss::gpu::randVal(10000, 30000);
46 numCentroids = std::sqrt((float) numAdd);
47 numTrain = numCentroids * 40;
48
49 pickEncoding(codes, dim);
50
51 bitsPerCode = faiss::gpu::randVal(3, 8);
52 nprobe = std::min(faiss::gpu::randVal(40, 1000), numCentroids);
53 numQuery = faiss::gpu::randVal(32, 256);
54 k = std::min(faiss::gpu::randVal(10, 50), numAdd / 40);
55 usePrecomputed = faiss::gpu::randBool();
56 indicesOpt = faiss::gpu::randSelect({
57 faiss::gpu::INDICES_CPU,
58 faiss::gpu::INDICES_32_BIT,
59 faiss::gpu::INDICES_64_BIT});
60 if (codes > 48) {
61 // large codes can only fit using float16
62 useFloat16 = true;
63 } else {
64 useFloat16 = faiss::gpu::randBool();
65 }
66
67 device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
68 }
69
70 std::string toString() const {
71 std::stringstream str;
72 str << "IVFPQ device " << device
73 << " numVecs " << numAdd
74 << " dim " << dim
75 << " numCentroids " << numCentroids
76 << " codes " << codes
77 << " bitsPerCode " << bitsPerCode
78 << " nprobe " << nprobe
79 << " numQuery " << numQuery
80 << " k " << k
81 << " usePrecomputed " << usePrecomputed
82 << " indicesOpt " << indicesOpt
83 << " useFloat16 " << useFloat16;
84
85 return str.str();
86 }
87
88 float getCompareEpsilon() const {
89 return 0.03f;
90 }
91
92 float getPctMaxDiff1() const {
93 return useFloat16 ? 0.30f : 0.10f;
94 }
95
96 float getPctMaxDiffN() const {
97 return useFloat16 ? 0.05f : 0.02f;
98 }
99
100 int numAdd;
101 int numCentroids;
102 int numTrain;
103 int codes;
104 int dim;
105 int bitsPerCode;
106 int nprobe;
107 int numQuery;
108 int k;
109 bool usePrecomputed;
110 faiss::gpu::IndicesOptions indicesOpt;
111 bool useFloat16;
112 int device;
113 };
114
115 TEST(TestGpuIndexIVFPQ, Query) {
116 for (int tries = 0; tries < 5; ++tries) {
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::GpuIndexIVFPQConfig config;
133 config.device = opt.device;
134 config.usePrecomputedTables = opt.usePrecomputed;
135 config.indicesOptions = opt.indicesOpt;
136 config.useFloat16LookupTables = opt.useFloat16;
137
138 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
139 gpuIndex.setNumProbes(opt.nprobe);
140
141 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
142 opt.numQuery, opt.dim, opt.k, opt.toString(),
143 opt.getCompareEpsilon(),
144 opt.getPctMaxDiff1(),
145 opt.getPctMaxDiffN());
146 }
147 }
148
149 TEST(TestGpuIndexIVFPQ, Add) {
150 for (int tries = 0; tries < 5; ++tries) {
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::GpuIndexIVFPQConfig config;
166 config.device = opt.device;
167 config.usePrecomputedTables = opt.usePrecomputed;
168 config.indicesOptions = opt.indicesOpt;
169 config.useFloat16LookupTables = opt.useFloat16;
170
171 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
172 gpuIndex.setNumProbes(opt.nprobe);
173
174 gpuIndex.add(opt.numAdd, addVecs.data());
175 cpuIndex.add(opt.numAdd, addVecs.data());
176
177 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
178 opt.numQuery, opt.dim, opt.k, opt.toString(),
179 opt.getCompareEpsilon(),
180 opt.getPctMaxDiff1(),
181 opt.getPctMaxDiffN());
182 }
183 }
184
185 TEST(TestGpuIndexIVFPQ, CopyTo) {
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::GpuIndexIVFPQConfig config;
194 config.device = opt.device;
195 config.usePrecomputedTables = opt.usePrecomputed;
196 config.indicesOptions = opt.indicesOpt;
197 config.useFloat16LookupTables = opt.useFloat16;
198
199 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res,
200 opt.dim,
201 opt.numCentroids,
202 opt.codes,
203 opt.bitsPerCode,
204 faiss::METRIC_L2,
205 config);
206 gpuIndex.setNumProbes(opt.nprobe);
207 gpuIndex.train(opt.numTrain, trainVecs.data());
208 gpuIndex.add(opt.numAdd, addVecs.data());
209
210 // Use garbage values to see if we overwrite them
211 faiss::IndexFlatL2 cpuQuantizer(1);
212 faiss::IndexIVFPQ cpuIndex(&cpuQuantizer, 1, 1, 1, 1);
213
214 gpuIndex.copyTo(&cpuIndex);
215
216 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
217 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
218
219 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
220 EXPECT_EQ(cpuIndex.d, opt.dim);
221 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
222 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
223 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
224 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
225 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
226 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
227
228 // Query both objects; results should be equivalent
229 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
230 opt.numQuery, opt.dim, opt.k, opt.toString(),
231 opt.getCompareEpsilon(),
232 opt.getPctMaxDiff1(),
233 opt.getPctMaxDiffN());
234 }
235
236 TEST(TestGpuIndexIVFPQ, CopyFrom) {
237 Options opt;
238 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
239 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
240
241 faiss::IndexFlatL2 coarseQuantizer(opt.dim);
242 faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
243 opt.codes, opt.bitsPerCode);
244 cpuIndex.nprobe = opt.nprobe;
245 cpuIndex.train(opt.numTrain, trainVecs.data());
246 cpuIndex.add(opt.numAdd, addVecs.data());
247
248 // Use garbage values to see if we overwrite them
249 faiss::gpu::StandardGpuResources res;
250 res.noTempMemory();
251
252 faiss::gpu::GpuIndexIVFPQConfig config;
253 config.device = opt.device;
254 config.usePrecomputedTables = opt.usePrecomputed;
255 config.indicesOptions = opt.indicesOpt;
256 config.useFloat16LookupTables = opt.useFloat16;
257
258 faiss::gpu::GpuIndexIVFPQ
259 gpuIndex(&res, 1, 1, 1, 1, faiss::METRIC_L2, config);
260 gpuIndex.setNumProbes(1);
261
262 gpuIndex.copyFrom(&cpuIndex);
263
264 // Make sure we are equivalent
265 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
266 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
267
268 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
269 EXPECT_EQ(cpuIndex.d, opt.dim);
270 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
271 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
272 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
273 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
274 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
275 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
276
277 // Query both objects; results should be equivalent
278 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
279 opt.numQuery, opt.dim, opt.k, opt.toString(),
280 opt.getCompareEpsilon(),
281 opt.getPctMaxDiff1(),
282 opt.getPctMaxDiffN());
283 }
284
285 TEST(TestGpuIndexIVFPQ, QueryNaN) {
286 Options opt;
287
288 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
289 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
290
291 faiss::gpu::StandardGpuResources res;
292 res.noTempMemory();
293
294 faiss::gpu::GpuIndexIVFPQConfig config;
295 config.device = opt.device;
296 config.usePrecomputedTables = opt.usePrecomputed;
297 config.indicesOptions = opt.indicesOpt;
298 config.useFloat16LookupTables = opt.useFloat16;
299
300 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res,
301 opt.dim,
302 opt.numCentroids,
303 opt.codes,
304 opt.bitsPerCode,
305 faiss::METRIC_L2,
306 config);
307
308 gpuIndex.setNumProbes(opt.nprobe);
309
310 gpuIndex.train(opt.numTrain, trainVecs.data());
311 gpuIndex.add(opt.numAdd, addVecs.data());
312
313 int numQuery = 10;
314 std::vector<float> nans(numQuery * opt.dim,
315 std::numeric_limits<float>::quiet_NaN());
316
317 std::vector<float> distances(numQuery * opt.k, 0);
318 std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
319
320 gpuIndex.search(numQuery,
321 nans.data(),
322 opt.k,
323 distances.data(),
324 indices.data());
325
326 for (int q = 0; q < numQuery; ++q) {
327 for (int k = 0; k < opt.k; ++k) {
328 EXPECT_EQ(indices[q * opt.k + k], -1);
329 EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
330 }
331 }
332 }
333
334 TEST(TestGpuIndexIVFPQ, AddNaN) {
335 Options opt;
336
337 faiss::gpu::StandardGpuResources res;
338 res.noTempMemory();
339
340 faiss::gpu::GpuIndexIVFPQConfig config;
341 config.device = opt.device;
342 config.usePrecomputedTables = opt.usePrecomputed;
343 config.indicesOptions = opt.indicesOpt;
344 config.useFloat16LookupTables = opt.useFloat16;
345
346 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res,
347 opt.dim,
348 opt.numCentroids,
349 opt.codes,
350 opt.bitsPerCode,
351 faiss::METRIC_L2,
352 config);
353
354 gpuIndex.setNumProbes(opt.nprobe);
355
356 int numNans = 10;
357 std::vector<float> nans(numNans * opt.dim,
358 std::numeric_limits<float>::quiet_NaN());
359
360 // Make one vector valid, which should actually add
361 for (int i = 0; i < opt.dim; ++i) {
362 nans[i] = 0.0f;
363 }
364
365 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
366 gpuIndex.train(opt.numTrain, trainVecs.data());
367
368 // should not crash
369 EXPECT_EQ(gpuIndex.ntotal, 0);
370 gpuIndex.add(numNans, nans.data());
371
372 // Only the single valid vector should have added
373 EXPECT_EQ(gpuIndex.ntotal, 1);
374
375 std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
376 std::vector<float> distance(opt.numQuery * opt.k, 0);
377 std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
378
379 // should not crash
380 gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
381 distance.data(), indices.data());
382
383 }
384
385 TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
386 // Construct on a random device to test multi-device, if we have
387 // multiple devices
388 int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
389
390 if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
391 return;
392 }
393
394 int dim = 256;
395
396 int numCentroids = 1024;
397 // Unfortunately it would take forever to add 24 GB in IVFPQ data,
398 // so just perform a small test with data allocated in the unified
399 // memory address space
400 size_t numAdd = 100000;
401 size_t numTrain = numCentroids * 40;
402 int numQuery = 10;
403 int k = 10;
404 int nprobe = 8;
405 int codes = 8;
406 int bitsPerCode = 8;
407
408 std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
409 std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
410
411 faiss::IndexFlatL2 quantizer(dim);
412 faiss::IndexIVFPQ cpuIndex(&quantizer, dim, numCentroids, codes, bitsPerCode);
413
414 cpuIndex.train(numTrain, trainVecs.data());
415 cpuIndex.add(numAdd, addVecs.data());
416 cpuIndex.nprobe = nprobe;
417
418 faiss::gpu::StandardGpuResources res;
419 res.noTempMemory();
420
421 faiss::gpu::GpuIndexIVFPQConfig config;
422 config.device = device;
423 config.memorySpace = faiss::gpu::MemorySpace::Unified;
424
425 faiss::gpu::GpuIndexIVFPQ gpuIndex(&res,
426 dim,
427 numCentroids,
428 codes,
429 bitsPerCode,
430 faiss::METRIC_L2,
431 config);
432 gpuIndex.copyFrom(&cpuIndex);
433 gpuIndex.setNumProbes(nprobe);
434
435 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
436 numQuery, dim, k, "Unified Memory",
437 0.015f,
438 0.1f,
439 0.015f);
440 }
441
442 int main(int argc, char** argv) {
443 testing::InitGoogleTest(&argc, argv);
444
445 // just run with a fixed test seed
446 faiss::gpu::setTestSeed(100);
447
448 return RUN_ALL_TESTS();
449 }
faiss::gpu::GpuIndexIVFPQConfig::usePrecomputedTables
bool usePrecomputedTables
Definition: GpuIndexIVFPQ.h:35
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:76
faiss::IndexFlatL2
Definition: IndexFlat.h:78
faiss::gpu::GpuIndexIVFPQConfig::useFloat16LookupTables
bool useFloat16LookupTables
Definition: GpuIndexIVFPQ.h:31
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::GpuIndexIVFPQConfig
Definition: GpuIndexIVFPQ.h:22
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:35
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:31
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:22
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:39
faiss::gpu::GpuIndexIVFConfig::indicesOptions
IndicesOptions indicesOptions
Index storage options for the GPU.
Definition: GpuIndexIVF.h:30
Options
Definition: TestGpuIndexIVFFlat.cpp:27