milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #324 from Yukikaze-CZR/0.6.0 

Support new Index type IVFPQ #127
pull/376/head
Jin Hai 2019-11-15 19:54:49 +08:00 committed by GitHub
parent 6c996b5a18 4d1345b453
commit df28a763e1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
16 changed files with 244 additions and 98 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
CHANGELOG.md
View File

@ -14,6 +14,7 @@ Please mark all change in change log and use the ticket from JIRA.
- \#12 - Pure CPU version for Milvus
- \#77 - Support table partition
- \#226 - Experimental shards middleware for Milvus
- \#127 - Support new Index type IVFPQ
## Improvement
- \#275 - Rename C++ SDK IndexType
@ -21,7 +22,7 @@ Please mark all change in change log and use the ticket from JIRA.
- \#260 - C++ SDK README
- \#314 - add Find FAISS in CMake
- \#310 - Add Q&A for 'protocol https not supported or disable in libcurl' issue
- \#322 - Add option to enable / disable prometheus
- \#322 - Add option to enable / disable prometheus
- \#358 - Add more information in build.sh and install.md
## Task

4
core/src/db/engine/ExecutionEngine.h
View File

@ -26,6 +26,7 @@
namespace milvus {
namespace engine {
// TODO(linxj): replace with VecIndex::IndexType
enum class EngineType {
INVALID = 0,
FAISS_IDMAP = 1,
@ -33,7 +34,8 @@ enum class EngineType {
FAISS_IVFSQ8,
NSG_MIX,
FAISS_IVFSQ8H,
MAX_VALUE = FAISS_IVFSQ8H,
FAISS_PQ,
MAX_VALUE = FAISS_PQ,
};
enum class MetricType {

8
core/src/db/engine/ExecutionEngineImpl.cpp
View File

@ -116,6 +116,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID);
break;
}
case EngineType::FAISS_PQ: {
#ifdef MILVUS_CPU_VERSION
index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_CPU);
#else
index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX);
#endif
break;
}
default: {
ENGINE_LOG_ERROR << "Unsupported index type";
return nullptr;

33
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
View File

@ -39,17 +39,19 @@ GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) {
GETTENSOR(dataset)
// TODO(linxj): set device here.
// TODO(linxj): set gpu resource here.
faiss::gpu::StandardGpuResources res;
faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, build_cfg->nlist, build_cfg->m, build_cfg->nbits,
GetMetricType(build_cfg->metric_type)); // IP not support
device_index.train(rows, (float*)p_data);
std::shared_ptr<faiss::Index> host_index = nullptr;
host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index));
return std::make_shared<IVFIndexModel>(host_index);
auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_);
if (temp_resource != nullptr) {
ResScope rs(temp_resource, gpu_id_, true);
auto device_index = new faiss::gpu::GpuIndexIVFPQ(temp_resource->faiss_res.get(), dim, build_cfg->nlist,
build_cfg->m, build_cfg->nbits,
GetMetricType(build_cfg->metric_type)); // IP not support
device_index->train(rows, (float*)p_data);
std::shared_ptr<faiss::Index> host_index = nullptr;
host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index));
return std::make_shared<IVFIndexModel>(host_index);
} else {
KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource");
}
}
std::shared_ptr<faiss::IVFSearchParameters>
@ -66,7 +68,14 @@ GPUIVFPQ::GenParams(const Config& config) {
VectorIndexPtr
GPUIVFPQ::CopyGpuToCpu(const Config& config) {
KNOWHERE_THROW_MSG("not support yet");
std::lock_guard<std::mutex> lk(mutex_);
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IVFPQ>(new_index);
}
} // namespace knowhere

5
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h
View File

@ -18,6 +18,7 @@
#pragma once
#include <memory>
#include <utility>
#include "IndexGPUIVF.h"
@ -28,6 +29,10 @@ class GPUIVFPQ : public GPUIVF {
explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) {
}
GPUIVFPQ(std::shared_ptr<faiss::Index> index, const int64_t& device_id, ResPtr& resource)
: GPUIVF(std::move(index), device_id, resource) {
}
IndexModelPtr
Train(const DatasetPtr& dataset, const Config& config) override;

