1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#include "../../IndexIVF.h"
#include "../../IndexIVFPQ.h"
#include "../../IndexFlat.h"
#include "../../index_io.h"
#include "../test/TestUtils.h"
#include <vector>
#include <gflags/gflags.h>
// For IVFPQ:
DEFINE_bool(ivfpq, false, "use IVFPQ encoding");
DEFINE_int32(codes, 4, "number of PQ codes per vector");
DEFINE_int32(bits_per_code, 8, "number of bits per PQ code");
// For IVFFlat:
DEFINE_bool(l2, true, "use L2 metric (versus IP metric)");
DEFINE_bool(ivfflat, false, "use IVF flat encoding");
// For both:
DEFINE_string(out, "/home/jhj/local/index.out", "index file for output");
DEFINE_int32(dim, 128, "vector dimension");
DEFINE_int32(num_coarse, 100, "number of coarse centroids");
DEFINE_int32(num, 100000, "total database size");
DEFINE_int32(num_train, -1, "number of database vecs to train on");
template <typename T>
void fillAndSave(T& index, int numTrain, int num, int dim) {
auto trainVecs = faiss::gpu::randVecs(numTrain, dim);
index.train(numTrain, trainVecs.data());
constexpr int kAddChunk = 1000000;
for (int i = 0; i < num; i += kAddChunk) {
int numRemaining = (num - i) < kAddChunk ? (num - i) : kAddChunk;
auto vecs = faiss::gpu::randVecs(numRemaining, dim);
printf("adding at %d: %d\n", i, numRemaining);
index.add(numRemaining, vecs.data());
}
faiss::write_index(&index, FLAGS_out.c_str());
}
int main(int argc, char** argv) {
google::ParseCommandLineFlags(&argc, &argv, true);
// Either ivfpq or ivfflat must be set
if ((FLAGS_ivfpq && FLAGS_ivfflat) ||
(!FLAGS_ivfpq && !FLAGS_ivfflat)) {
printf("must specify either ivfpq or ivfflat\n");
return 1;
}
auto dim = FLAGS_dim;
auto numCentroids = FLAGS_num_coarse;
auto num = FLAGS_num;
auto numTrain = FLAGS_num_train;
numTrain = numTrain == -1 ? std::max((num / 4), 1) : numTrain;
numTrain = std::min(num, numTrain);
if (FLAGS_ivfpq) {
faiss::IndexFlatL2 quantizer(dim);
faiss::IndexIVFPQ index(&quantizer, dim, numCentroids,
FLAGS_codes, FLAGS_bits_per_code);
index.verbose = true;
printf("IVFPQ: codes %d bits per code %d\n",
FLAGS_codes, FLAGS_bits_per_code);
printf("Lists: %d\n", numCentroids);
printf("Database: dim %d num vecs %d trained on %d\n", dim, num, numTrain);
printf("output file: %s\n", FLAGS_out.c_str());
fillAndSave(index, numTrain, num, dim);
} else if (FLAGS_ivfflat) {
faiss::IndexFlatL2 quantizerL2(dim);
faiss::IndexFlatIP quantizerIP(dim);
faiss::IndexFlat* quantizer = FLAGS_l2 ?
(faiss::IndexFlat*) &quantizerL2 :
(faiss::IndexFlat*) &quantizerIP;
faiss::IndexIVFFlat index(quantizer, dim, numCentroids,
FLAGS_l2 ? faiss::METRIC_L2 :
faiss::METRIC_INNER_PRODUCT);
printf("IVFFlat: metric %s\n", FLAGS_l2 ? "L2" : "IP");
printf("Lists: %d\n", numCentroids);
printf("Database: dim %d num vecs %d trained on %d\n", dim, num, numTrain);
printf("output file: %s\n", FLAGS_out.c_str());
fillAndSave(index, numTrain, num, dim);
}
return 0;
}