#89 code format (b6bf40e5) · Commits · SUMMER2020 / students / proj-2018085

core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp

+136 −146

Original line number	Diff line number	Diff line
		@@ -17,28 +17,28 @@

		#include <gtest/gtest.h>

		#include <cassert>
		#include <cmath>
		#include <cstdio>
		#include <cstdlib>
		#include <cassert>
		#include <cstring>

		#include <faiss/AutoTune.h>
		#include <faiss/Index.h>
		#include <faiss/IndexIVF.h>
		#include <faiss/gpu/StandardGpuResources.h>
		#include <faiss/gpu/GpuAutoTune.h>
		#include <faiss/gpu/GpuIndexFlat.h>
		#include <faiss/gpu/StandardGpuResources.h>
		#include <faiss/index_io.h>
		#include <faiss/utils.h>

		#include <hdf5.h>

		#include <vector>
		#include <sys/stat.h>
		#include <sys/time.h>
		#include <sys/types.h>
		#include <unistd.h>
		#include <vector>

		/*****************************************************
		* To run this test, please download the HDF5 from
		@@ -46,17 +46,15 @@
		* and install it to /usr/local/hdf5 .
		*****************************************************/

		double elapsed() {
		double
		elapsed() {
		struct timeval tv;
		gettimeofday(&tv, nullptr);
		return tv.tv_sec + tv.tv_usec * 1e-6;
		}

		void* hdf5_read(const char *file_name,
		const char *dataset_name,
		H5T_class_t dataset_class,
		size_t &d_out,
		size_t &n_out) {
		void*
		hdf5_read(const char* file_name, const char* dataset_name, H5T_class_t dataset_class, size_t& d_out, size_t& n_out) {
		hid_t file, dataset, datatype, dataspace, memspace;
		H5T_class_t t_class; /* data type class */
		H5T_order_t order; /* data order */
		@@ -149,9 +147,8 @@ void* hdf5_read(const char *file_name,
		return data_out;
		}

		std::string get_index_file_name(const std::string& ann_test_name,
		const std::string& index_key,
		int32_t data_loops) {
		std::string
		get_index_file_name(const std::string& ann_test_name, const std::string& index_key, int32_t data_loops) {
		size_t pos = index_key.find_first_of(',', 0);
		std::string file_name = ann_test_name;
		file_name = file_name + "_" + index_key.substr(0, pos) + "_" + index_key.substr(pos + 1);
		@@ -159,17 +156,19 @@ std::string get_index_file_name(const std::string& ann_test_name,
		return file_name;
		}

		bool parse_ann_test_name(const std::string& ann_test_name,
		size_t &dim,
		faiss::MetricType &metric_type) {
		bool
		parse_ann_test_name(const std::string& ann_test_name, size_t& dim, faiss::MetricType& metric_type) {
		size_t pos1, pos2;

		if (ann_test_name.empty()) return false;
		if (ann_test_name.empty())
		return false;

		pos1 = ann_test_name.find_first_of('-', 0);
		if (pos1 == std::string::npos) return false;
		if (pos1 == std::string::npos)
		return false;
		pos2 = ann_test_name.find_first_of('-', pos1 + 1);
		if (pos2 == std::string::npos) return false;
		if (pos2 == std::string::npos)
		return false;

		dim = std::stoi(ann_test_name.substr(pos1 + 1, pos2 - pos1 - 1));
		std::string metric_str = ann_test_name.substr(pos2 + 1);
		@@ -184,9 +183,8 @@ bool parse_ann_test_name(const std::string& ann_test_name,
		return true;
		}

		void test_ann_hdf5(const std::string& ann_test_name,
		const std::string& index_key,
		int32_t index_add_loops,
		void
		test_ann_hdf5(const std::string& ann_test_name, const std::string& index_key, int32_t index_add_loops,
		const std::vector<size_t>& nprobes) {
		double t0 = elapsed();

		@@ -207,8 +205,7 @@ void test_ann_hdf5(const std::string& ann_test_name,
		try {
		index = faiss::read_index(index_file_name.c_str());
		d = dim;
		}
		catch (...) {
		} catch (...) {
		printf("Cannot read index file: %s\n", index_file_name.c_str());

		printf("[%.3f s] Loading train set\n", elapsed() - t0);
		@@ -217,8 +214,7 @@ void test_ann_hdf5(const std::string& ann_test_name,
		float* xb = (float*)hdf5_read(ann_file_name.c_str(), "train", H5T_FLOAT, d, nb);
		assert(d == dim \|\| !"dataset does not have correct dimension");

		printf ("[%.3f s] Preparing index \"%s\" d=%ld\n",
		elapsed() - t0, index_key.c_str(), d);
		printf("[%.3f s] Preparing index \"%s\" d=%ld\n", elapsed() - t0, index_key.c_str(), d);

		index = faiss::index_factory(d, index_key.c_str(), metric_type);

		@@ -267,7 +263,6 @@ void test_ann_hdf5(const std::string& ann_test_name,
		}

		for (auto nprobe : nprobes) {

		faiss::ParameterSpace params;

		printf("[%.3f s] Setting parameter configuration 'nprobe=%lu' on index\n", elapsed() - t0, nprobe);
		@@ -309,11 +304,9 @@ void test_ann_hdf5(const std::string& ann_test_name,
		}
		}
		}
		printf("nq = %4ld, k = %4ld, elapse = %fs (quant = %fs, search = %fs), R@ = %.4f\n",
		t_nq, t_k, (t_end - t_start),
		faiss::indexIVF_stats.quantization_time / 1000,
		faiss::indexIVF_stats.search_time / 1000,
		(hit / float(t_nq * k / index_add_loops)));
		printf("nq = %4ld, k = %4ld, elapse = %fs (quant = %fs, search = %fs), R@ = %.4f\n", t_nq, t_k,
		(t_end - t_start), faiss::indexIVF_stats.quantization_time / 1000,
		faiss::indexIVF_stats.search_time / 1000, (hit / float(t_nq * k / index_add_loops)));
		}
		}
		printf("============================================================================================\n");
		@@ -333,9 +326,12 @@ void test_ann_hdf5(const std::string& ann_test_name,
		int gt_nn = gt[i * k];
		for (int j = 0; j < k; j++) {
		if (I[i * k + j] == gt_nn) {
		if(j < 1) n_1++;
		if(j < 10) n_10++;
		if(j < 100) n_100++;
		if (j < 1)
		n_1++;
		if (j < 10)
		n_10++;
		if (j < 100)
		n_100++;
		}
		}
		}
		@@ -356,9 +352,8 @@ void test_ann_hdf5(const std::string& ann_test_name,
		}

		#ifdef CUSTOMIZATION
		void test_ivfsq8h_gpu(const std::string& ann_test_name,
		int32_t index_add_loops,
		const std::vector<size_t>& nprobes){
		void
		test_ivfsq8h_gpu(const std::string& ann_test_name, int32_t index_add_loops, const std::vector<size_t>& nprobes) {
		double t0 = elapsed();

		const std::string ann_file_name = ann_test_name + ".hdf5";
		@@ -383,8 +378,7 @@ void test_ivfsq8h_gpu(const std::string& ann_test_name,
		try {
		cpu_index = faiss::read_index(index_file_name.c_str());
		d = dim;
		}
		catch (...){
		} catch (...) {
		printf("Cannot read index file: %s\n", index_file_name.c_str());

		printf("[%.3f s] Loading train set\n", elapsed() - t0);
		@@ -459,8 +453,7 @@ void test_ivfsq8h_gpu(const std::string& ann_test_name,
		}

		for (auto nprobe : nprobes) {
		printf ("[%.3f s] Setting parameter configuration 'nprobe=%lu' on index\n",
		elapsed() - t0, nprobe);
		printf("[%.3f s] Setting parameter configuration 'nprobe=%lu' on index\n", elapsed() - t0, nprobe);

		auto ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
		ivf_index->nprobe = nprobe;
		@@ -503,11 +496,9 @@ void test_ivfsq8h_gpu(const std::string& ann_test_name,
		}
		}
		}
		printf("nq = %4ld, k = %4ld, elapse = %fs (quant = %fs, search = %fs), R@ = %.4f\n",
		t_nq, t_k, (t_end - t_start),
		faiss::indexIVF_stats.quantization_time / 1000,
		faiss::indexIVF_stats.search_time / 1000,
		(hit / float(t_nq * k / index_add_loops)));
		printf("nq = %4ld, k = %4ld, elapse = %fs (quant = %fs, search = %fs), R@ = %.4f\n", t_nq, t_k,
		(t_end - t_start), faiss::indexIVF_stats.quantization_time / 1000,
		faiss::indexIVF_stats.search_time / 1000, (hit / float(t_nq * k / index_add_loops)));
		}
		}
		printf("============================================================================================\n");
		@@ -553,4 +544,3 @@ TEST(FAISSTEST, BENCHMARK) {
		test_ivfsq8h_gpu("glove-200-angular", 1, {8, 128});
		#endif
		}