Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
demo_ivfpq_indexing_gpu.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
13 #include <cmath>
14 #include <cstdio>
15 #include <cstdlib>
16
17 #include <sys/time.h>
18
19
20 #include "../StandardGpuResources.h"
21 #include "../GpuIndexIVFPQ.h"
22
23 #include "../GpuAutoTune.h"
24 #include "../../index_io.h"
25
26 double elapsed ()
27 {
28 struct timeval tv;
29 gettimeofday (&tv, NULL);
30 return tv.tv_sec + tv.tv_usec * 1e-6;
31 }
32
33
34 int main ()
35 {
36
37 double t0 = elapsed();
38
39 // dimension of the vectors to index
40 int d = 128;
41
42 // size of the database we plan to index
43 size_t nb = 200 * 1000;
44
45 // make a set of nt training vectors in the unit cube
46 // (could be the database)
47 size_t nt = 100 * 1000;
48
49 int dev_no = 0;
50 /*
51 printf ("[%.3f s] Begin d=%d nb=%ld nt=%nt dev_no=%d\n",
52 elapsed() - t0, d, nb, nt, dev_no);
53 */
54 // a reasonable number of centroids to index nb vectors
55 int ncentroids = int (4 * sqrt (nb));
56
57 faiss::gpu::StandardGpuResources resources;
58
59
60 // the coarse quantizer should not be dealloced before the index
61 // 4 = nb of bytes per code (d must be a multiple of this)
62 // 8 = nb of bits per sub-code (almost always 8)
63 faiss::gpu::GpuIndexIVFPQ index (
64 &resources, dev_no, d,
65 ncentroids, 4, 8, true,
66 faiss::gpu::INDICES_64_BIT,
67 false,
68 faiss::METRIC_L2);
69
70
71 { // training
72 printf ("[%.3f s] Generating %ld vectors in %dD for training\n",
73 elapsed() - t0, nt, d);
74
75 std::vector <float> trainvecs (nt * d);
76 for (size_t i = 0; i < nt * d; i++) {
77 trainvecs[i] = drand48();
78 }
79
80 printf ("[%.3f s] Training the index\n",
81 elapsed() - t0);
82 index.verbose = true;
83
84 index.train (nt, trainvecs.data());
85 }
86
87 { // I/O demo
88 const char *outfilename = "/tmp/index_trained.faissindex";
89 printf ("[%.3f s] storing the pre-trained index to %s\n",
90 elapsed() - t0, outfilename);
91
92 faiss::Index * cpu_index = faiss::gpu::index_gpu_to_cpu (&index);
93
94 write_index (cpu_index, outfilename);
95
96 delete cpu_index;
97 }
98
99 size_t nq;
100 std::vector<float> queries;
101
102 { // populating the database
103 printf ("[%.3f s] Building a dataset of %ld vectors to index\n",
104 elapsed() - t0, nb);
105
106 std::vector <float> database (nb * d);
107 for (size_t i = 0; i < nb * d; i++) {
108 database[i] = drand48();
109 }
110
111 printf ("[%.3f s] Adding the vectors to the index\n",
112 elapsed() - t0);
113
114 index.add (nb, database.data());
115
116 printf ("[%.3f s] done\n", elapsed() - t0);
117
118 // remember a few elements from the database as queries
119 int i0 = 1234;
120 int i1 = 1243;
121
122 nq = i1 - i0;
123 queries.resize (nq * d);
124 for (int i = i0; i < i1; i++) {
125 for (int j = 0; j < d; j++) {
126 queries [(i - i0) * d + j] = database [i * d + j];
127 }
128 }
129
130 }
131
132 { // searching the database
133 int k = 5;
134 printf ("[%.3f s] Searching the %d nearest neighbors "
135 "of %ld vectors in the index\n",
136 elapsed() - t0, k, nq);
137
138 std::vector<faiss::Index::idx_t> nns (k * nq);
139 std::vector<float> dis (k * nq);
140
141 index.search (nq, queries.data(), k, dis.data(), nns.data());
142
143 printf ("[%.3f s] Query results (vector ids, then distances):\n",
144 elapsed() - t0);
145
146 for (int i = 0; i < nq; i++) {
147 printf ("query %2d: ", i);
148 for (int j = 0; j < k; j++) {
149 printf ("%7ld ", nns[j + i * k]);
150 }
151 printf ("\n dis: ");
152 for (int j = 0; j < k; j++) {
153 printf ("%7g ", dis[j + i * k]);
154 }
155 printf ("\n");
156 }
157
158 printf ("note that the nearest neighbor is not at "
159 "distance 0 due to quantization errors\n");
160 }
161
162 return 0;
163 }
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:24
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:25
faiss::Index
Definition: Index.h:61