Faiss
 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::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:31
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
Generated by   doxygen 1.8.5