Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:72
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:30
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:24
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:25
Options
Definition: TestGpuIndexIVFFlat.cpp:27