Pass "nbits" as an argument to IVF_PQ create_index (#4109)

* add nbits choices between 1 and 16

Signed-off-by: sahuang <xiaohaix@student.unimelb.edu.au>

* fix python case

Signed-off-by: sahuang <xiaohai.xu@zilliz.com>

* Fix case

Signed-off-by: sahuang <xiaohai.xu@zilliz.com>
pull/4127/head
Xiaohai Xu 2020-10-29 09:59:09 +08:00 committed by GitHub
parent 1cd0169a82
commit 97b180271a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 48 additions and 73 deletions

View File

@ -743,6 +743,9 @@ ExecutionEngineImpl::BuildKnowhereIndex(const std::string& field_name, const Col
conf[knowhere::meta::ROWS] = row_count;
conf[knowhere::meta::DEVICEID] = gpu_num_;
conf[knowhere::Metric::TYPE] = index_info.metric_name_;
if (!conf.contains(knowhere::IndexParams::nbits)) {
conf[knowhere::IndexParams::nbits] = 8;
}
LOG_ENGINE_DEBUG_ << "Index params: " << conf.dump();
knowhere::IndexMode mode = knowhere::IndexMode::MODE_CPU;
@ -752,7 +755,8 @@ ExecutionEngineImpl::BuildKnowhereIndex(const std::string& field_name, const Col
}
if (index_info.index_type_ == milvus::knowhere::IndexEnum::INDEX_FAISS_IVFPQ) {
auto m = conf[knowhere::IndexParams::m].get<int64_t>();
knowhere::IVFPQConfAdapter::GetValidM(dimension, m, mode);
auto nbits = conf[knowhere::IndexParams::nbits].get<int64_t>();
knowhere::IVFPQConfAdapter::CheckPQParams(dimension, m, nbits, mode);
}
#endif
auto adapter = knowhere::AdapterMgr::GetInstance().GetAdapter(index_info.index_type_);

View File

@ -24,6 +24,8 @@
namespace milvus {
namespace knowhere {
static const int64_t MIN_NBITS = 1;
static const int64_t MAX_NBITS = 16;
static const int64_t MIN_NLIST = 1;
static const int64_t MAX_NLIST = 65536;
static const int64_t MIN_NPROBE = 1;
@ -138,49 +140,42 @@ IVFSQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
bool
IVFPQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
const int64_t DEFAULT_NBITS = 8;
oricfg[knowhere::IndexParams::nbits] = DEFAULT_NBITS;
CheckStrByValues(knowhere::Metric::TYPE, METRICS);
CheckIntByRange(knowhere::meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM);
CheckIntByRange(knowhere::meta::ROWS, DEFAULT_MIN_ROWS, DEFAULT_MAX_ROWS);
CheckIntByRange(knowhere::IndexParams::nlist, MIN_NLIST, MAX_NLIST);
// int64_t nlist = oricfg[knowhere::IndexParams::nlist];
// CheckIntByRange(knowhere::meta::ROWS, nlist, DEFAULT_MAX_ROWS);
CheckIntByRange(knowhere::IndexParams::nbits, MIN_NBITS, MAX_NBITS);
// auto tune params
oricfg[knowhere::IndexParams::nlist] =
MatchNlist(oricfg[knowhere::meta::ROWS].get<int64_t>(), oricfg[knowhere::IndexParams::nlist].get<int64_t>());
auto m = oricfg[knowhere::IndexParams::m].get<int64_t>();
auto nbits = oricfg[knowhere::IndexParams::nbits].get<int64_t>();
auto dimension = oricfg[knowhere::meta::DIM].get<int64_t>();
// Best Practice
// static int64_t MIN_POINTS_PER_CENTROID = 40;
// static int64_t MAX_POINTS_PER_CENTROID = 256;
// CheckIntByRange(knowhere::meta::ROWS, MIN_POINTS_PER_CENTROID * nlist, MAX_POINTS_PER_CENTROID * nlist);
/*std::vector<int64_t> resset;
IVFPQConfAdapter::GetValidCPUM(dimension, resset);*/
IndexMode ivfpq_mode = mode;
return GetValidM(dimension, m, ivfpq_mode);
return CheckPQParams(dimension, m, nbits, ivfpq_mode);
}
bool
IVFPQConfAdapter::GetValidM(int64_t dimension, int64_t m, IndexMode& mode) {
IVFPQConfAdapter::CheckPQParams(int64_t dimension, int64_t m, int64_t nbits, IndexMode& mode) {
#ifdef MILVUS_GPU_VERSION
if (mode == knowhere::IndexMode::MODE_GPU && !IVFPQConfAdapter::GetValidGPUM(dimension, m)) {
if (mode == knowhere::IndexMode::MODE_GPU && !IVFPQConfAdapter::CheckGPUPQParams(dimension, m, nbits)) {
mode = knowhere::IndexMode::MODE_CPU;
}
#endif
if (mode == knowhere::IndexMode::MODE_CPU && !IVFPQConfAdapter::GetValidCPUM(dimension, m)) {
if (mode == knowhere::IndexMode::MODE_CPU && !IVFPQConfAdapter::CheckCPUPQParams(dimension, m)) {
return false;
}
return true;
}
bool
IVFPQConfAdapter::GetValidGPUM(int64_t dimension, int64_t m) {
IVFPQConfAdapter::CheckGPUPQParams(int64_t dimension, int64_t m, int64_t nbits) {
/*
* Faiss 1.6
* Only 1, 2, 3, 4, 6, 8, 10, 12, 16, 20, 24, 28, 32 dims per sub-quantizer are currently supported with
@ -193,7 +188,8 @@ IVFPQConfAdapter::GetValidGPUM(int64_t dimension, int64_t m) {
return (std::find(std::begin(support_subquantizer), std::end(support_subquantizer), m) !=
support_subquantizer.end()) &&
(std::find(std::begin(support_dim_per_subquantizer), std::end(support_dim_per_subquantizer), sub_dim) !=
support_dim_per_subquantizer.end());
support_dim_per_subquantizer.end()) &&
(nbits == 8);
/*resset.clear();
for (const auto& dimperquantizer : support_dim_per_subquantizer) {
@ -208,7 +204,7 @@ IVFPQConfAdapter::GetValidGPUM(int64_t dimension, int64_t m) {
}
bool
IVFPQConfAdapter::GetValidCPUM(int64_t dimension, int64_t m) {
IVFPQConfAdapter::CheckCPUPQParams(int64_t dimension, int64_t m) {
return (dimension % m == 0);
}
@ -289,7 +285,7 @@ RHNSWPQConfAdapter::CheckTrain(Config& oricfg, const IndexMode mode) {
auto dimension = oricfg[knowhere::meta::DIM].get<int64_t>();
IVFPQConfAdapter::GetValidCPUM(dimension, oricfg[knowhere::IndexParams::PQM].get<int64_t>());
IVFPQConfAdapter::CheckCPUPQParams(dimension, oricfg[knowhere::IndexParams::PQM].get<int64_t>());
return ConfAdapter::CheckTrain(oricfg, mode);
}

View File

@ -52,13 +52,13 @@ class IVFPQConfAdapter : public IVFConfAdapter {
CheckTrain(Config& oricfg, const IndexMode mode) override;
static bool
GetValidM(int64_t dimension, int64_t m, IndexMode& mode);
CheckPQParams(int64_t dimension, int64_t m, int64_t nbits, IndexMode& mode);
static bool
GetValidGPUM(int64_t dimension, int64_t m);
CheckGPUPQParams(int64_t dimension, int64_t m, int64_t nbits);
static bool
GetValidCPUM(int64_t dimension, int64_t m);
CheckCPUPQParams(int64_t dimension, int64_t m);
};
class NSGConfAdapter : public IVFConfAdapter {

View File

@ -51,7 +51,8 @@ IVFPQ::CopyCpuToGpu(const int64_t device_id, const Config& config) {
auto ivfpq_index = dynamic_cast<faiss::IndexIVFPQ*>(index_.get());
int64_t dim = ivfpq_index->d;
int64_t m = ivfpq_index->pq.M;
if (!IVFPQConfAdapter::GetValidGPUM(dim, m)) {
int64_t nbits = ivfpq_index->pq.nbits;
if (!IVFPQConfAdapter::CheckGPUPQParams(dim, m, nbits)) {
return nullptr;
}
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {

View File

@ -41,7 +41,7 @@ IndexIVFPQ::IndexIVFPQ (Index * quantizer, size_t d, size_t nlist,
IndexIVF (quantizer, d, nlist, 0, metric),
pq (d, M, nbits_per_idx)
{
FAISS_THROW_IF_NOT (nbits_per_idx <= 8);
// FAISS_THROW_IF_NOT (nbits_per_idx <= 8);
code_size = pq.code_size;
invlists->code_size = code_size;
is_trained = false;

View File

@ -283,36 +283,20 @@ ValidateIndexParams(const milvus::json& index_params, int64_t dimension, const s
return status;
}
if (index_params.find(knowhere::IndexParams::nbits) != index_params.end()) {
status = CheckParameterRange(index_params, knowhere::IndexParams::nbits, 1, 16);
if (!status.ok()) {
return status;
}
}
// special check for 'm' parameter
int64_t m_value = index_params[knowhere::IndexParams::m];
if (!milvus::knowhere::IVFPQConfAdapter::GetValidCPUM(dimension, m_value)) {
std::string msg = "Invalid m, dimension can't not be divided by m ";
if (!milvus::knowhere::IVFPQConfAdapter::CheckCPUPQParams(dimension, m_value)) {
std::string msg = "Invalid m, dimension cannot be divided by m ";
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_ARGUMENT, msg);
}
/*std::vector<int64_t> resset;
milvus::knowhere::IVFPQConfAdapter::GetValidMList(dimension, resset);
int64_t m_value = index_params[knowhere::IndexParams::m];
if (resset.empty()) {
std::string msg = "Invalid collection dimension, unable to get reasonable values for 'm'";
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_COLLECTION_DIMENSION, msg);
}
auto iter = std::find(std::begin(resset), std::end(resset), m_value);
if (iter == std::end(resset)) {
std::string msg =
"Invalid " + std::string(knowhere::IndexParams::m) + ", must be one of the following values: ";
for (size_t i = 0; i < resset.size(); i++) {
if (i != 0) {
msg += ",";
}
msg += std::to_string(resset[i]);
}
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_ARGUMENT, msg);
}*/
} else if (index_type == knowhere::IndexEnum::INDEX_NSG) {
auto status = CheckParameterRange(index_params, knowhere::IndexParams::search_length, 10, 300);
if (!status.ok()) {
@ -347,35 +331,20 @@ ValidateIndexParams(const milvus::json& index_params, int64_t dimension, const s
if (!status.ok()) {
return status;
}
if (index_params.find(knowhere::IndexParams::nbits) != index_params.end()) {
status = CheckParameterRange(index_params, knowhere::IndexParams::nbits, 1, 16);
if (!status.ok()) {
return status;
}
}
// special check for 'PQM' parameter
int64_t pqm_value = index_params[knowhere::IndexParams::PQM];
if (!milvus::knowhere::IVFPQConfAdapter::GetValidCPUM(dimension, pqm_value)) {
std::string msg = "Invalid m, dimension can't not be divided by m ";
if (!milvus::knowhere::IVFPQConfAdapter::CheckCPUPQParams(dimension, pqm_value)) {
std::string msg = "Invalid m, dimension cannot be divided by m ";
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_ARGUMENT, msg);
}
/*int64_t pqm_value = index_params[knowhere::IndexParams::PQM];
if (resset.empty()) {
std::string msg = "Invalid collection dimension, unable to get reasonable values for 'PQM'";
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_COLLECTION_DIMENSION, msg);
}
auto iter = std::find(std::begin(resset), std::end(resset), pqm_value);
if (iter == std::end(resset)) {
std::string msg =
"Invalid " + std::string(knowhere::IndexParams::PQM) + ", must be one of the following values: ";
for (size_t i = 0; i < resset.size(); i++) {
if (i != 0) {
msg += ",";
}
msg += std::to_string(resset[i]);
}
LOG_SERVER_ERROR_ << msg;
return Status(SERVER_INVALID_ARGUMENT, msg);
}*/
}
} else if (index_type == knowhere::IndexEnum::INDEX_ANNOY) {
auto status = CheckParameterRange(index_params, knowhere::IndexParams::n_trees, 1, 1024);

View File

@ -676,6 +676,9 @@ def gen_invalid_index():
for nlist in gen_invalid_params():
index_param = {"index_type": "IVF_FLAT", "params": {"nlist": nlist}}
index_params.append(index_param)
for nbits in gen_invalid_params() + [0, 17]:
index_param = {"index_type": "IVF_PQ", "params": {"nlist": 1024, "m": 16, "nbits": nbits}}
index_params.append(index_param)
for M in gen_invalid_params():
index_param = {"index_type": "HNSW", "params": {"M": M, "efConstruction": 100}}
index_param = {"index_type": "RHNSW_PQ", "params": {"M": M, "efConstruction": 100}}
@ -717,6 +720,7 @@ def gen_invalid_index():
def gen_index():
nlists = [1, 1024, 16384]
pq_ms = [128, 64, 32, 16, 8, 4]
pq_nbits = [1, 2, 4, 8, 9]
Ms = [5, 24, 48]
efConstructions = [100, 300, 500]
search_lengths = [10, 100, 300]
@ -733,9 +737,10 @@ def gen_index():
for nlist in nlists]
index_params.extend(ivf_params)
elif index_type == "IVF_PQ":
IVFPQ_params = [{"index_type": index_type, "index_param": {"nlist": nlist, "m": m}} \
IVFPQ_params = [{"index_type": index_type, "index_param": {"nlist": nlist, "m": m, "nbits": nbits}} \
for nlist in nlists \
for m in pq_ms]
for m in pq_ms \
for nbits in pq_nbits]
index_params.extend(IVFPQ_params)
elif index_type in ["HNSW", "RHNSW_SQ", "RHNSW_PQ"]:
hnsw_params = [{"index_type": index_type, "index_param": {"M": M, "efConstruction": efConstruction}} \