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