Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
milvus
milvus
提交
77b28a39
M
milvus
项目概览
milvus
/
milvus
大约 1 年 前同步成功
通知
261
Star
22476
Fork
2472
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
M
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
77b28a39
编写于
10月 10, 2019
作者:
X
xiaojun.lin
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
enable hybrid IVFSQ8
Former-commit-id: 9bbf835e63818ef6cc5f792d0ab11247794d0c3a
上级
d4babbca
变更
8
隐藏空白更改
内联
并排
Showing
8 changed file
with
172 addition
and
61 deletion
+172
-61
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
...knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
+43
-8
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h
...e/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h
+11
-3
cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
...knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
+5
-0
cpp/src/core/test/CMakeLists.txt
cpp/src/core/test/CMakeLists.txt
+1
-0
cpp/src/core/test/test_ivf.cpp
cpp/src/core/test/test_ivf.cpp
+82
-48
cpp/src/wrapper/VecImpl.cpp
cpp/src/wrapper/VecImpl.cpp
+21
-0
cpp/src/wrapper/VecImpl.h
cpp/src/wrapper/VecImpl.h
+3
-0
cpp/src/wrapper/VecIndex.h
cpp/src/wrapper/VecIndex.h
+6
-2
未找到文件。
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
浏览文件 @
77b28a39
...
...
@@ -63,7 +63,7 @@ VectorIndexPtr
IVFSQHybrid
::
CopyGpuToCpu
(
const
Config
&
config
)
{
std
::
lock_guard
<
std
::
mutex
>
lk
(
mutex_
);
if
(
auto
device_idx
=
std
::
dynamic_pointer_cast
<
faiss
::
gpu
::
Gpu
IndexIVF
>
(
index_
))
{
if
(
auto
device_idx
=
std
::
dynamic_pointer_cast
<
faiss
::
IndexIVF
>
(
index_
))
{
faiss
::
Index
*
device_index
=
index_
.
get
();
faiss
::
Index
*
host_index
=
faiss
::
gpu
::
index_gpu_to_cpu
(
device_index
);
...
...
@@ -113,6 +113,7 @@ IVFSQHybrid::search_impl(int64_t n,
if
(
gpu_mode
)
{
GPUIVF
::
search_impl
(
n
,
data
,
k
,
distances
,
labels
,
cfg
);
}
else
{
ResScope
rs
(
res_
,
gpu_id_
);
IVF
::
search_impl
(
n
,
data
,
k
,
distances
,
labels
,
cfg
);
}
}
...
...
@@ -121,7 +122,9 @@ QuantizerPtr
IVFSQHybrid
::
LoadQuantizer
(
const
Config
&
conf
)
{
auto
quantizer_conf
=
std
::
dynamic_pointer_cast
<
QuantizerCfg
>
(
conf
);
if
(
quantizer_conf
!=
nullptr
)
{
quantizer_conf
->
CheckValid
();
// throw exception
if
(
quantizer_conf
->
mode
!=
1
)
{
KNOWHERE_THROW_MSG
(
"mode only support 1 in this func"
);
}
}
gpu_id_
=
quantizer_conf
->
gpu_id
;
...
...
@@ -133,13 +136,14 @@ IVFSQHybrid::LoadQuantizer(const Config &conf) {
auto
index_composition
=
new
faiss
::
IndexComposition
;
index_composition
->
index
=
index_
.
get
();
index_composition
->
quantizer
=
nullptr
;
index_composition
->
mode
=
quantizer_conf
->
mode
;
//
1 or 2
index_composition
->
mode
=
quantizer_conf
->
mode
;
//
only 1
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
gpu_id_
,
index_composition
,
&
option
);
delete
gpu_index
;
std
::
shared_ptr
<
FaissIVFQuantizer
>
q
;
q
->
quantizer
=
index_composition
;
auto
q
=
std
::
make_shared
<
FaissIVFQuantizer
>
();
q
->
quantizer
=
index_composition
->
quantizer
;
res_
=
res
;
return
q
;
}
else
{
KNOWHERE_THROW_MSG
(
"CopyCpuToGpu Error, can't get gpu_resource"
);
...
...
@@ -153,15 +157,13 @@ IVFSQHybrid::SetQuantizer(const QuantizerPtr& q) {
KNOWHERE_THROW_MSG
(
"Quantizer type error"
);
}
if
(
ivf_quantizer
->
quantizer
->
mode
==
2
)
gpu_mode
=
true
;
// all in gpu
faiss
::
IndexIVF
*
ivf_index
=
dynamic_cast
<
faiss
::
IndexIVF
*>
(
index_
.
get
());
faiss
::
gpu
::
GpuIndexFlat
*
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
=
ivf_quantizer
->
quantizer
->
quantizer
;
ivf_index
->
quantizer
=
ivf_quantizer
->
quantizer
;
}
}
...
...
@@ -175,5 +177,38 @@ IVFSQHybrid::UnsetQuantizer() {
ivf_index
->
quantizer
=
nullptr
;
}
void
IVFSQHybrid
::
LoadData
(
const
knowhere
::
QuantizerPtr
&
q
,
const
Config
&
conf
)
{
auto
quantizer_conf
=
std
::
dynamic_pointer_cast
<
QuantizerCfg
>
(
conf
);
if
(
quantizer_conf
!=
nullptr
)
{
if
(
quantizer_conf
->
mode
!=
2
)
{
KNOWHERE_THROW_MSG
(
"mode only support 2 in this func"
);
}
}
if
(
quantizer_conf
->
gpu_id
!=
gpu_id_
)
{
KNOWHERE_THROW_MSG
(
"quantizer and data must on the same gpu card"
);
}
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
))
{
ResScope
rs
(
res
,
gpu_id_
,
false
);
faiss
::
gpu
::
GpuClonerOptions
option
;
option
.
allInGpu
=
true
;
auto
ivf_quantizer
=
std
::
dynamic_pointer_cast
<
FaissIVFQuantizer
>
(
q
);
if
(
ivf_quantizer
==
nullptr
)
KNOWHERE_THROW_MSG
(
"quantizer type not faissivfquantizer"
);
auto
index_composition
=
new
faiss
::
IndexComposition
;
index_composition
->
index
=
index_
.
get
();
index_composition
->
quantizer
=
ivf_quantizer
->
quantizer
;
index_composition
->
mode
=
quantizer_conf
->
mode
;
// only 2
auto
gpu_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
gpu_id_
,
index_composition
,
&
option
);
index_
.
reset
(
gpu_index
);
gpu_mode
=
true
;
// all in gpu
}
else
{
KNOWHERE_THROW_MSG
(
"CopyCpuToGpu Error, can't get gpu_resource"
);
}
}
}
}
cpp/src/core/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.h
浏览文件 @
77b28a39
...
...
@@ -29,15 +29,20 @@ namespace zilliz {
namespace
knowhere
{
struct
FaissIVFQuantizer
:
public
Quantizer
{
faiss
::
IndexComposition
*
quantizer
=
nullptr
;
faiss
::
gpu
::
GpuIndexFlat
*
quantizer
=
nullptr
;
};
using
FaissIVFQuantizerPtr
=
std
::
shared_ptr
<
FaissIVFQuantizer
>
;
class
IVFSQHybrid
:
public
GPUIVFSQ
{
public:
explicit
IVFSQHybrid
(
const
int
&
device_id
)
:
GPUIVFSQ
(
device_id
)
{}
explicit
IVFSQHybrid
(
const
int
&
device_id
)
:
GPUIVFSQ
(
device_id
)
{
gpu_mode
=
false
;
}
explicit
IVFSQHybrid
(
std
::
shared_ptr
<
faiss
::
Index
>
index
)
:
GPUIVFSQ
(
-
1
)
{
gpu_mode
=
false
;}
explicit
IVFSQHybrid
(
std
::
shared_ptr
<
faiss
::
Index
>
index
)
:
GPUIVFSQ
(
-
1
)
{
index_
=
index
;
gpu_mode
=
false
;
}
explicit
IVFSQHybrid
(
std
::
shared_ptr
<
faiss
::
Index
>
index
,
const
int64_t
&
device_id
,
ResPtr
&
resource
)
:
GPUIVFSQ
(
index
,
device_id
,
resource
)
{
...
...
@@ -54,6 +59,9 @@ class IVFSQHybrid : public GPUIVFSQ {
void
UnsetQuantizer
();
void
LoadData
(
const
knowhere
::
QuantizerPtr
&
q
,
const
Config
&
conf
);
IndexModelPtr
Train
(
const
DatasetPtr
&
dataset
,
const
Config
&
config
)
override
;
...
...
cpp/src/core/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h
浏览文件 @
77b28a39
...
...
@@ -100,6 +100,11 @@ public:
Lock
();
}
ResScope
(
ResWPtr
&
res
,
const
int64_t
&
device_id
,
const
bool
&
isown
)
:
resource
(
res
),
device_id
(
device_id
),
move
(
true
),
own
(
isown
)
{
Lock
();
}
// specif for search
// get the ownership of gpuresource and gpu
ResScope
(
ResWPtr
&
res
,
const
int64_t
&
device_id
)
...
...
cpp/src/core/test/CMakeLists.txt
浏览文件 @
77b28a39
...
...
@@ -41,6 +41,7 @@ set(ivf_srcs
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
${
CORE_SOURCE_DIR
}
/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp
)
...
...
cpp/src/core/test/test_ivf.cpp
浏览文件 @
77b28a39
...
...
@@ -65,6 +65,8 @@ IVFIndexPtr IndexFactory(const std::string &type) {
return
std
::
make_shared
<
IVFSQ
>
();
}
else
if
(
type
==
"GPUIVFSQ"
)
{
return
std
::
make_shared
<
GPUIVFSQ
>
(
device_id
);
}
else
if
(
type
==
"IVFSQHybrid"
)
{
return
std
::
make_shared
<
IVFSQHybrid
>
(
device_id
);
}
}
...
...
@@ -72,6 +74,7 @@ enum class ParameterType {
ivf
,
ivfpq
,
ivfsq
,
ivfsqhybrid
,
nsg
,
};
...
...
@@ -105,7 +108,7 @@ class ParamGenerator {
tempconf
->
metric_type
=
METRICTYPE
::
L2
;
return
tempconf
;
}
else
if
(
type
==
ParameterType
::
ivfsq
)
{
else
if
(
type
==
ParameterType
::
ivfsq
||
type
==
ParameterType
::
ivfsqhybrid
)
{
auto
tempconf
=
std
::
make_shared
<
IVFSQCfg
>
();
tempconf
->
d
=
DIM
;
tempconf
->
gpu_id
=
device_id
;
...
...
@@ -160,6 +163,7 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
// std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std
::
make_tuple
(
"IVFSQ"
,
ParameterType
::
ivfsq
),
std
::
make_tuple
(
"GPUIVFSQ"
,
ParameterType
::
ivfsq
)
std
::
make_tuple
(
"IVFSQHybrid"
,
ParameterType
::
ivfsqhybrid
)
)
);
...
...
@@ -210,36 +214,60 @@ TEST_P(IVFTest, ivf_basic) {
//PrintResult(result, nq, k);
}
//TEST_P(IVFTest, hybrid) {
// assert(!xb.empty());
//
// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
// index_->set_preprocessor(preprocessor);
//
// auto model = index_->Train(base_dataset, conf);
// index_->set_index_model(model);
// index_->Add(base_dataset, conf);
// EXPECT_EQ(index_->Count(), nb);
// EXPECT_EQ(index_->Dimension(), dim);
//
//// auto new_idx = ChooseTodo();
//// auto result = new_idx->Search(query_dataset, conf);
//// AssertAnns(result, nq, conf->k);
//
// auto iss_idx = std::make_shared<IVFSQHybrid>(device_id);
//
// auto binaryset = index_->Serialize();
// iss_idx->Load(binaryset);
//
// auto quantizer_conf = std::make_shared<QuantizerCfg>();
// quantizer_conf->mode = 1;
// quantizer_conf->gpu_id = 1;
// auto q = iss_idx->LoadQuantizer(quantizer_conf);
// iss_idx->SetQuantizer(q);
// auto result = iss_idx->Search(query_dataset, conf);
TEST_P
(
IVFTest
,
hybrid
)
{
if
(
index_type
!=
"IVFSQHybrid"
)
{
return
;
}
assert
(
!
xb
.
empty
());
auto
preprocessor
=
index_
->
BuildPreprocessor
(
base_dataset
,
conf
);
index_
->
set_preprocessor
(
preprocessor
);
auto
model
=
index_
->
Train
(
base_dataset
,
conf
);
index_
->
set_index_model
(
model
);
index_
->
Add
(
base_dataset
,
conf
);
EXPECT_EQ
(
index_
->
Count
(),
nb
);
EXPECT_EQ
(
index_
->
Dimension
(),
dim
);
// auto new_idx = ChooseTodo();
// auto result = new_idx->Search(query_dataset, conf);
// AssertAnns(result, nq, conf->k);
// //PrintResult(result, nq, k);
//}
{
auto
hybrid_1_idx
=
std
::
make_shared
<
IVFSQHybrid
>
(
device_id
);
auto
binaryset
=
index_
->
Serialize
();
hybrid_1_idx
->
Load
(
binaryset
);
auto
quantizer_conf
=
std
::
make_shared
<
QuantizerCfg
>
();
quantizer_conf
->
mode
=
1
;
quantizer_conf
->
gpu_id
=
device_id
;
auto
q
=
hybrid_1_idx
->
LoadQuantizer
(
quantizer_conf
);
hybrid_1_idx
->
SetQuantizer
(
q
);
auto
result
=
hybrid_1_idx
->
Search
(
query_dataset
,
conf
);
AssertAnns
(
result
,
nq
,
conf
->
k
);
PrintResult
(
result
,
nq
,
k
);
}
{
auto
hybrid_2_idx
=
std
::
make_shared
<
IVFSQHybrid
>
(
device_id
);
auto
binaryset
=
index_
->
Serialize
();
hybrid_2_idx
->
Load
(
binaryset
);
auto
quantizer_conf
=
std
::
make_shared
<
QuantizerCfg
>
();
quantizer_conf
->
mode
=
1
;
quantizer_conf
->
gpu_id
=
device_id
;
auto
q
=
hybrid_2_idx
->
LoadQuantizer
(
quantizer_conf
);
quantizer_conf
->
mode
=
2
;
hybrid_2_idx
->
LoadData
(
q
,
quantizer_conf
);
auto
result
=
hybrid_2_idx
->
Search
(
query_dataset
,
conf
);
AssertAnns
(
result
,
nq
,
conf
->
k
);
PrintResult
(
result
,
nq
,
k
);
}
}
//TEST_P(IVFTest, gpu_to_cpu) {
// if (index_type.find("GPU") == std::string::npos) { return; }
...
...
@@ -350,29 +378,35 @@ TEST_P(IVFTest, clone_test) {
}
};
// {
// // clone in place
// std::vector<std::string> support_idx_vec{"IVF", "GPUIVF", "IVFPQ", "IVFSQ", "GPUIVFSQ"};
// auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
// if (finder != support_idx_vec.cend()) {
// EXPECT_NO_THROW({
// auto clone_index = index_->Clone();
// auto clone_result = clone_index->Search(query_dataset, conf);
// //AssertAnns(result, nq, conf->k);
// AssertEqual(result, clone_result);
// std::cout << "inplace clone [" << index_type << "] success" << std::endl;
// });
// } else {
// EXPECT_THROW({
// std::cout << "inplace clone [" << index_type << "] failed" << std::endl;
// auto clone_index = index_->Clone();
// }, KnowhereException);
// }
// }
{
// clone in place
std
::
vector
<
std
::
string
>
support_idx_vec
{
"IVF"
,
"GPUIVF"
,
"IVFPQ"
,
"IVFSQ"
,
"GPUIVFSQ"
};
auto
finder
=
std
::
find
(
support_idx_vec
.
cbegin
(),
support_idx_vec
.
cend
(),
index_type
);
if
(
finder
!=
support_idx_vec
.
cend
())
{
EXPECT_NO_THROW
({
auto
clone_index
=
index_
->
Clone
();
auto
clone_result
=
clone_index
->
Search
(
query_dataset
,
conf
);
//AssertAnns(result, nq, conf->k);
AssertEqual
(
result
,
clone_result
);
std
::
cout
<<
"inplace clone ["
<<
index_type
<<
"] success"
<<
std
::
endl
;
});
}
else
{
EXPECT_THROW
({
std
::
cout
<<
"inplace clone ["
<<
index_type
<<
"] failed"
<<
std
::
endl
;
auto
clone_index
=
index_
->
Clone
();
},
KnowhereException
);
if
(
index_type
==
"IVFSQHybrid"
)
{
return
;
}
}
{
// copy from gpu to cpu
std
::
vector
<
std
::
string
>
support_idx_vec
{
"GPUIVF"
,
"GPUIVFSQ"
};
std
::
vector
<
std
::
string
>
support_idx_vec
{
"GPUIVF"
,
"GPUIVFSQ"
,
"IVFSQHybrid"
};
auto
finder
=
std
::
find
(
support_idx_vec
.
cbegin
(),
support_idx_vec
.
cend
(),
index_type
);
if
(
finder
!=
support_idx_vec
.
cend
())
{
EXPECT_NO_THROW
({
...
...
@@ -412,7 +446,7 @@ TEST_P(IVFTest, clone_test) {
TEST_P
(
IVFTest
,
seal_test
)
{
//FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
std
::
vector
<
std
::
string
>
support_idx_vec
{
"GPUIVF"
,
"GPUIVFSQ"
};
std
::
vector
<
std
::
string
>
support_idx_vec
{
"GPUIVF"
,
"GPUIVFSQ"
,
"IVFSQHybrid"
};
auto
finder
=
std
::
find
(
support_idx_vec
.
cbegin
(),
support_idx_vec
.
cend
(),
index_type
);
if
(
finder
==
support_idx_vec
.
cend
())
{
return
;
...
...
cpp/src/wrapper/VecImpl.cpp
浏览文件 @
77b28a39
...
...
@@ -304,6 +304,7 @@ IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
return
Status
::
OK
();
}
Status
...
...
@@ -323,6 +324,26 @@ IVFHybridIndex::UnsetQuantizer() {
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
return
Status
::
OK
();
}
Status
IVFHybridIndex
::
LoadData
(
const
knowhere
::
QuantizerPtr
&
q
,
const
Config
&
conf
)
{
try
{
// TODO(linxj): Hardcode here
if
(
auto
new_idx
=
std
::
dynamic_pointer_cast
<
knowhere
::
IVFSQHybrid
>
(
index_
))
{
new_idx
->
LoadData
(
q
,
conf
);
}
else
{
WRAPPER_LOG_ERROR
<<
"Hybrid mode not support for index type: "
<<
int
(
type
);
return
Status
(
KNOWHERE_ERROR
,
"not support"
);
}
}
catch
(
knowhere
::
KnowhereException
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_UNEXPECTED_ERROR
,
e
.
what
());
}
catch
(
std
::
exception
&
e
)
{
WRAPPER_LOG_ERROR
<<
e
.
what
();
return
Status
(
KNOWHERE_ERROR
,
e
.
what
());
}
return
Status
::
OK
();
}
}
// namespace engine
...
...
cpp/src/wrapper/VecImpl.h
浏览文件 @
77b28a39
...
...
@@ -111,6 +111,9 @@ class IVFHybridIndex : public IVFMixIndex {
Status
UnsetQuantizer
()
override
;
Status
LoadData
(
const
knowhere
::
QuantizerPtr
&
q
,
const
Config
&
conf
)
override
;
};
class
BFIndex
:
public
VecIndexImpl
{
...
...
cpp/src/wrapper/VecIndex.h
浏览文件 @
77b28a39
...
...
@@ -104,15 +104,19 @@ class VecIndex {
Load
(
const
zilliz
::
knowhere
::
BinarySet
&
index_binary
)
=
0
;
// TODO(linxj): refactor later
////////////////
virtual
knowhere
::
QuantizerPtr
LoadQuantizer
(
const
Config
&
conf
)
{
return
Status
::
OK
();
}
LoadQuantizer
(
const
Config
&
conf
)
{
return
nullptr
;
}
virtual
Status
LoadData
(
const
knowhere
::
QuantizerPtr
&
q
,
const
Config
&
conf
)
{
return
Status
::
OK
();
}
// TODO(linxj): refactor later
virtual
Status
SetQuantizer
(
const
knowhere
::
QuantizerPtr
&
q
)
{
return
Status
::
OK
();
}
virtual
Status
UnsetQuantizer
()
{
return
Status
::
OK
();
}
////////////////
};
extern
Status
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录