Add an optional parameter 'nbtis' for IVF_PQ (#4709)

* Add an optional parameter 'nbtis' for IVF_PQ

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>

* fix test

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>
pull/4733/head
shengjun.li 2021-02-19 15:29:25 +08:00
parent e8afb2ba65
commit 64ddc2a556
7 changed files with 38 additions and 9 deletions

View File

@ -12,6 +12,7 @@ Please mark all change in change log and use the issue from GitHub
## Improvement
- \#1970 Improve the performance of BinaryFlat by AVX2
- \#3920 Add an optional parameter 'nbtis' for IVF_PQ
- \#4676,#4614 Support configurable metric labels cluster and instance for Prometheus (implemented by IQIYI)
## Task

View File

@ -40,6 +40,14 @@ namespace knowhere {
return false; \
}
#define CheckIntByRangeIfExist(key, min, max) \
if (oricfg.contains(key)) { \
if (!oricfg[key].is_number_integer() || oricfg[key].get<int64_t>() > max || \
oricfg[key].get<int64_t>() < min) { \
return false; \
} \
}
#define CheckIntByValues(key, container) \
if (!oricfg.contains(key) || !oricfg[key].is_number_integer()) { \
return false; \
@ -157,24 +165,29 @@ IVFPQConfAdapter::CheckTrain(Config& oricfg, IndexMode& mode) {
static int64_t DEFAULT_NBITS = 8;
static int64_t MAX_NLIST = 65536;
static int64_t MIN_NLIST = 1;
static int64_t MAX_NBITS = 16;
static int64_t MIN_NBITS = 1;
static std::vector<std::string> METRICS{Metric::L2, Metric::IP};
oricfg[IndexParams::nbits] = DEFAULT_NBITS;
CheckStrByValues(Metric::TYPE, METRICS);
CheckIntByRange(meta::DIM, DEFAULT_MIN_DIM, DEFAULT_MAX_DIM);
CheckIntByRange(meta::ROWS, DEFAULT_MIN_ROWS, DEFAULT_MAX_ROWS);
CheckIntByRange(IndexParams::nlist, MIN_NLIST, MAX_NLIST);
CheckIntByRangeIfExist(IndexParams::nbits, MIN_NBITS, MAX_NBITS);
auto rows = oricfg[meta::ROWS].get<int64_t>();
auto nlist = oricfg[IndexParams::nlist].get<int64_t>();
auto dimension = oricfg[meta::DIM].get<int64_t>();
auto m = oricfg[IndexParams::m].get<int64_t>();
auto nbits = oricfg.count(IndexParams::nbits) ? oricfg[IndexParams::nbits].get<int64_t>() : DEFAULT_NBITS;
// auto tune params
oricfg[IndexParams::nlist] = MatchNlist(rows, nlist);
oricfg[IndexParams::nbits] = nbits = MatchNbits(rows, nbits);
#ifdef MILVUS_GPU_VERSION
if (mode == IndexMode::MODE_GPU) {
if (IsValidForGPU(dimension, m)) {
if (IsValidForGPU(dimension, m, nbits)) {
return true;
}
// else try CPU Mode
@ -185,7 +198,7 @@ IVFPQConfAdapter::CheckTrain(Config& oricfg, IndexMode& mode) {
}
bool
IVFPQConfAdapter::IsValidForGPU(int64_t dimension, int64_t m) {
IVFPQConfAdapter::IsValidForGPU(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

View File

@ -52,7 +52,7 @@ class IVFPQConfAdapter : public IVFConfAdapter {
CheckTrain(Config& oricfg, IndexMode& mode) override;
static bool
IsValidForGPU(int64_t dimension, int64_t m);
IsValidForGPU(int64_t dimension, int64_t m, int64_t nbits);
static bool
IsValidForCPU(int64_t dimension, int64_t m);

View File

@ -52,7 +52,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::IsValidForGPU(dim, m)) {
int64_t nbits = ivfpq_index->pq.nbits;
if (!IVFPQConfAdapter::IsValidForGPU(dim, m, nbits)) {
return nullptr;
}

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

@ -50,8 +50,12 @@ constexpr int64_t MAX_INSERT_DATA_SIZE = 256 * M_BYTE;
Status
CheckParameterRange(const milvus::json& json_params, const std::string& param_name, int64_t min, int64_t max,
bool min_close = true, bool max_closed = true) {
bool can_ignore = false, bool min_close = true, bool max_closed = true) {
if (json_params.find(param_name) == json_params.end()) {
if (can_ignore) {
return Status::OK();
}
std::string msg = "Parameter list must contain: ";
msg += param_name;
LOG_SERVER_ERROR_ << msg;
@ -213,6 +217,11 @@ ValidationUtil::ValidateIndexParams(const milvus::json& index_params,
return status;
}
status = CheckParameterRange(index_params, knowhere::IndexParams::nbits, 1, 16, true);
if (!status.ok()) {
return status;
}
status = CheckParameterExistence(index_params, knowhere::IndexParams::m);
if (!status.ok()) {
return status;

View File

@ -528,6 +528,9 @@ def gen_invalid_index():
for nlist in gen_invalid_params():
index_param = {"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}
index_params.append(index_param)
for nbits in gen_invalid_params():
index_param = {"index_type": IndexType.IVF_PQ, "index_param": {"nlist": 1024, "m": 16, "nbits": nbits}}
index_params.append(index_param)
for M in gen_invalid_params():
index_param = {"index_type": IndexType.HNSW, "index_param": {"M": M, "efConstruction": 100}}
index_params.append(index_param)
@ -563,6 +566,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]
@ -579,9 +583,10 @@ def gen_index():
for nlist in nlists]
index_params.extend(ivf_params)
elif index_type == IndexType.IVF_PQ:
ivf_pq_params = [{"index_type": index_type, "index_param": {"nlist": nlist, "m": m}} \
ivf_pq_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(ivf_pq_params)
elif index_type == IndexType.HNSW:
hnsw_params = [{"index_type": index_type, "index_param": {"M": M, "efConstruction": efConstruction}} \