Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexIVFFlat.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 "../../IndexIVF.h"
14 #include "../GpuIndexIVFFlat.h"
15 #include "../StandardGpuResources.h"
16 #include "../utils/DeviceUtils.h"
17 #include "../test/TestUtils.h"
18 #include <gtest/gtest.h>
19 #include <sstream>
20 #include <vector>
21
22 // FIXME: figure out a better way to test fp16
23 constexpr float kF16MaxRelErr = 0.3f;
24 constexpr float kF32MaxRelErr = 0.03f;
25
26
27 struct Options {
28 Options() {
29 numAdd = faiss::gpu::randVal(4000, 20000);
30 dim = faiss::gpu::randVal(64, 200);
31
32 numCentroids = std::sqrt((float) numAdd);
33 numTrain = numCentroids * 40;
34 nprobe = faiss::gpu::randVal(10, numCentroids);
35 numQuery = faiss::gpu::randVal(32, 100);
36 k = std::min(faiss::gpu::randVal(10, 30), numAdd / 40);
37 indicesOpt = faiss::gpu::randSelect({
38 faiss::gpu::INDICES_CPU,
39 faiss::gpu::INDICES_32_BIT,
40 faiss::gpu::INDICES_64_BIT});
41
42 device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
43 }
44
45 std::string toString() const {
46 std::stringstream str;
47 str << "IVFFlat device " << device
48 << " numVecs " << numAdd
49 << " dim " << dim
50 << " numCentroids " << numCentroids
51 << " nprobe " << nprobe
52 << " numQuery " << numQuery
53 << " k " << k
54 << " indicesOpt " << indicesOpt;
55
56 return str.str();
57 }
58
59 int numAdd;
60 int dim;
61 int numCentroids;
62 int numTrain;
63 int nprobe;
64 int numQuery;
65 int k;
66 int device;
67 faiss::gpu::IndicesOptions indicesOpt;
68 };
69
70 void queryTest(faiss::MetricType metricType,
71 bool useFloat16CoarseQuantizer,
72 bool useFloat16,
73 int dimOverride = -1) {
74 for (int tries = 0; tries < 3; ++tries) {
75 faiss::gpu::newTestSeed();
76
77 Options opt;
78 opt.dim = dimOverride != -1 ? dimOverride : opt.dim;
79
80 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
81 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
82
83 faiss::IndexFlatL2 quantizerL2(opt.dim);
84 faiss::IndexFlatIP quantizerIP(opt.dim);
85 faiss::Index* quantizer =
86 metricType == faiss::METRIC_L2 ?
87 (faiss::Index*) &quantizerL2 : (faiss::Index*) &quantizerIP;
88
89 faiss::IndexIVFFlat cpuIndex(quantizer,
90 opt.dim, opt.numCentroids, metricType);
91 cpuIndex.train(opt.numTrain, trainVecs.data());
92 cpuIndex.add(opt.numAdd, addVecs.data());
93 cpuIndex.nprobe = opt.nprobe;
94
95 faiss::gpu::StandardGpuResources res;
96 res.noTempMemory();
97
98 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
99 opt.device,
100 useFloat16CoarseQuantizer,
101 useFloat16,
102 cpuIndex.d,
103 cpuIndex.nlist,
104 opt.indicesOpt,
105 cpuIndex.metric_type);
106 gpuIndex.copyFrom(&cpuIndex);
107 gpuIndex.setNumProbes(opt.nprobe);
108
109 bool compFloat16 = useFloat16CoarseQuantizer || useFloat16;
110 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
111 opt.numQuery, opt.dim, opt.k, opt.toString(),
112 compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
113 // FIXME: the fp16 bounds are
114 // useless when math (the accumulator) is
115 // in fp16. Figure out another way to test
116 compFloat16 ? 0.99f : 0.1f,
117 compFloat16 ? 0.65f : 0.015f);
118 }
119 }
120
121 void addTest(faiss::MetricType metricType,
122 bool useFloat16CoarseQuantizer,
123 bool useFloat16) {
124 for (int tries = 0; tries < 5; ++tries) {
125 faiss::gpu::newTestSeed();
126
127 Options opt;
128
129 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
130 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
131
132 faiss::IndexFlatL2 quantizerL2(opt.dim);
133 faiss::IndexFlatIP quantizerIP(opt.dim);
134 faiss::Index* quantizer =
135 metricType == faiss::METRIC_L2 ?
136 (faiss::Index*) &quantizerL2 : (faiss::Index*) &quantizerIP;
137
138 faiss::IndexIVFFlat cpuIndex(quantizer,
139 opt.dim,
140 opt.numCentroids,
141 metricType);
142 cpuIndex.train(opt.numTrain, trainVecs.data());
143 cpuIndex.nprobe = opt.nprobe;
144
145 faiss::gpu::StandardGpuResources res;
146 res.noTempMemory();
147
148 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
149 opt.device,
150 useFloat16CoarseQuantizer,
151 useFloat16,
152 cpuIndex.d,
153 cpuIndex.nlist,
154 opt.indicesOpt,
155 cpuIndex.metric_type);
156 gpuIndex.copyFrom(&cpuIndex);
157 gpuIndex.setNumProbes(opt.nprobe);
158
159 cpuIndex.add(opt.numAdd, addVecs.data());
160 gpuIndex.add(opt.numAdd, addVecs.data());
161
162 bool compFloat16 = useFloat16CoarseQuantizer || useFloat16;
163 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
164 opt.numQuery, opt.dim, opt.k, opt.toString(),
165 compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
166 compFloat16 ? 0.70f : 0.1f,
167 compFloat16 ? 0.30f : 0.015f);
168 }
169 }
170
171 void copyToTest(bool useFloat16CoarseQuantizer,
172 bool useFloat16) {
173 faiss::gpu::newTestSeed();
174
175 Options opt;
176 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
177 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
178
179 faiss::gpu::StandardGpuResources res;
180 res.noTempMemory();
181
182 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
183 opt.device,
184 useFloat16CoarseQuantizer,
185 useFloat16,
186 opt.dim,
187 opt.numCentroids,
188 opt.indicesOpt,
189 faiss::METRIC_L2);
190 gpuIndex.train(opt.numTrain, trainVecs.data());
191 gpuIndex.add(opt.numAdd, addVecs.data());
192 gpuIndex.setNumProbes(opt.nprobe);
193
194 // use garbage values to see if we overwrite then
195 faiss::IndexFlatL2 cpuQuantizer(1);
196 faiss::IndexIVFFlat cpuIndex(&cpuQuantizer, 1, 1, faiss::METRIC_L2);
197 cpuIndex.nprobe = 1;
198
199 gpuIndex.copyTo(&cpuIndex);
200
201 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
202 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
203
204 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
205 EXPECT_EQ(cpuIndex.d, opt.dim);
206 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
207 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
208
209 // Query both objects; results should be equivalent
210 bool compFloat16 = useFloat16CoarseQuantizer || useFloat16;
211 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
212 opt.numQuery, opt.dim, opt.k, opt.toString(),
213 compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
214 compFloat16 ? 0.70f : 0.1f,
215 compFloat16 ? 0.30f : 0.015f);
216 }
217
218 void copyFromTest(bool useFloat16CoarseQuantizer,
219 bool useFloat16) {
220 faiss::gpu::newTestSeed();
221
222 Options opt;
223 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
224 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
225
226 faiss::IndexFlatL2 cpuQuantizer(opt.dim);
227 faiss::IndexIVFFlat cpuIndex(&cpuQuantizer,
228 opt.dim,
229 opt.numCentroids,
230 faiss::METRIC_L2);
231 cpuIndex.nprobe = opt.nprobe;
232 cpuIndex.train(opt.numTrain, trainVecs.data());
233 cpuIndex.add(opt.numAdd, addVecs.data());
234
235 // use garbage values to see if we overwrite then
236 faiss::gpu::StandardGpuResources res;
237 res.noTempMemory();
238
239 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
240 opt.device,
241 useFloat16CoarseQuantizer,
242 useFloat16,
243 1,
244 1,
245 opt.indicesOpt,
246 faiss::METRIC_L2);
247 gpuIndex.setNumProbes(1);
248
249 gpuIndex.copyFrom(&cpuIndex);
250
251 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
252 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
253
254 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
255 EXPECT_EQ(cpuIndex.d, opt.dim);
256 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
257 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
258
259 // Query both objects; results should be equivalent
260 bool compFloat16 = useFloat16CoarseQuantizer || useFloat16;
261 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
262 opt.numQuery, opt.dim, opt.k, opt.toString(),
263 compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
264 compFloat16 ? 0.70f : 0.1f,
265 compFloat16 ? 0.30f : 0.015f);
266 }
267
268 TEST(TestGpuIndexIVFFlat, Float32_32_Add_L2) {
269 addTest(faiss::METRIC_L2, false, false);
270 }
271
272 TEST(TestGpuIndexIVFFlat, Float32_32_Add_IP) {
273 addTest(faiss::METRIC_INNER_PRODUCT, false, false);
274 }
275
276 TEST(TestGpuIndexIVFFlat, Float32_16_Add_L2) {
277 addTest(faiss::METRIC_L2, false, true);
278 }
279
280 TEST(TestGpuIndexIVFFlat, Float32_16_Add_IP) {
281 addTest(faiss::METRIC_INNER_PRODUCT, false, true);
282 }
283
284 TEST(TestGpuIndexIVFFlat, Float16_32_Add_L2) {
285 addTest(faiss::METRIC_L2, true, false);
286 }
287
288 TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
289 addTest(faiss::METRIC_INNER_PRODUCT, true, false);
290 }
291
292 //
293 // General query tests
294 //
295
296 TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
297 queryTest(faiss::METRIC_L2, false, false);
298 }
299
300 TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
301 queryTest(faiss::METRIC_INNER_PRODUCT, false, false);
302 }
303
304 TEST(TestGpuIndexIVFFlat, Float16_Query_L2) {
305 queryTest(faiss::METRIC_L2, false, true);
306 }
307
308 TEST(TestGpuIndexIVFFlat, Float16_Query_IP) {
309 queryTest(faiss::METRIC_INNER_PRODUCT, false, true);
310 }
311
312 // float16 coarse quantizer
313
314 TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
315 queryTest(faiss::METRIC_L2, true, false);
316 }
317
318 TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
319 queryTest(faiss::METRIC_INNER_PRODUCT, true, false);
320 }
321
322 //
323 // There are IVF list scanning specializations for 64-d and 128-d that we
324 // make sure we explicitly test here
325 //
326
327 TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
328 queryTest(faiss::METRIC_L2, false, false, 64);
329 }
330
331 TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
332 queryTest(faiss::METRIC_INNER_PRODUCT, false, false, 64);
333 }
334
335 TEST(TestGpuIndexIVFFlat, Float16_Query_L2_64) {
336 queryTest(faiss::METRIC_L2, false, true, 64);
337 }
338
339 TEST(TestGpuIndexIVFFlat, Float16_Query_IP_64) {
340 queryTest(faiss::METRIC_INNER_PRODUCT, false, true, 64);
341 }
342
343 TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
344 queryTest(faiss::METRIC_L2, false, false, 128);
345 }
346
347 TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
348 queryTest(faiss::METRIC_INNER_PRODUCT, false, false, 128);
349 }
350
351 TEST(TestGpuIndexIVFFlat, Float16_Query_L2_128) {
352 queryTest(faiss::METRIC_L2, false, true, 128);
353 }
354
355 TEST(TestGpuIndexIVFFlat, Float16_Query_IP_128) {
356 queryTest(faiss::METRIC_INNER_PRODUCT, false, true, 128);
357 }
358
359 // For 256-d, only float16 is specialized
360
361 TEST(TestGpuIndexIVFFlat, Float16_Query_L2_256) {
362 queryTest(faiss::METRIC_L2, false, true, 256);
363 }
364
365 TEST(TestGpuIndexIVFFlat, Float16_Query_IP_256) {
366 queryTest(faiss::METRIC_INNER_PRODUCT, false, true, 256);
367 }
368
369 //
370 // Copy tests
371 //
372
373 TEST(TestGpuIndexIVFFlat, Float32_16_CopyTo) {
374 copyToTest(false, true);
375 }
376
377 TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
378 copyToTest(false, false);
379 }
380
381 //
382 // NaN tests
383 //
384
385 TEST(TestGpuIndexIVFFlat, QueryNaN) {
386 faiss::gpu::newTestSeed();
387
388 Options opt;
389
390 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
391 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
392
393 faiss::gpu::StandardGpuResources res;
394 res.noTempMemory();
395
396 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
397 opt.device,
398 faiss::gpu::randBool(),
399 faiss::gpu::randBool(),
400 opt.dim,
401 opt.numCentroids,
402 opt.indicesOpt,
403 faiss::METRIC_L2);
404 gpuIndex.setNumProbes(opt.nprobe);
405
406 gpuIndex.train(opt.numTrain, trainVecs.data());
407 gpuIndex.add(opt.numAdd, addVecs.data());
408
409 int numQuery = 10;
410 std::vector<float> nans(numQuery * opt.dim,
411 std::numeric_limits<float>::quiet_NaN());
412
413 std::vector<float> distances(numQuery * opt.k, 0);
414 std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
415
416 gpuIndex.search(numQuery,
417 nans.data(),
418 opt.k,
419 distances.data(),
420 indices.data());
421
422 for (int q = 0; q < numQuery; ++q) {
423 for (int k = 0; k < opt.k; ++k) {
424 EXPECT_EQ(indices[q * opt.k + k], -1);
425 EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
426 }
427 }
428 }
429
430 TEST(TestGpuIndexIVFFlat, AddNaN) {
431 faiss::gpu::newTestSeed();
432
433 Options opt;
434
435 faiss::gpu::StandardGpuResources res;
436 res.noTempMemory();
437
438 faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
439 opt.device,
440 faiss::gpu::randBool(),
441 faiss::gpu::randBool(),
442 opt.dim,
443 opt.numCentroids,
444 opt.indicesOpt,
445 faiss::METRIC_L2);
446 gpuIndex.setNumProbes(opt.nprobe);
447
448 int numNans = 10;
449 std::vector<float> nans(numNans * opt.dim,
450 std::numeric_limits<float>::quiet_NaN());
451
452 // Make one vector valid, which should actually add
453 for (int i = 0; i < opt.dim; ++i) {
454 nans[i] = 0.0f;
455 }
456
457 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
458 gpuIndex.train(opt.numTrain, trainVecs.data());
459
460 // should not crash
461 EXPECT_EQ(gpuIndex.ntotal, 0);
462 gpuIndex.add(numNans, nans.data());
463
464 // Only the single valid vector should have added
465 EXPECT_EQ(gpuIndex.ntotal, 1);
466
467 std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
468 std::vector<float> distance(opt.numQuery * opt.k, 0);
469 std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
470
471 // should not crash
472 gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
473 distance.data(), indices.data());
474
475 }
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:72
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::IndexIVFFlat
Definition: IndexIVF.h:113
faiss::IndexFlatIP
Definition: IndexFlat.h:75
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:24
faiss::Index
Definition: Index.h:61
Options
Definition: TestGpuIndexIVFFlat.cpp:27
faiss::gpu::GpuIndexIVFFlat
Definition: GpuIndexIVFFlat.h:25
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44