Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PerfClustering.cpp
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #include "../../utils.h"
10 #include "../../Clustering.h"
11 #include "../GpuIndexFlat.h"
12 #include "../StandardGpuResources.h"
13 #include "IndexWrapper.h"
14 #include "../utils/DeviceUtils.h"
15 #include "../utils/Timer.h"
16 #include <gflags/gflags.h>
17 #include <memory>
18 #include <vector>
19
20 #include <cuda_profiler_api.h>
21
22 DEFINE_int32(num, 10000, "# of vecs");
23 DEFINE_int32(k, 100, "# of clusters");
24 DEFINE_int32(dim, 128, "# of dimensions");
25 DEFINE_int32(niter, 10, "# of iterations");
26 DEFINE_bool(L2_metric, true, "If true, use L2 metric. If false, use IP metric");
27 DEFINE_bool(use_float16, false, "use float16 vectors and math");
28 DEFINE_bool(transposed, false, "transposed vector storage");
29 DEFINE_bool(verbose, false, "turn on clustering logging");
30 DEFINE_int64(seed, -1, "specify random seed");
31 DEFINE_int32(num_gpus, 1, "number of gpus to use");
32 DEFINE_int64(min_paging_size, -1, "minimum size to use CPU -> GPU paged copies");
33 DEFINE_int64(pinned_mem, -1, "pinned memory allocation to use");
34 DEFINE_int32(max_points, -1, "max points per centroid");
35
36 using namespace faiss::gpu;
37
38 int main(int argc, char** argv) {
39 gflags::ParseCommandLineFlags(&argc, &argv, true);
40
41 cudaProfilerStop();
42
43 auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
44 printf("using seed %ld\n", seed);
45
46 std::vector<float> vecs((size_t) FLAGS_num * FLAGS_dim);
47 faiss::float_rand(vecs.data(), vecs.size(), seed);
48
49 printf("K-means metric %s dim %d centroids %d num train %d niter %d\n",
50 FLAGS_L2_metric ? "L2" : "IP",
51 FLAGS_dim, FLAGS_k, FLAGS_num, FLAGS_niter);
52 printf("float16 math %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
53 printf("transposed storage %s\n", FLAGS_transposed ? "enabled" : "disabled");
54 printf("verbose %s\n", FLAGS_verbose ? "enabled" : "disabled");
55
56 auto initFn = [](faiss::gpu::GpuResources* res, int dev) ->
57 std::unique_ptr<faiss::gpu::GpuIndexFlat> {
58 if (FLAGS_pinned_mem >= 0) {
59 ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
60 FLAGS_pinned_mem);
61 }
62
63 GpuIndexFlatConfig config;
64 config.device = dev;
65 config.useFloat16 = FLAGS_use_float16;
66 config.storeTransposed = FLAGS_transposed;
67
68 auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
69 FLAGS_L2_metric ?
70 (faiss::gpu::GpuIndexFlat*)
71 new faiss::gpu::GpuIndexFlatL2(res, FLAGS_dim, config) :
72 (faiss::gpu::GpuIndexFlat*)
73 new faiss::gpu::GpuIndexFlatIP(res, FLAGS_dim, config));
74
75 if (FLAGS_min_paging_size >= 0) {
76 p->setMinPagingSize(FLAGS_min_paging_size);
77 }
78 return p;
79 };
80
81 IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
82
83 CUDA_VERIFY(cudaProfilerStart());
84 faiss::gpu::synchronizeAllDevices();
85
86 float gpuTime = 0.0f;
87
88 faiss::ClusteringParameters cp;
89 cp.niter = FLAGS_niter;
90 cp.verbose = FLAGS_verbose;
91
92 if (FLAGS_max_points > 0) {
93 cp.max_points_per_centroid = FLAGS_max_points;
94 }
95
96 faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
97
98 // Time k-means
99 {
100 CpuTimer timer;
101
102 kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
103
104 // There is a device -> host copy above, so no need to time
105 // additional synchronization with the GPU
106 gpuTime = timer.elapsedMilliseconds();
107 }
108
109 CUDA_VERIFY(cudaProfilerStop());
110 printf("k-means time %.3f ms\n", gpuTime);
111
112 CUDA_VERIFY(cudaDeviceSynchronize());
113
114 return 0;
115 }
faiss::gpu::CpuTimer::elapsedMilliseconds
float elapsedMilliseconds()
Returns elapsed time in milliseconds.
Definition: Timer.cpp:50
faiss::ClusteringParameters::niter
int niter
clustering iterations
Definition: Clustering.h:23
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:46
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::CpuTimer
CPU wallclock elapsed timer.
Definition: Timer.h:40
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:33
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:25
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:25
faiss::Clustering
Definition: Clustering.h:56
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:21
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:158
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:133
faiss::ClusteringParameters
Definition: Clustering.h:22
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:52
faiss::ClusteringParameters::max_points_per_centroid
int max_points_per_centroid
to limit size of dataset
Definition: Clustering.h:33
faiss::gpu::IndexWrapper
Definition: IndexWrapper.h:23