Skip to content

milvus
milvus

BaiXuePrincess / milvus
与 Fork 源项目一致

从无法访问的项目Fork

通知 7
Star 4
Fork 0
milvus
milvus

体验新版 GitCode，发现更多精彩内容 >>

提交 b46168cf 编写于 作者: J jinhai
浏览文件
操作

Merge branch 'dev-sq8' into 'branch-0.3.1-xiaojun' 

Dev sq8

See merge request megasearch/milvus!269

Former-commit-id: 7f2c9ce4b52d0436e5cea77f57b379b5edb8ff4a
上级 168ca0a5 7b3886fd
显示空白变更内容
内联 并排
Showing with 62 addition and 32 deletion
cpp/src/db/ExecutionEngineImpl.cpp
浏览文件 @ b46168cf
......@@ -30,7 +30,10 @@ ExecutionEngineImpl::ExecutionEngineImpl(uint16_t dimension,
index_ = CreatetVecIndex(EngineType::FAISS_IDMAP);
if (!index_) throw Exception("Create Empty VecIndex");
auto ec = std::static_pointer_cast<BFIndex>(index_)->Build(dimension);
Config build_cfg;
build_cfg["dim"] = dimension;
AutoGenParams(index_->GetType(), 0, build_cfg);
auto ec = std::static_pointer_cast<BFIndex>(index_)->Build(build_cfg);
if (ec != server::KNOWHERE_SUCCESS) { throw Exception("Build index error"); }
}
......@@ -69,7 +72,7 @@ VecIndexPtr ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
}
Status ExecutionEngineImpl::AddWithIds(long n, const float *xdata, const long *xids) {
auto ec = index_->Add(n, xdata, xids, Config::object{{"dim", dim}});
auto ec = index_->Add(n, xdata, xids);
if (ec != server::KNOWHERE_SUCCESS) {
return Status::Error("Add error");
}
......@@ -171,10 +174,15 @@ ExecutionEngineImpl::BuildIndex(const std::string &location) {
throw Exception("Create Empty VecIndex");
}
Config build_cfg;
build_cfg["dim"] = Dimension();
build_cfg["gpu_id"] = gpu_num;
AutoGenParams(to_index->GetType(), Count(), build_cfg);
auto ec = to_index->BuildAll(Count(),
from_index->GetRawVectors(),
from_index->GetRawIds(),
Config::object{{"dim", Dimension()}, {"gpu_id", gpu_num}});
build_cfg);
if (ec != server::KNOWHERE_SUCCESS) { throw Exception("Build index error"); }
return std::make_shared<ExecutionEngineImpl>(to_index, location, build_type);
......
cpp/src/wrapper/knowhere/vec_impl.cpp
浏览文件 @ b46168cf
......@@ -32,9 +32,7 @@ server::KnowhereError VecIndexImpl::BuildAll(const long &nb,
auto preprocessor = index_->BuildPreprocessor(dataset, cfg);
index_->set_preprocessor(preprocessor);
auto nlist = int(nb / 1000000.0 * 16384);
auto cfg_t = Config::object{{"nlist", nlist}, {"dim", dim}};
auto model = index_->Train(dataset, cfg_t);
auto model = index_->Train(dataset, cfg);
index_->set_index_model(model);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
......@@ -52,8 +50,7 @@ server::KnowhereError VecIndexImpl::BuildAll(const long &nb,
server::KnowhereError VecIndexImpl::Add(const long &nb, const float *xb, const long *ids, const Config &cfg) {
try {
auto d = cfg.get_with_default("dim", dim);
auto dataset = GenDatasetWithIds(nb, d, xb, ids);
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
......@@ -72,8 +69,7 @@ server::KnowhereError VecIndexImpl::Add(const long &nb, const float *xb, const l
server::KnowhereError VecIndexImpl::Search(const long &nq, const float *xq, float *dist, long *ids, const Config &cfg) {
try {
auto k = cfg["k"].as<int>();
auto d = cfg.get_with_default("dim", dim);
auto dataset = GenDataset(nq, d, xq);
auto dataset = GenDataset(nq, dim, xq);
Config search_cfg;
auto res = index_->Search(dataset, cfg);
......@@ -148,10 +144,10 @@ int64_t *BFIndex::GetRawIds() {
return std::static_pointer_cast<IDMAP>(index_)->GetRawIds();
}
server::KnowhereError BFIndex::Build(const int64_t &d) {
server::KnowhereError BFIndex::Build(const Config &cfg) {
try {
dim = d;
std::static_pointer_cast<IDMAP>(index_)->Train(dim);
dim = cfg["dim"].as<int>();
std::static_pointer_cast<IDMAP>(index_)->Train(cfg);
} catch (KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return server::KNOWHERE_UNEXPECTED_ERROR;
......@@ -175,7 +171,7 @@ server::KnowhereError BFIndex::BuildAll(const long &nb,
dim = cfg["dim"].as<int>();
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
std::static_pointer_cast<IDMAP>(index_)->Train(dim);
std::static_pointer_cast<IDMAP>(index_)->Train(cfg);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
......@@ -203,9 +199,7 @@ server::KnowhereError IVFMixIndex::BuildAll(const long &nb,
auto preprocessor = index_->BuildPreprocessor(dataset, cfg);
index_->set_preprocessor(preprocessor);
auto nlist = int(nb / 1000000.0 * 16384);
auto cfg_t = Config::object{{"nlist", nlist}, {"dim", dim}};
auto model = index_->Train(dataset, cfg_t);
auto model = index_->Train(dataset, cfg);
index_->set_index_model(model);
index_->Add(dataset, cfg);
......
cpp/src/wrapper/knowhere/vec_impl.h
浏览文件 @ b46168cf
......@@ -41,8 +41,9 @@ class VecIndexImpl : public VecIndex {
class IVFMixIndex : public VecIndexImpl {
public:
explicit IVFMixIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index) : VecIndexImpl(std::move(index),
IndexType::FAISS_IVFFLAT_MIX) {};
explicit IVFMixIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index, const IndexType &type)
: VecIndexImpl(std::move(index), type) {};
server::KnowhereError BuildAll(const long &nb,
const float *xb,
const long *ids,
......@@ -56,7 +57,7 @@ class BFIndex : public VecIndexImpl {
public:
explicit BFIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index) : VecIndexImpl(std::move(index),
IndexType::FAISS_IDMAP) {};
server::KnowhereError Build(const int64_t &d);
server::KnowhereError Build(const Config& cfg);
float *GetRawVectors();
server::KnowhereError BuildAll(const long &nb,
const float *xb,
......
cpp/src/wrapper/knowhere/vec_index.cpp
浏览文件 @ b46168cf
......@@ -85,7 +85,7 @@ VecIndexPtr GetVecIndexFactory(const IndexType &type) {
}
case IndexType::FAISS_IVFFLAT_MIX: {
index = std::make_shared<zilliz::knowhere::GPUIVF>(0);
return std::make_shared<IVFMixIndex>(index);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFFLAT_MIX);
}
case IndexType::FAISS_IVFPQ_CPU: {
index = std::make_shared<zilliz::knowhere::IVFPQ>();
......@@ -98,6 +98,10 @@ VecIndexPtr GetVecIndexFactory(const IndexType &type) {
case IndexType::SPTAG_KDT_RNT_CPU: {
index = std::make_shared<zilliz::knowhere::CPUKDTRNG>();
break;
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<zilliz::knowhere::GPUIVFSQ>(0);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
}
//case IndexType::NSG: { // TODO(linxj): bug.
// index = std::make_shared<zilliz::knowhere::NSG>();
......@@ -176,13 +180,28 @@ server::KnowhereError write_index(VecIndexPtr index, const std::string &location
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return server::KNOWHERE_UNEXPECTED_ERROR;
} catch (std::exception& e) {
} catch (std::exception &e) {
WRAPPER_LOG_ERROR << e.what();
return server::KNOWHERE_ERROR;
}
return server::KNOWHERE_SUCCESS;
}
// TODO(linxj): redo here.
void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Config &cfg) {
if (!cfg.contains("nlist")) { cfg["nlist"] = int(size / 1000000.0 * 16384); }
if (!cfg.contains("gpu_id")) { cfg["gpu_id"] = int(0); }
if (!cfg.contains("metric_type")) { cfg["metric_type"] = "IP"; } // TODO: remove
switch (type) {
case IndexType::FAISS_IVFSQ8_MIX: {
if (!cfg.contains("nbits")) { cfg["nbits"] = int(8); }
break;
}
}
}
}
}
}
cpp/src/wrapper/knowhere/vec_index.h
浏览文件 @ b46168cf
......@@ -31,6 +31,7 @@ enum class IndexType {
FAISS_IVFPQ_CPU,
FAISS_IVFPQ_GPU,
SPTAG_KDT_RNT_CPU,
FAISS_IVFSQ8_MIX,
//NSG,
};
......@@ -75,6 +76,8 @@ extern VecIndexPtr GetVecIndexFactory(const IndexType &type);
extern VecIndexPtr LoadVecIndex(const IndexType &index_type, const zilliz::knowhere::BinarySet &index_binary);
extern void AutoGenParams(const IndexType& type, const long& size, Config& cfg);
}
}
}
knowhere @ f866ac4e
Subproject commit b0b9dd18fadbf9dc0fccaad815e14e578a92993e
Subproject commit f866ac4e297dea477ec591a62679cf5cdd219cc8
cpp/unittest/index_wrapper/knowhere_test.cpp
浏览文件 @ b46168cf
......@@ -41,7 +41,7 @@ class KnowhereWrapperTest
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(ids[i * k], gt_ids[i * k]);
EXPECT_EQ(dis[i * k], gt_dis[i * k]);
//EXPECT_EQ(dis[i * k], gt_dis[i * k]);
}
int match = 0;
......@@ -84,11 +84,11 @@ class KnowhereWrapperTest
INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
Values(
//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"]
std::make_tuple(IndexType::FAISS_IVFFLAT_CPU, "Default",
64, 100000, 10, 10,
Config::object{{"nlist", 100}, {"dim", 64}},
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 10}}
),
//std::make_tuple(IndexType::FAISS_IVFFLAT_CPU, "Default",
// 64, 100000, 10, 10,
// Config::object{{"nlist", 100}, {"dim", 64}},
// Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 10}}
//),
//std::make_tuple(IndexType::FAISS_IVFFLAT_GPU, "Default",
// 64, 10000, 10, 10,
// Config::object{{"nlist", 100}, {"dim", 64}},
......@@ -96,13 +96,18 @@ INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
//),
std::make_tuple(IndexType::FAISS_IVFFLAT_MIX, "Default",
64, 100000, 10, 10,
Config::object{{"nlist", 100}, {"dim", 64}},
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 10}}
Config::object{{"nlist", 1000}, {"dim", 64}, {"metric_type", "L2"}},
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
),
std::make_tuple(IndexType::FAISS_IDMAP, "Default",
64, 100000, 10, 10,
Config::object{{"dim", 64}},
Config::object{{"dim", 64}, {"metric_type", "L2"}},
Config::object{{"dim", 64}, {"k", 10}}
),
std::make_tuple(IndexType::FAISS_IVFSQ8_MIX, "Default",
64, 100000, 10, 10,
Config::object{{"dim", 64}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}},
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
)
//std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default",
// 64, 10000, 10, 10,
......
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或