Faiss benchmark sample code https://github.com/facebookresearch/faiss

faiss_benchmark_sample.cpp

#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;
    }
}