26
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
View File

@ -17,11 +17,19 @@
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQ.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuCloner.h>
#endif
#include <memory>
#include <utility>
#include "knowhere/adapter/VectorAdapter.h"
#include "knowhere/common/Exception.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#endif
#include "knowhere/index/vector_index/IndexIVFPQ.h"
namespace knowhere {
@ -60,4 +68,22 @@ IVFPQ::Clone_impl(const std::shared_ptr<faiss::Index>& index) {
return std::make_shared<IVFPQ>(index);
}
VectorIndexPtr
IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
#ifdef MILVUS_GPU_VERSION
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
ResScope rs(res, device_id, false);
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
std::shared_ptr<faiss::Index> device_index;
device_index.reset(gpu_index);
return std::make_shared<GPUIVFPQ>(device_index, device_id, res);
} else {
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
}
#else
KNOWHERE_THROW_MSG("Calling IVFPQ::CopyCpuToGpu when we are using CPU version");
#endif
}
} // namespace knowhere

3
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h
View File

@ -34,6 +34,9 @@ class IVFPQ : public IVF {
IndexModelPtr
Train(const DatasetPtr& dataset, const Config& config) override;
VectorIndexPtr
CopyCpuToGpu(const int64_t& device_id, const Config& config) override;
protected:
std::shared_ptr<faiss::IVFSearchParameters>
GenParams(const Config& config) override;

2
core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
View File

@ -49,7 +49,7 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config
if (auto cpu_index = std::dynamic_pointer_cast<IVFSQ>(index)) {
return cpu_index->CopyCpuToGpu(device_id, config);
} else if (auto cpu_index = std::dynamic_pointer_cast<IVFPQ>(index)) {
KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu");
return cpu_index->CopyCpuToGpu(device_id, config);
} else if (auto cpu_index = std::dynamic_pointer_cast<IVF>(index)) {
return cpu_index->CopyCpuToGpu(device_id, config);
} else if (auto cpu_index = std::dynamic_pointer_cast<IDMAP>(index)) {

4
core/src/index/unittest/test_ivf.cpp
View File

@ -213,7 +213,7 @@ TEST_P(IVFTest, clone_test) {
{
// copy from gpu to cpu
std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "GPUIVFPQ", "IVFSQHybrid"};
auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
if (finder != support_idx_vec.cend()) {
EXPECT_NO_THROW({
@ -238,7 +238,7 @@ TEST_P(IVFTest, clone_test) {
{
// copy to gpu
std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ"};
std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ", "IVFPQ", "GPUIVFPQ"};
auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
if (finder != support_idx_vec.cend()) {
EXPECT_NO_THROW({

61
core/src/wrapper/ConfAdapter.cpp
View File

@ -16,6 +16,7 @@
// under the License.
#include "wrapper/ConfAdapter.h"
#include "WrapperException.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "utils/Log.h"
@ -76,7 +77,7 @@ IVFConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) {
if (size <= TYPICAL_COUNT / 16384 + 1) {
// handle less row count, avoid nlist set to 0
return 1;
} else if (int(size / TYPICAL_COUNT) * nlist == 0) {
} else if (int(size / TYPICAL_COUNT) * nlist <= 0) {
// calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT
return int(size / TYPICAL_COUNT * 16384);
}
@ -87,7 +88,11 @@ knowhere::Config
IVFConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
auto conf = std::make_shared<knowhere::IVFCfg>();
conf->k = metaconf.k;
conf->nprobe = metaconf.nprobe;
if (metaconf.nprobe <= 0)
conf->nprobe = 16; // hardcode here
else
conf->nprobe = metaconf.nprobe;
switch (type) {
case IndexType::FAISS_IVFFLAT_GPU:
@ -123,11 +128,47 @@ IVFPQConfAdapter::Match(const TempMetaConf& metaconf) {
conf->metric_type = metaconf.metric_type;
conf->gpu_id = conf->gpu_id;
conf->nbits = 8;
conf->m = 8;
if (!(conf->d % 4))
conf->m = conf->d / 4; // compression radio = 16
else if (!(conf->d % 2))
conf->m = conf->d / 2; // compression radio = 8
else if (!(conf->d % 3))
conf->m = conf->d / 3; // compression radio = 12
else
conf->m = conf->d; // same as SQ8, compression radio = 4
MatchBase(conf);
return conf;
}
knowhere::Config
IVFPQConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) {
auto conf = std::make_shared<knowhere::IVFPQCfg>();
conf->k = metaconf.k;
if (metaconf.nprobe <= 0) {
WRAPPER_LOG_ERROR << "The nprobe of PQ is wrong!";
throw WrapperException("The nprobe of PQ is wrong!");
} else {
conf->nprobe = metaconf.nprobe;
}
return conf;
}
int64_t
IVFPQConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) {
if (size <= TYPICAL_COUNT / 16384 + 1) {
// handle less row count, avoid nlist set to 0
return 1;
} else if (int(size / TYPICAL_COUNT) * nlist <= 0) {
// calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT
return int(size / TYPICAL_COUNT * 16384);
}
return nlist;
}
knowhere::Config
NSGConfAdapter::Match(const TempMetaConf& metaconf) {
auto conf = std::make_shared<knowhere::NSGCfg>();
@ -136,13 +177,14 @@ NSGConfAdapter::Match(const TempMetaConf& metaconf) {
conf->metric_type = metaconf.metric_type;
conf->gpu_id = conf->gpu_id;
double factor = metaconf.size / TYPICAL_COUNT;
auto scale_factor = round(metaconf.dim / 128.0);
scale_factor = scale_factor >= 4 ? 4 : scale_factor;
conf->nprobe = 6 + 10 * scale_factor;
conf->knng = 100 + 100 * scale_factor;
conf->search_length = 40 + 5 * scale_factor;
conf->out_degree = 50 + 5 * scale_factor;
conf->candidate_pool_size = 200 + 100 * scale_factor;
conf->nprobe = conf->nlist > 10000 ? conf->nlist * 0.02 : conf->nlist * 0.1;
conf->knng = (100 + 100 * scale_factor) * factor;
conf->search_length = (40 + 5 * scale_factor) * factor;
conf->out_degree = (50 + 5 * scale_factor) * factor;
conf->candidate_pool_size = (200 + 100 * scale_factor) * factor;
MatchBase(conf);
// WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist
@ -156,6 +198,9 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type)
auto conf = std::make_shared<knowhere::NSGCfg>();
conf->k = metaconf.k;
conf->search_length = metaconf.search_length;
if (metaconf.search_length == TEMPMETA_DEFAULT_VALUE) {
conf->search_length = 30; // TODO(linxj): hardcode here.
}
return conf;
}

7
core/src/wrapper/ConfAdapter.h
View File

@ -79,6 +79,13 @@ class IVFPQConfAdapter : public IVFConfAdapter {
public:
knowhere::Config
Match(const TempMetaConf& metaconf) override;
knowhere::Config
MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override;
protected:
static int64_t
MatchNlist(const int64_t& size, const int64_t& nlist);
};
class NSGConfAdapter : public IVFConfAdapter {

1
core/src/wrapper/ConfAdapterMgr.cpp
View File

@ -53,6 +53,7 @@ AdapterMgr::RegisterAdapter() {
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_CPU, ivfpq_cpu);
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_GPU, ivfpq_gpu);
REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix);
REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix);
}

13
core/src/wrapper/VecIndex.cpp
View File

@ -145,6 +145,10 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
break;
}
case IndexType::FAISS_IVFPQ_MIX: {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFPQ_MIX);
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
@ -276,6 +280,10 @@ ConvertToCpuIndexType(const IndexType& type) {
case IndexType::FAISS_IVFSQ8_MIX: {
return IndexType::FAISS_IVFSQ8_CPU;
}
case IndexType::FAISS_IVFPQ_GPU:
case IndexType::FAISS_IVFPQ_MIX: {
return IndexType::FAISS_IVFPQ_CPU;
}
default: { return type; }
}
}
@ -291,9 +299,12 @@ ConvertToGpuIndexType(const IndexType& type) {
case IndexType::FAISS_IVFSQ8_CPU: {
return IndexType::FAISS_IVFSQ8_GPU;
}
case IndexType::FAISS_IVFPQ_MIX:
case IndexType::FAISS_IVFPQ_CPU: {
return IndexType::FAISS_IVFPQ_GPU;
}
default: { return type; }
}
}
} // namespace engine
} // namespace milvus

2
core/src/wrapper/VecIndex.h
View File

@ -33,6 +33,7 @@ namespace engine {
using Config = knowhere::Config;
// TODO(linxj): replace with string, Do refactor serialization
enum class IndexType {
INVALID = 0,
FAISS_IDMAP = 1,
@ -47,6 +48,7 @@ enum class IndexType {
FAISS_IVFSQ8_GPU,
FAISS_IVFSQ8_HYBRID, // only support build on gpu.
NSG_MIX,
FAISS_IVFPQ_MIX,
};
class VecIndex;

63
core/unittest/wrapper/test_wrapper.cpp
View File

@ -29,33 +29,40 @@
INITIALIZE_EASYLOGGINGPP
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
class KnowhereWrapperTest
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
protected:
void
SetUp() override {
void SetUp() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
#endif
std::string generator_type;
std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam();
GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
milvus::engine::TempMetaConf tempconf;
tempconf.metric_type = knowhere::METRICTYPE::L2;
tempconf.gpu_id = DEVICEID;
tempconf.size = nb;
tempconf.dim = dim;
tempconf.k = k;
tempconf.nprobe = 16;
index_ = GetVecIndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(index_type);
conf->k = k;
conf->d = dim;
conf->gpu_id = DEVICEID;
conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf);
searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf);
// conf->k = k;
// conf->d = dim;
// conf->gpu_id = DEVICEID;
}
void
TearDown() override {
void TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
@ -65,24 +72,27 @@ class KnowhereWrapperTest
milvus::engine::IndexType index_type;
milvus::engine::VecIndexPtr index_ = nullptr;
knowhere::Config conf;
knowhere::Config searchconf;
};
INSTANTIATE_TEST_CASE_P(
WrapperParam, KnowhereWrapperTest,
Values(
//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
Values(
//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
#ifdef MILVUS_GPU_VERSION
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 100000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 1000, 10, 10),
// std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB,
// 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", 64, 1000, 10, 10),
// std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10),
#endif
// std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 100000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 1000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 1000, 10, 10),
std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10)));
TEST_P(KnowhereWrapperTest, BASE_TEST) {
@ -93,12 +103,11 @@ TEST_P(KnowhereWrapperTest, BASE_TEST) {
std::vector<float> res_dis(elems);
index_->BuildAll(nb, xb.data(), ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
AssertResult(res_ids, res_dis);
}
#ifdef MILVUS_GPU_VERSION
TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
EXPECT_EQ(index_->GetType(), index_type);
@ -107,13 +116,13 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
std::vector<float> res_dis(elems);
index_->BuildAll(nb, xb.data(), ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
AssertResult(res_ids, res_dis);
{
auto dev_idx = index_->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
}
AssertResult(res_ids, res_dis);
}
@ -125,7 +134,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
auto dev_idx = new_index->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
}
AssertResult(res_ids, res_dis);
}
@ -139,7 +148,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
std::vector<int64_t> res_ids(elems);
std::vector<float> res_dis(elems);
index_->BuildAll(nb, xb.data(), ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
AssertResult(res_ids, res_dis);
{
@ -152,7 +161,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
std::vector<int64_t> res_ids(elems);
std::vector<float> res_dis(elems);
new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
AssertResult(res_ids, res_dis);
}
@ -166,7 +175,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
std::vector<int64_t> res_ids(elems);
std::vector<float> res_dis(elems);
new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf);
AssertResult(res_ids, res_dis);
}
}

