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