Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
CompareIVFFlat.cu
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9
10 #include "../../IndexIVFFlat.h"
11 #include "../../index_io.h"
12 #include "../../utils.h"
13
14 #include "../GpuIndexIVFFlat.h"
15 #include "IndexWrapper.h"
16 #include "../test/TestUtils.h"
17 #include "../utils/DeviceTensor.cuh"
18 #include "../utils/DeviceUtils.h"
19 #include "../utils/HostTensor.cuh"
20 #include "../utils/Timer.h"
21 #include <gflags/gflags.h>
22 #include <map>
23 #include <memory>
24 #include <vector>
25
26 #include <cuda_profiler_api.h>
27
28 DEFINE_int32(nprobe, 5, "number of coarse centroids to probe");
29 DEFINE_int32(k, 3, "final number of closest results returned");
30 DEFINE_int32(num_queries, 3, "number of query vectors");
31 DEFINE_string(in, "/home/jhj/local/index.out", "index file for input");
32 DEFINE_bool(diff, true, "show exact distance + index output discrepancies");
33 DEFINE_bool(use_float16, false, "use encodings in float16");
34 DEFINE_bool(use_float16_coarse, false, "coarse quantizer in float16");
35 DEFINE_int64(seed, -1, "specify random seed");
36 DEFINE_int32(num_gpus, 1, "number of gpus to use");
37 DEFINE_int32(index, 2, "0 = no indices on GPU; 1 = 32 bit, 2 = 64 bit on GPU");
38
39 using namespace faiss::gpu;
40
41 int main(int argc, char** argv) {
42 gflags::ParseCommandLineFlags(&argc, &argv, true);
43
44 cudaProfilerStop();
45
46 auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
47 printf("using seed %ld\n", seed);
48
49 auto numQueries = FLAGS_num_queries;
50
51 auto index = std::unique_ptr<faiss::IndexIVFFlat>(
52 dynamic_cast<faiss::IndexIVFFlat*>(faiss::read_index(FLAGS_in.c_str())));
53 FAISS_ASSERT((bool) index);
54 index->nprobe = FLAGS_nprobe;
55
56 auto dim = index->d;
57
58 printf("Database: dim %d num vecs %ld\n", dim, index->ntotal);
59 printf("Coarse centroids: %ld\n", index->quantizer->ntotal);
60 printf("L2 lookup: %d queries, nprobe %d, total k %d\n",
61 numQueries, FLAGS_nprobe, FLAGS_k);
62 printf("float16 coarse quantizer %s\n",
63 FLAGS_use_float16_coarse ? "enabled" : "disabled");
64 printf("float16 encoding %s\n",
65 FLAGS_use_float16 ? "enabled" : "disabled");
66
67 // Convert to GPU index
68 printf("Copying index to %d GPU(s)...\n", FLAGS_num_gpus);
69
70 auto initFn = [&index](faiss::gpu::GpuResources* res, int dev) ->
71 std::unique_ptr<faiss::gpu::GpuIndexIVFFlat> {
72 GpuIndexIVFFlatConfig config;
73 config.device = dev;
74 config.indicesOptions = (faiss::gpu::IndicesOptions) FLAGS_index;
75 config.flatConfig.useFloat16 = FLAGS_use_float16_coarse;
76 config.useFloat16IVFStorage = FLAGS_use_float16;
77
78 auto p = std::unique_ptr<faiss::gpu::GpuIndexIVFFlat>(
79 new faiss::gpu::GpuIndexIVFFlat(res,
80 index->d,
81 index->nlist,
82 index->metric_type,
83 config));
84 p->copyFrom(index.get());
85 return p;
86 };
87
88 IndexWrapper<faiss::gpu::GpuIndexIVFFlat> gpuIndex(FLAGS_num_gpus, initFn);
89 gpuIndex.setNumProbes(FLAGS_nprobe);
90 printf("copy done\n");
91
92 // Build query vectors
93 HostTensor<float, 2, true> cpuQuery({numQueries, dim});
94 faiss::float_rand(cpuQuery.data(), cpuQuery.numElements(), seed);
95
96 // Time faiss CPU
97 HostTensor<float, 2, true> cpuDistances({numQueries, FLAGS_k});
98 HostTensor<faiss::Index::idx_t, 2, true> cpuIndices({numQueries, FLAGS_k});
99
100 float cpuTime = 0.0f;
101
102 {
103 CpuTimer timer;
104 index->search(numQueries,
105 cpuQuery.data(),
106 FLAGS_k,
107 cpuDistances.data(),
108 cpuIndices.data());
109
110 cpuTime = timer.elapsedMilliseconds();
111 }
112
113 printf("CPU time %.3f ms\n", cpuTime);
114
115 HostTensor<float, 2, true> gpuDistances({numQueries, FLAGS_k});
116 HostTensor<faiss::Index::idx_t, 2, true> gpuIndices({numQueries, FLAGS_k});
117
118 CUDA_VERIFY(cudaProfilerStart());
119 faiss::gpu::synchronizeAllDevices();
120
121 float gpuTime = 0.0f;
122
123 // Time GPU
124 {
125 CpuTimer timer;
126
127 gpuIndex.getIndex()->search(cpuQuery.getSize(0),
128 cpuQuery.data(),
129 FLAGS_k,
130 gpuDistances.data(),
131 gpuIndices.data());
132
133 // There is a device -> host copy above, so no need to time
134 // additional synchronization with the GPU
135 gpuTime = timer.elapsedMilliseconds();
136 }
137
138 CUDA_VERIFY(cudaProfilerStop());
139 printf("GPU time %.3f ms\n", gpuTime);
140
141 compareLists(cpuDistances.data(), cpuIndices.data(),
142 gpuDistances.data(), gpuIndices.data(),
143 numQueries, FLAGS_k,
144 "", true, FLAGS_diff, false);
145
146 CUDA_VERIFY(cudaDeviceSynchronize());
147 // printf("\ncudaMalloc usage %zd\n",
148 // resources.getMemoryManager().getHighWaterCudaMalloc());
149
150 return 0;
151 }
faiss::gpu::CpuTimer::elapsedMilliseconds
float elapsedMilliseconds()
Returns elapsed time in milliseconds.
Definition: Timer.cpp:51
faiss::gpu::GpuIndexIVFFlatConfig::useFloat16IVFStorage
bool useFloat16IVFStorage
Definition: GpuIndexIVFFlat.h:28
faiss::gpu::GpuResources
Definition: GpuResources.h:22
faiss::gpu::CpuTimer
CPU wallclock elapsed timer.
Definition: Timer.h:41
faiss::IndexIVFFlat
Definition: IndexIVFFlat.h:25
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:34
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::GpuIndexIVFConfig::flatConfig
GpuIndexFlatConfig flatConfig
Configuration for the coarse quantizer object.
Definition: GpuIndexIVF.h:33
faiss::gpu::HostTensor
Definition: HostTensor.cuh:21
faiss::gpu::GpuIndexIVFFlatConfig
Definition: GpuIndexIVFFlat.h:21
faiss::gpu::GpuIndexIVFConfig::indicesOptions
IndicesOptions indicesOptions
Index storage options for the GPU.
Definition: GpuIndexIVF.h:30
faiss::gpu::GpuIndexIVFFlat
Definition: GpuIndexIVFFlat.h:33
faiss::gpu::IndexWrapper
Definition: IndexWrapper.h:24