107
core/unittest/wrapper/utils.h
View File

@ -15,18 +15,22 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <gtest/gtest.h>
#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <memory>
#include <vector>
#include <cstdlib>
#include <cstdio>
#include <fstream>
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include "wrapper/ConfAdapterMgr.h"
#include "wrapper/ConfAdapter.h"
class DataGenBase;
@ -40,29 +44,31 @@ constexpr int64_t PINMEM = 1024 * 1024 * 200;
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
constexpr int64_t RESNUM = 2;
static const char* CONFIG_PATH = "/tmp/milvus_test";
static const char* CONFIG_FILE = "/server_config.yaml";
static const char *CONFIG_PATH = "/tmp/milvus_test";
static const char *CONFIG_FILE = "/server_config.yaml";
class KnowhereTest : public ::testing::Test {
protected:
void
SetUp() override;
void
TearDown() override;
void SetUp() override;
void TearDown() override;
};
class DataGenBase {
public:
virtual void
GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids, const int& k,
int64_t* gt_ids, float* gt_dis);
virtual void GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids,
const int& k, int64_t* gt_ids, float* gt_dis);
virtual void
GenData(const int& dim, const int& nb, const int& nq, std::vector<float>& xb, std::vector<float>& xq,
std::vector<int64_t>& ids, const int& k, std::vector<int64_t>& gt_ids, std::vector<float>& gt_dis);
virtual void GenData(const int& dim,
const int& nb,
const int& nq,
std::vector<float>& xb,
std::vector<float>& xq,
std::vector<int64_t>& ids,
const int& k,
std::vector<int64_t>& gt_ids,
std::vector<float>& gt_dis);
void
AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
int dim = DIM;
int nb = NB;
@ -79,14 +85,22 @@ class DataGenBase {
class ParamGenerator {
public:
static ParamGenerator&
GetInstance() {
static ParamGenerator& GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config
Gen(const milvus::engine::IndexType& type) {
knowhere::Config GenSearchConf(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) {
auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type);
return adapter->MatchSearch(conf, type);
}
knowhere::Config GenBuild(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) {
auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type);
return adapter->Match(conf);
}
knowhere::Config Gen(const milvus::engine::IndexType& type) {
switch (type) {
case milvus::engine::IndexType::FAISS_IDMAP: {
auto tempconf = std::make_shared<knowhere::Cfg>();
@ -113,34 +127,37 @@ class ParamGenerator {
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
// case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
// case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
// auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
// tempconf->nlist = 100;
// tempconf->nprobe = 16;
// tempconf->nbits = 8;
// tempconf->m = 8;
// tempconf->metric_type = knowhere::METRICTYPE::L2;
// return tempconf;
// }
// case milvus::engine::IndexType::NSG_MIX: {
// auto tempconf = std::make_shared<knowhere::NSGCfg>();
// tempconf->nlist = 100;
// tempconf->nprobe = 16;
// tempconf->search_length = 8;
// tempconf->knng = 200;
// tempconf->search_length = 40; // TODO(linxj): be 20 when search
// tempconf->out_degree = 60;
// tempconf->candidate_pool_size = 200;
// tempconf->metric_type = knowhere::METRICTYPE::L2;
// return tempconf;
// }
case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
case milvus::engine::IndexType::FAISS_IVFPQ_GPU:
case milvus::engine::IndexType::FAISS_IVFPQ_MIX: {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->m = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::NSG_MIX: {
auto tempconf = std::make_shared<knowhere::NSGCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->search_length = 8;
tempconf->knng = 200;
tempconf->search_length = 40; // TODO(linxj): be 20 when search
tempconf->out_degree = 60;
tempconf->candidate_pool_size = 200;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
}
};
// class SanityCheck : public DataGenBase {
//class SanityCheck : public DataGenBase {
// public:
// void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, long *ids,
// const int &k, long *gt_ids, float *gt_dis) override;
//};