Skip to content

milvus
milvus

qq_35024513 / milvus
与 Fork 源项目一致

Fork自 milvus / milvus

通知 14
Star 0
Fork 0
milvus
milvus

前往新版Gitcode，体验更适合开发者的 AI 搜索 >>

提交 7f6092b6 编写于 作者: Y yudong.cai
浏览文件
操作

#89 add SQ8Hybrid MIX test and pure-GPU test 


Former-commit-id: 25b5f419d3deedd4ec39c014ea47822a83a0af38
上级 ca0f5e46
隐藏空白更改
内联 并排
Showing with 135 addition and 77 deletion
core/src/index/unittest/faiss_benchmark/faiss_benchmark_test.cpp
浏览文件 @ 7f6092b6
......@@ -28,6 +28,7 @@
#include <faiss/IndexIVF.h>
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/GpuIndexIVFSQHybrid.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/index_io.h>
#include <faiss/utils.h>
......@@ -183,9 +184,31 @@ parse_ann_test_name(const std::string& ann_test_name, size_t& dim, faiss::Metric
return true;
}
int32_t
GetResultHitCount(const faiss::Index::idx_t* ground_index, const faiss::Index::idx_t* index, size_t ground_k, size_t k,
size_t nq, int32_t index_add_loops) {
assert(ground_k <= k);
int hit = 0;
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++) {
int r_c = index[i * k + j_c];
int j_g = 0;
for (; j_g < ground_k / index_add_loops; j_g++) {
if (ground_index[i * ground_k + j_g] == r_c) {
hit++;
continue;
}
}
}
}
return hit;
}
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) {
const std::vector<size_t>& nprobes, int32_t search_loops) {
double t0 = elapsed();
const std::string ann_file_name = ann_test_name + ".hdf5";
......@@ -265,8 +288,6 @@ test_ann_hdf5(const std::string& ann_test_name, const std::string& index_key, in
for (auto nprobe : nprobes) {
faiss::ParameterSpace params;
printf("[%.3f s] Setting parameter configuration 'nprobe=%lu' on index\n", elapsed() - t0, nprobe);
std::string nprobe_str = "nprobe=" + std::to_string(nprobe);
params.set_index_parameters(index, nprobe_str.c_str());
......@@ -277,39 +298,28 @@ test_ann_hdf5(const std::string& ann_test_name, const std::string& index_key, in
float* D = new float[NQ * K];
printf("\n%s | %s | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
printf("============================================================================================\n");
printf("======================================================================================\n");
for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) { // nq = {10, 100, 1000}
for (size_t t_k = 100; t_k <= K; t_k *= 10) { // k = {100, 1000}
faiss::indexIVF_stats.quantization_time = 0.0;
faiss::indexIVF_stats.search_time = 0.0;
double t_start = elapsed(), t_end;
index->search(t_nq, xq, t_k, D, I);
for (int i = 0; i < search_loops; i++) {
index->search(t_nq, xq, t_k, D, I);
}
t_end = elapsed();
// k = 100 for ground truth
int hit = 0;
for (int i = 0; i < t_nq; i++) {
// count the num of results exist in ground truth result set
// consider: each result replicates DATA_LOOPS times
for (int j_c = 0; j_c < k; j_c++) {
int r_c = I[i * t_k + j_c];
for (int j_g = 0; j_g < k / index_add_loops; j_g++) {
if (gt[i * k + j_g] == r_c) {
hit++;
continue;
}
}
}
}
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)));
int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);
printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq, t_k,
(t_end - t_start) / search_loops, faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
faiss::indexIVF_stats.search_time / 1000 / search_loops,
(hit / float(t_nq * k / index_add_loops)));
}
}
printf("============================================================================================\n");
printf("======================================================================================\n");
#else
printf("[%.3f s] Perform a search on %ld queries\n", elapsed() - t0, nq);
......@@ -353,7 +363,8 @@ test_ann_hdf5(const std::string& ann_test_name, const std::string& index_key, in
#ifdef CUSTOMIZATION
void
test_ivfsq8h_gpu(const std::string& ann_test_name, int32_t index_add_loops, const std::vector<size_t>& nprobes) {
test_ivfsq8h(const std::string& ann_test_name, int32_t index_add_loops, const std::vector<size_t>& nprobes,
bool pure_gpu_mode, int32_t search_loops) {
double t0 = elapsed();
const std::string ann_file_name = ann_test_name + ".hdf5";
......@@ -423,9 +434,18 @@ test_ivfsq8h_gpu(const std::string& ann_test_name, int32_t index_add_loops, cons
index_composition.quantizer = nullptr;
index_composition.mode = 1;
double copy_time = elapsed();
auto index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
delete index;
if (pure_gpu_mode) {
index_composition.mode = 2; // 0: all data, 1: copy quantizer, 2: copy data
index = faiss::gpu::index_cpu_to_gpu(&res, 0, &index_composition, &option);
}
copy_time = elapsed() - copy_time;
printf("[%.3f s] Copy quantizer completed, cost %f s\n", elapsed() - t0, copy_time);
size_t nq;
float* xq;
{
......@@ -446,67 +466,98 @@ test_ivfsq8h_gpu(const std::string& ann_test_name, int32_t index_add_loops, cons
assert(nq2 == nq || !"incorrect nb of ground truth entries");
gt = new faiss::Index::idx_t[k * nq];
for (unsigned long i = 0; i < k * nq; ++i) {
for (uint64_t i = 0; i < k * nq; ++i) {
gt[i] = gt_int[i];
}
delete[] gt_int;
}
for (auto nprobe : nprobes) {
printf("[%.3f s] Setting parameter configuration 'nprobe=%lu' on index\n", elapsed() - t0, nprobe);
const size_t NQ = 1000, K = 1000;
if (!pure_gpu_mode) {
for (auto nprobe : nprobes) {
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
ivf_index->nprobe = nprobe;
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(cpu_index);
ivf_index->nprobe = nprobe;
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;
}
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;
}
int64_t* I = new faiss::Index::idx_t[NQ * K];
float* D = new float[NQ * K];
const size_t NQ = 1000, K = 1000;
long* I = new faiss::Index::idx_t[NQ * K];
float* D = new float[NQ * K];
printf("\n%s | %s-MIX | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
printf("======================================================================================\n");
for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) { // nq = {10, 100, 1000}
for (size_t t_k = 100; t_k <= K; t_k *= 10) { // k = {100, 1000}
faiss::indexIVF_stats.quantization_time = 0.0;
faiss::indexIVF_stats.search_time = 0.0;
printf("\n%s | %s-gpu | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
printf("============================================================================================\n");
for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) { // nq = {10, 100, 1000}
for (size_t t_k = 100; t_k <= K; t_k *= 10) { // k = {100, 1000}
faiss::indexIVF_stats.quantization_time = 0.0;
faiss::indexIVF_stats.search_time = 0.0;
double t_start = elapsed(), t_end;
for (int32_t i = 0; i < search_loops; i++) {
cpu_index->search(t_nq, xq, t_k, D, I);
}
t_end = elapsed();
double t_start = elapsed(), t_end;
// k = 100 for ground truth
int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);
cpu_index->search(t_nq, xq, t_k, D, I);
printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq,
t_k, (t_end - t_start) / search_loops,
faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
faiss::indexIVF_stats.search_time / 1000 / search_loops,
(hit / float(t_nq * k / index_add_loops)));
}
}
printf("======================================================================================\n");
t_end = elapsed();
printf("[%.3f s] Search test done\n\n", elapsed() - t0);
// k = 100 for ground truth
int hit = 0;
for (unsigned long i = 0; i < t_nq; i++) {
// count the num of results exist in ground truth result set
// consider: each result replicates DATA_LOOPS times
for (unsigned long j_c = 0; j_c < k; j_c++) {
int r_c = I[i * t_k + j_c];
for (unsigned long j_g = 0; j_g < k / index_add_loops; j_g++) {
if (gt[i * k + j_g] == r_c) {
hit++;
continue;
}
}
delete[] I;
delete[] D;
}
} else {
std::shared_ptr<faiss::Index> gpu_index_ivf_ptr = std::shared_ptr<faiss::Index>(index);
for (auto nprobe : nprobes) {
faiss::gpu::GpuIndexIVFSQHybrid* gpu_index_ivf_hybrid =
dynamic_cast<faiss::gpu::GpuIndexIVFSQHybrid*>(gpu_index_ivf_ptr.get());
gpu_index_ivf_hybrid->setNumProbes(nprobe);
int64_t* I = new faiss::Index::idx_t[NQ * K];
float* D = new float[NQ * K];
printf("\n%s | %s-GPU | nprobe=%lu\n", ann_test_name.c_str(), index_key.c_str(), nprobe);
printf("======================================================================================\n");
for (size_t t_nq = 10; t_nq <= NQ; t_nq *= 10) { // nq = {10, 100, 1000}
for (size_t t_k = 100; t_k <= K; t_k *= 10) { // k = {100, 1000}
faiss::indexIVF_stats.quantization_time = 0.0;
faiss::indexIVF_stats.search_time = 0.0;
double t_start = elapsed(), t_end;
for (int32_t i = 0; i < search_loops; i++) {
gpu_index_ivf_ptr->search(nq, xq, k, D, I);
}
t_end = elapsed();
// k = 100 for ground truth
int32_t hit = GetResultHitCount(gt, I, k, t_k, t_nq, index_add_loops);
printf("nq = %4ld, k = %4ld, elapse = %.4fs (quant = %.4fs, search = %.4fs), R@ = %.4f\n", t_nq,
t_k, (t_end - t_start) / search_loops,
faiss::indexIVF_stats.quantization_time / 1000 / search_loops,
faiss::indexIVF_stats.search_time / 1000 / search_loops,
(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");
printf("======================================================================================\n");
printf("[%.3f s] Search test done\n\n", elapsed() - t0);
printf("[%.3f s] Search test done\n\n", elapsed() - t0);
delete[] I;
delete[] D;
delete[] I;
delete[] D;
}
}
delete[] xq;
......@@ -530,17 +581,24 @@ test_ivfsq8h_gpu(const std::string& ann_test_name, int32_t index_add_loops, cons
*************************************************************************************/
TEST(FAISSTEST, BENCHMARK) {
test_ann_hdf5("sift-128-euclidean", "IVF4096,Flat", 2, {8, 128});
test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8", 2, {8, 128});
std::vector<size_t> param_nprobes = {8, 128};
const int32_t SEARCH_LOOPS = 5;
const int32_t SIFT_INSERT_LOOPS = 2; // insert twice to get ~1G data set
const int32_t GLOVE_INSERT_LOOPS = 1;
test_ann_hdf5("sift-128-euclidean", "IVF4096,Flat", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
#ifdef CUSTOMIZATION
test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8Hybrid", 2, {8, 128});
test_ivfsq8h_gpu("sift-128-euclidean", 2, {8, 128});
test_ann_hdf5("sift-128-euclidean", "IVF16384,SQ8Hybrid", SIFT_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ivfsq8h("sift-128-euclidean", SIFT_INSERT_LOOPS, param_nprobes, false, SEARCH_LOOPS);
test_ivfsq8h("sift-128-euclidean", SIFT_INSERT_LOOPS, param_nprobes, true, SEARCH_LOOPS);
#endif
test_ann_hdf5("glove-200-angular", "IVF4096,Flat", 1, {8, 128});
test_ann_hdf5("glove-200-angular", "IVF16384,SQ8", 1, {8, 128});
test_ann_hdf5("glove-200-angular", "IVF4096,Flat", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ann_hdf5("glove-200-angular", "IVF16384,SQ8", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
#ifdef CUSTOMIZATION
test_ann_hdf5("glove-200-angular", "IVF16384,SQ8Hybrid", 1, {8, 128});
test_ivfsq8h_gpu("glove-200-angular", 1, {8, 128});
test_ann_hdf5("glove-200-angular", "IVF16384,SQ8Hybrid", GLOVE_INSERT_LOOPS, param_nprobes, SEARCH_LOOPS);
test_ivfsq8h("glove-200-angular", GLOVE_INSERT_LOOPS, param_nprobes, false, SEARCH_LOOPS);
test_ivfsq8h("glove-200-angular", GLOVE_INSERT_LOOPS, param_nprobes, true, SEARCH_LOOPS);
#endif
}
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或