Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexIVFPQ.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 "../../IndexIVFPQ.h"
14 #include "../GpuIndexIVFPQ.h"
15 #include "../StandardGpuResources.h"
16 #include "../utils/DeviceUtils.h"
17 #include "../test/TestUtils.h"
18 #include <cmath>
19 #include <gtest/gtest.h>
20 #include <sstream>
21 #include <vector>
22 
23 
24 void pickEncoding(int& codes, int& dim) {
25  std::vector<int> codeSizes{
26  3, 4, 8, 12, 16, 20, 24,
27  28, 32, 40, 48, 56, 64, 96
28  };
29 
30  std::vector<int> dimSizes{4, 8, 16, 32};
31 
32  codes = codeSizes[faiss::gpu::randVal(0, codeSizes.size() - 1)];
33 
34  while (true) {
35  dim = codes * dimSizes[faiss::gpu::randVal(0, dimSizes.size() - 1)];
36 
37  // for such a small test, super-low or high dim is more likely to
38  // generate comparison errors
39  if (dim < 512 && dim >= 64) {
40  return;
41  }
42  }
43 }
44 
45 struct Options {
46  Options() {
47  numAdd = faiss::gpu::randVal(10000, 30000);
48  numCentroids = std::sqrt((float) numAdd);
49  numTrain = numCentroids * 40;
50 
51  pickEncoding(codes, dim);
52 
53  bitsPerCode = faiss::gpu::randVal(3, 8);
54  nprobe = std::min(faiss::gpu::randVal(40, 1000), numCentroids);
55  numQuery = faiss::gpu::randVal(32, 256);
56  k = std::min(faiss::gpu::randVal(10, 50), numAdd / 40);
57  usePrecomputed = faiss::gpu::randBool();
58  indicesOpt = faiss::gpu::randSelect({
59  faiss::gpu::INDICES_CPU,
60  faiss::gpu::INDICES_32_BIT,
61  faiss::gpu::INDICES_64_BIT});
62  if (codes > 48) {
63  // large codes can only fit using float16
64  useFloat16 = true;
65  } else {
66  useFloat16 = faiss::gpu::randBool();
67  }
68 
69  device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
70  }
71 
72  std::string toString() const {
73  std::stringstream str;
74  str << "IVFPQ device " << device
75  << " numVecs " << numAdd
76  << " dim " << dim
77  << " numCentroids " << numCentroids
78  << " codes " << codes
79  << " bitsPerCode " << bitsPerCode
80  << " nprobe " << nprobe
81  << " numQuery " << numQuery
82  << " k " << k
83  << " usePrecomputed " << usePrecomputed
84  << " indicesOpt " << indicesOpt
85  << " useFloat16 " << useFloat16;
86 
87  return str.str();
88  }
89 
90  float getCompareEpsilon() const {
91  return 0.015f;
92  }
93 
94  float getPctMaxDiff1() const {
95  return useFloat16 ? 0.30f : 0.10f;
96  }
97 
98  float getPctMaxDiffN() const {
99  return useFloat16 ? 0.05f : 0.015f;
100  }
101 
102  int numAdd;
103  int numCentroids;
104  int numTrain;
105  int codes;
106  int dim;
107  int bitsPerCode;
108  int nprobe;
109  int numQuery;
110  int k;
111  bool usePrecomputed;
112  faiss::gpu::IndicesOptions indicesOpt;
113  bool useFloat16;
114  int device;
115 };
116 
117 TEST(TestGpuIndexIVFPQ, Query) {
118  for (int tries = 0; tries < 5; ++tries) {
119  faiss::gpu::newTestSeed();
120 
121  Options opt;
122 
123  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
124  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
125 
126  faiss::IndexFlatL2 coarseQuantizer(opt.dim);
127  faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
128  opt.codes, opt.bitsPerCode);
129  cpuIndex.nprobe = opt.nprobe;
130  cpuIndex.train(opt.numTrain, trainVecs.data());
131  cpuIndex.add(opt.numAdd, addVecs.data());
132 
133  faiss::gpu::StandardGpuResources res;
134  res.noTempMemory();
135 
136  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
137  opt.indicesOpt,
138  opt.useFloat16,
139  &cpuIndex);
140  gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
141  gpuIndex.setNumProbes(opt.nprobe);
142 
143  faiss::gpu::compareIndices(cpuIndex, gpuIndex,
144  opt.numQuery, opt.dim, opt.k, opt.toString(),
145  opt.getCompareEpsilon(),
146  opt.getPctMaxDiff1(),
147  opt.getPctMaxDiffN());
148  }
149 }
150 
151 TEST(TestGpuIndexIVFPQ, Add) {
152  for (int tries = 0; tries < 5; ++tries) {
153  faiss::gpu::newTestSeed();
154 
155  Options opt;
156 
157  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
158  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
159 
160  faiss::IndexFlatL2 coarseQuantizer(opt.dim);
161  faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
162  opt.codes, opt.bitsPerCode);
163  cpuIndex.nprobe = opt.nprobe;
164  cpuIndex.train(opt.numTrain, trainVecs.data());
165 
166  faiss::gpu::StandardGpuResources res;
167  res.noTempMemory();
168 
169  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
170  opt.indicesOpt,
171  opt.useFloat16,
172  &cpuIndex);
173  gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
174  gpuIndex.setNumProbes(opt.nprobe);
175 
176  gpuIndex.add(opt.numAdd, addVecs.data());
177  cpuIndex.add(opt.numAdd, addVecs.data());
178 
179  faiss::gpu::compareIndices(cpuIndex, gpuIndex,
180  opt.numQuery, opt.dim, opt.k, opt.toString(),
181  opt.getCompareEpsilon(),
182  opt.getPctMaxDiff1(),
183  opt.getPctMaxDiffN());
184  }
185 }
186 
187 TEST(TestGpuIndexIVFPQ, CopyTo) {
188  faiss::gpu::newTestSeed();
189 
190  Options opt;
191  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
192  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
193 
194  faiss::gpu::StandardGpuResources res;
195  res.noTempMemory();
196 
197  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
198  opt.dim,
199  opt.numCentroids,
200  opt.codes,
201  opt.bitsPerCode,
202  opt.usePrecomputed,
203  opt.indicesOpt,
204  opt.useFloat16,
205  faiss::METRIC_L2);
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  faiss::gpu::newTestSeed();
238 
239  Options opt;
240  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
241  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
242 
243  faiss::IndexFlatL2 coarseQuantizer(opt.dim);
244  faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, opt.dim, opt.numCentroids,
245  opt.codes, opt.bitsPerCode);
246  cpuIndex.nprobe = opt.nprobe;
247  cpuIndex.train(opt.numTrain, trainVecs.data());
248  cpuIndex.add(opt.numAdd, addVecs.data());
249 
250  // Use garbage values to see if we overwrite them
251  faiss::gpu::StandardGpuResources res;
252  res.noTempMemory();
253 
254  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
255  1, 1, 1, 1, false, opt.indicesOpt,
256  opt.useFloat16, faiss::METRIC_L2);
257  gpuIndex.setNumProbes(1);
258 
259  gpuIndex.copyFrom(&cpuIndex);
260 
261  // Make sure we are equivalent
262  EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
263  EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
264 
265  EXPECT_EQ(cpuIndex.d, gpuIndex.d);
266  EXPECT_EQ(cpuIndex.d, opt.dim);
267  EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
268  EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
269  EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
270  EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
271  EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
272  EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
273 
274  // Query both objects; results should be equivalent
275  faiss::gpu::compareIndices(cpuIndex, gpuIndex,
276  opt.numQuery, opt.dim, opt.k, opt.toString(),
277  opt.getCompareEpsilon(),
278  opt.getPctMaxDiff1(),
279  opt.getPctMaxDiffN());
280 }
281 
282 TEST(TestGpuIndexIVFPQ, QueryNaN) {
283  faiss::gpu::newTestSeed();
284 
285  Options opt;
286 
287  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
288  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
289 
290  faiss::gpu::StandardGpuResources res;
291  res.noTempMemory();
292 
293  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
294  opt.dim,
295  opt.numCentroids,
296  opt.codes,
297  opt.bitsPerCode,
298  opt.usePrecomputed,
299  opt.indicesOpt,
300  opt.useFloat16,
301  faiss::METRIC_L2);
302 
303  gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
304  gpuIndex.setNumProbes(opt.nprobe);
305 
306  gpuIndex.train(opt.numTrain, trainVecs.data());
307  gpuIndex.add(opt.numAdd, addVecs.data());
308 
309  int numQuery = 10;
310  std::vector<float> nans(numQuery * opt.dim,
311  std::numeric_limits<float>::quiet_NaN());
312 
313  std::vector<float> distances(numQuery * opt.k, 0);
314  std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
315 
316  gpuIndex.search(numQuery,
317  nans.data(),
318  opt.k,
319  distances.data(),
320  indices.data());
321 
322  for (int q = 0; q < numQuery; ++q) {
323  for (int k = 0; k < opt.k; ++k) {
324  EXPECT_EQ(indices[q * opt.k + k], -1);
325  EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
326  }
327  }
328 }
329 
330 TEST(TestGpuIndexIVFPQ, AddNaN) {
331  faiss::gpu::newTestSeed();
332 
333  Options opt;
334 
335  faiss::gpu::StandardGpuResources res;
336  res.noTempMemory();
337 
338  faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, opt.device,
339  opt.dim,
340  opt.numCentroids,
341  opt.codes,
342  opt.bitsPerCode,
343  opt.usePrecomputed,
344  opt.indicesOpt,
345  opt.useFloat16,
346  faiss::METRIC_L2);
347 
348  gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
349  gpuIndex.setNumProbes(opt.nprobe);
350 
351  int numNans = 10;
352  std::vector<float> nans(numNans * opt.dim,
353  std::numeric_limits<float>::quiet_NaN());
354 
355  // Make one vector valid, which should actually add
356  for (int i = 0; i < opt.dim; ++i) {
357  nans[i] = 0.0f;
358  }
359 
360  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
361  gpuIndex.train(opt.numTrain, trainVecs.data());
362 
363  // should not crash
364  EXPECT_EQ(gpuIndex.ntotal, 0);
365  gpuIndex.add(numNans, nans.data());
366 
367  // Only the single valid vector should have added
368  EXPECT_EQ(gpuIndex.ntotal, 1);
369 
370  std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
371  std::vector<float> distance(opt.numQuery * opt.k, 0);
372  std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
373 
374  // should not crash
375  gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
376  distance.data(), indices.data());
377 
378 }
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:25
Generated by   doxygen 1.8.5