Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
Faiss benchmark sample code https://github.com/facebookresearch/faiss
#include <iostream>
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/GpuIndexIVFPQ.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/index_io.h>
#include <faiss/utils.h>
#include <boost/program_options.hpp>
using namespace std;
namespace po = boost::program_options;
unique_ptr<float[]> makeData(int dim, int n) {
unique_ptr<float[]> data(new float[dim * n]);
faiss::float_rand(data.get(), dim*n, 1);
return data;
}
int main(int argc, char** argv) {
try {
po::options_description desc("benchmark program option");
desc.add_options()
("help,h", "benchmark program")
("gpu,g", "using gpu [default cpu mode]")
("database,d", po::value<int>()->default_value(100000), "database size")
("query,q", po::value<int>()->default_value(10000), "number of queries")
("nk,k", po::value<int>()->default_value(5), "number of k")
("verbose,v", "verbose output");
po::variables_map vm;
po::store(po::parse_command_line(argc, argv, desc), vm);
po::notify(vm);
if (vm.count("help")) {
cout << desc << endl;
return 1;
}
const int d = 128; // dimension
const int nb = vm["database"].as<int>(); // database size
const int nq = vm["query"].as<int>(); // number of queries
unique_ptr<float[]> xb = makeData(d, nb);
unique_ptr<float[]> xq = makeData(d, nq);
const int nlist = 100;
const int k = vm["nk"].as<int>();
const int m = 8;
const int nbits = 8;
cout << "database size: " << nb << endl
<< "number of queries: " << nq << endl
<< "k: " << k << endl;
double start, end;
if (!vm.count("gpu")) {
cout << "<--- CPU mode --->" << endl;
faiss::IndexFlatL2 quantizer(d);
faiss::IndexIVFPQ index(&quantizer, d, nlist, m, nbits);
if (vm.count("verbose")) {
index.verbose = true;
}
start = faiss::getmillisecs();
index.train(nb, xb.get());
end = faiss::getmillisecs();
cout << "training the index: " << (end - start)/1000 << " sec" << endl;
start = faiss::getmillisecs();
index.add(nb, xb.get());
end = faiss::getmillisecs();
cout << "adding the vectors to the index: " << (end - start)/1000 << " sec" << endl;
unique_ptr<long[]> I(new long[k * nq]);
unique_ptr<float[]> D(new float[k * nq]);
//index.nprobe = 100;
cout << "nprobe: " << index.nprobe << endl;
double start = faiss::getmillisecs();
index.search(nq, xq.get(), k, D.get(), I.get());
double end = faiss::getmillisecs();
cout << "searching the " << k << " nearest neighbors: "
<< (end - start)/1000 << " sec" << endl;
cout << faiss::get_mem_usage_kb()/1024 << " MB" << endl;
if (vm.count("verbose")) {
cout << "I=" << endl;
for(int i = 0; i < 10; i++) {
for(int j = 0; j < k; j++) printf("%5ld ", I[i * k + j]);
cout << endl;
}
}
} else {
cout << "<--- GPU mode --->" << endl;
faiss::gpu::StandardGpuResources res;
size_t memSize = 1073741824; // 1024MB (1024 * 1024 * 1024)
res.setTempMemory(memSize);
faiss::gpu::GpuIndexIVFPQ index(
&res, 0, d, nlist, m, nbits,
true, faiss::gpu::INDICES_32_BIT, true, faiss::METRIC_L2
);
faiss::gpu::DeviceMemory& deviceMemory = res.getMemoryManager(0);
if (vm.count("verbose")) {
index.verbose = true;
}
start = faiss::getmillisecs();
index.train(nb, xb.get());
end = faiss::getmillisecs();
cout << "training the index: " << (end - start)/1000 << " sec" << endl;
start = faiss::getmillisecs();
index.add(nb, xb.get());
end = faiss::getmillisecs();
cout << "adding the vectors to the index: " << (end - start)/1000 << " sec" << endl;
unique_ptr<long[]> I(new long[k * nq]);
unique_ptr<float[]> D(new float[k * nq]);
// index.setNumProbes(10);
cout << "nprobe: " << index.getNumProbes() << endl;
start = faiss::getmillisecs();
index.search(nq, xb.get(), k, D.get(), I.get());
end = faiss::getmillisecs();
cout << "searching the " << k << " nearest neighbors: "
<< (end - start)/1000 << " sec" << endl;
cout << faiss::get_mem_usage_kb()/1024 << " MB" << endl;
if (vm.count("verbose")) {
cout << "I=" << endl;
for(int i = 0; i < 10; i++) {
for(int j = 0; j < k; j++) printf("%5ld ", I[i * k + j]);
cout << endl;
}
cout << deviceMemory.getSizeAvailable()/1024/1024 << " MB" << endl;
long highWaterMark = deviceMemory.getHighWaterCudaMalloc();
cout << "HWM: " << highWaterMark/1024/1024 << " MB" << endl;
cout << deviceMemory.toString() << endl;
}
}
return 0;
} catch(const po::error& e) {
cerr << e.what() << endl;
return -1;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment