Faiss
 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::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44
Generated by   doxygen 1.8.5