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