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