admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-29 06:55:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

add mix mode for faiss baseline

Browse Source 
make baseline compatible for different faiss version


Former-commit-id: 90f96b4d437136539e7aa805f58c49e2d8a23ed1
This commit is contained in:
Heisenberg 2019-10-31 14:46:13 +08:00
parent 90e961649c
commit c8da2c3d64
1 changed files with 88 additions and 58 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

146
core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp
View File

@ -17,31 +17,40 @@
#include <gtest/gtest.h>
#include <hdf5.h>
#include <sys/time.h>
#include <sys/types.h>
#include <unistd.h>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <vector>
#define USE_FAISS_V1_5_3 0
#if USE_FAISS_V1_5_3
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/utils.h>
#include <sys/stat.h>
#include <cstdlib>
#include <cstring>
#else
#include <faiss/gpu/GpuCloner.h>
#include <faiss/index_factory.h>
#include <faiss/utils/distances.h>
#endif
#include <faiss/AutoTune.h>
#include <faiss/Index.h>
#include <faiss/IndexIVF.h>
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/index_io.h>
#ifdef CUSTOMIZATION
#include <faiss/gpu/GpuIndexIVFSQHybrid.h>
#endif
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/index_io.h>
#include <faiss/utils.h>
#include <hdf5.h>
#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
@ -56,6 +65,8 @@ const char HDF5_DATASET_TEST[] = "test";
const char HDF5_DATASET_NEIGHBORS[] = "neighbors";
const char HDF5_DATASET_DISTANCES[] = "distances";
const int32_t GPU_DEVICE_IDX = 0;
enum QueryMode { MODE_CPU = 0, MODE_MIX, MODE_GPU };
double
@ -196,7 +207,7 @@ GetResultHitCount(const faiss::Index::idx_t* ground_index, const faiss::Index::i
for (int i = 0; i < nq; i++) {
// count the num of results exist in ground truth result set
// each result replicates INDEX_ADD_LOOPS times
for (int j_c = 0; j_c < ground_k; j_c++) {
for (int j_c = 0; j_c < k; j_c++) {
int r_c = index[i * k + j_c];
for (int j_g = 0; j_g < ground_k / index_add_loops; j_g++) {
if (ground_index[i * ground_k + j_g] == r_c) {
@ -239,7 +250,6 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
double t0 = elapsed();
const std::string ann_file_name = ann_test_name + HDF5_POSTFIX;
const int GPU_DEVICE_IDX = 0;
faiss::Index *cpu_index = nullptr, *gpu_index = nullptr;
faiss::distance_compute_blas_threshold = 800;
@ -249,30 +259,6 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
try {
printf("[%.3f s] Reading index file: %s\n", elapsed() - t0, index_file_name.c_str());
cpu_index = faiss::read_index(index_file_name.c_str());
if (mode != MODE_CPU) {
faiss::gpu::GpuClonerOptions option;
option.allInGpu = true;
faiss::IndexComposition index_composition;
index_composition.index = cpu_index;
index_composition.quantizer = nullptr;
switch (mode) {
case MODE_CPU:
assert(false);
break;
case MODE_MIX:
index_composition.mode = 1; // 0: all data, 1: copy quantizer, 2: copy data
break;
case MODE_GPU:
index_composition.mode = 0; // 0: all data, 1: copy quantizer, 2: copy data
break;
}
printf("[%.3f s] Cloning CPU index to GPU\n", elapsed() - t0);
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
}
} catch (...) {
size_t nb, d;
printf("[%.3f s] Loading HDF5 file: %s\n", elapsed() - t0, ann_file_name.c_str());
@ -289,6 +275,7 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
printf("[%.3f s] Cloning CPU index to GPU\n", elapsed() - t0);
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, cpu_index);
delete cpu_index;
printf("[%.3f s] Training on %ld vectors\n", elapsed() - t0, nb);
gpu_index->train(nb, xb);
@ -300,8 +287,9 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
}
printf("[%.3f s] Coping GPU index to CPU\n", elapsed() - t0);
delete cpu_index;
cpu_index = faiss::gpu::index_gpu_to_cpu(gpu_index);
delete gpu_index;
faiss::IndexIVF* cpu_ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
if (cpu_ivf_index != nullptr) {
@ -314,21 +302,7 @@ load_base_data(faiss::Index*& index, const std::string& ann_test_name, const std
delete[] xb;
}
switch (mode) {
case MODE_CPU:
case MODE_MIX:
index = cpu_index;
if (gpu_index) {
delete gpu_index;
}
break;
case MODE_GPU:
index = gpu_index;
if (cpu_index) {
delete cpu_index;
}
break;
}
index = cpu_index;
}
void
@ -379,10 +353,12 @@ load_ground_truth(faiss::Index::idx_t*& gt, size_t& k, const std::string& ann_te
}
void
test_with_nprobes(const std::string& ann_test_name, const std::string& index_key, faiss::Index* index,
test_with_nprobes(const std::string& ann_test_name, const std::string& index_key, faiss::Index* cpu_index,
faiss::gpu::StandardGpuResources& res, const QueryMode query_mode, const faiss::Index::distance_t* xq,
const faiss::Index::idx_t* gt, const std::vector<size_t> nprobes, const int32_t index_add_loops,
const int32_t search_loops) {
double t0 = elapsed();
const size_t NQ = 1000, NQ_START = 10, NQ_STEP = 10;
const size_t K = 1000, K_START = 100, K_STEP = 10;
const size_t GK = 100; // topk of ground truth
@ -390,6 +366,59 @@ test_with_nprobes(const std::string& ann_test_name, const std::string& index_key
std::unordered_map<size_t, std::string> mode_str_map = {
{MODE_CPU, "MODE_CPU"}, {MODE_MIX, "MODE_MIX"}, {MODE_GPU, "MODE_GPU"}};
faiss::Index *gpu_index, *index;
if (query_mode != MODE_CPU) {
faiss::gpu::GpuClonerOptions option;
option.allInGpu = true;
faiss::IndexComposition index_composition;
index_composition.index = cpu_index;
index_composition.quantizer = nullptr;
double copy_time;
switch (query_mode) {
case MODE_MIX: {
index_composition.mode = 1; // 0: all data, 1: copy quantizer, 2: copy data
// warm up the transmission
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
delete gpu_index;
copy_time = elapsed();
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
delete gpu_index;
copy_time = elapsed() - copy_time;
printf("[%.3f s] Copy quantizer completed, cost %f s\n", elapsed() - t0, copy_time);
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
auto is_gpu_flat_index = dynamic_cast<faiss::gpu::GpuIndexFlat*>(ivf_index->quantizer);
if (is_gpu_flat_index == nullptr) {
delete ivf_index->quantizer;
ivf_index->quantizer = index_composition.quantizer;
}
index = cpu_index;
break;
}
case MODE_GPU:
index_composition.mode = 0; // 0: all data, 1: copy quantizer, 2: copy data
// warm up the transmission
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
delete gpu_index;
copy_time = elapsed();
gpu_index = faiss::gpu::index_cpu_to_gpu(&res, GPU_DEVICE_IDX, &index_composition, &option);
copy_time = elapsed() - copy_time;
printf("[%.3f s] Copy data completed, cost %f s\n", elapsed() - t0, copy_time);
delete cpu_index;
index = gpu_index;
break;
}
} else {
index = cpu_index;
}
for (auto nprobe : nprobes) {
switch (query_mode) {
case MODE_CPU:
@ -445,6 +474,8 @@ test_with_nprobes(const std::string& ann_test_name, const std::string& index_key
delete[] I;
delete[] D;
}
delete index;
}
void
@ -488,7 +519,6 @@ test_ann_hdf5(const std::string& ann_test_name, const std::string& index_type, c
delete[] xq;
delete[] gt;
delete index;
}
/************************************************************************************
@ -520,8 +550,8 @@ TEST(FAISSTEST, BENCHMARK) {
#ifdef CUSTOMIZATION
test_ann_hdf5("sift-128-euclidean", "SQ8Hybrid", MODE_CPU, SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("sift-128-euclidean", "SQ8Hybrid", MODE_MIX, SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("sift-128-euclidean", "SQ8Hybrid", MODE_GPU, SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
// test_ann_hdf5("sift-128-euclidean", "SQ8Hybrid", MODE_MIX, SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
#endif
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@ -535,7 +565,7 @@ TEST(FAISSTEST, BENCHMARK) {
#ifdef CUSTOMIZATION
test_ann_hdf5("glove-200-angular", "SQ8Hybrid", MODE_CPU, GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("glove-200-angular", "SQ8Hybrid", MODE_MIX, GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("glove-200-angular", "SQ8Hybrid", MODE_GPU, GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
// test_ann_hdf5("glove-200-angular", "SQ8Hybrid", MODE_MIX, GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
#endif
}