Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
milvus
milvus
提交
df28a763
M
milvus
项目概览
milvus
/
milvus
12 个月 前同步成功
通知
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,发现更多精彩内容 >>
未验证
提交
df28a763
编写于
11月 15, 2019
作者:
J
Jin Hai
提交者:
GitHub
11月 15, 2019
浏览文件
操作
浏览文件
下载
差异文件
Merge pull request #324 from Yukikaze-CZR/0.6.0
Support new Index type IVFPQ #127
上级
6c996b5a
4d1345b4
变更
16
隐藏空白更改
内联
并排
Showing
16 changed file
with
244 addition
and
98 deletion
+244
-98
CHANGELOG.md
CHANGELOG.md
+2
-1
core/src/db/engine/ExecutionEngine.h
core/src/db/engine/ExecutionEngine.h
+3
-1
core/src/db/engine/ExecutionEngineImpl.cpp
core/src/db/engine/ExecutionEngineImpl.cpp
+8
-0
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
...ex/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
+21
-12
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h
...ndex/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h
+5
-0
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
...index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
+26
-0
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h
...c/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h
+3
-0
core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
...x/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
+1
-1
core/src/index/unittest/test_ivf.cpp
core/src/index/unittest/test_ivf.cpp
+2
-2
core/src/wrapper/ConfAdapter.cpp
core/src/wrapper/ConfAdapter.cpp
+53
-8
core/src/wrapper/ConfAdapter.h
core/src/wrapper/ConfAdapter.h
+7
-0
core/src/wrapper/ConfAdapterMgr.cpp
core/src/wrapper/ConfAdapterMgr.cpp
+1
-0
core/src/wrapper/VecIndex.cpp
core/src/wrapper/VecIndex.cpp
+12
-1
core/src/wrapper/VecIndex.h
core/src/wrapper/VecIndex.h
+2
-0
core/unittest/wrapper/test_wrapper.cpp
core/unittest/wrapper/test_wrapper.cpp
+36
-27
core/unittest/wrapper/utils.h
core/unittest/wrapper/utils.h
+62
-45
未找到文件。
CHANGELOG.md
浏览文件 @
df28a763
...
...
@@ -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
...
...
core/src/db/engine/ExecutionEngine.h
浏览文件 @
df28a763
...
...
@@ -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
{
...
...
core/src/db/engine/ExecutionEngineImpl.cpp
浏览文件 @
df28a763
...
...
@@ -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
;
...
...
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
浏览文件 @
df28a763
...
...
@@ -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
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h
浏览文件 @
df28a763
...
...
@@ -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
;
...
...
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
浏览文件 @
df28a763
...
...
@@ -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
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h
浏览文件 @
df28a763
...
...
@@ -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
;
...
...
core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
浏览文件 @
df28a763
...
...
@@ -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
))
{
...
...
core/src/index/unittest/test_ivf.cpp
浏览文件 @
df28a763
...
...
@@ -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
({
...
...
core/src/wrapper/ConfAdapter.cpp
浏览文件 @
df28a763
...
...
@@ -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
;
}
...
...
core/src/wrapper/ConfAdapter.h
浏览文件 @
df28a763
...
...
@@ -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
{
...
...
core/src/wrapper/ConfAdapterMgr.cpp
浏览文件 @
df28a763
...
...
@@ -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
);
}
...
...
core/src/wrapper/VecIndex.cpp
浏览文件 @
df28a763
...
...
@@ -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
core/src/wrapper/VecIndex.h
浏览文件 @
df28a763
...
...
@@ -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
;
...
...
core/unittest/wrapper/test_wrapper.cpp
浏览文件 @
df28a763
...
...
@@ -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
,
1000
00
,
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
,
1000
00
,
10
,
10
),
std
::
make_tuple
(
milvus
::
engine
::
IndexType
::
FAISS_IVFFLAT_CPU
,
"Default"
,
64
,
1000
00
,
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
(),
search
conf
);
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
(),
search
conf
);
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
(),
search
conf
);
}
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
(),
search
conf
);
}
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
(),
search
conf
);
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
(),
search
conf
);
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
(),
search
conf
);
AssertResult
(
res_ids
,
res_dis
);
}
}
...
...
core/unittest/wrapper/utils.h
浏览文件 @
df28a763
...
...
@@ -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;
//};
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录