Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
BaiXuePrincess
milvus
提交
878b69c6
milvus
项目概览
BaiXuePrincess
/
milvus
与 Fork 源项目一致
从无法访问的项目Fork
通知
7
Star
4
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
milvus
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
提交
878b69c6
编写于
9月 10, 2019
作者:
X
xj.lin
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
MS-538 1. update unittest
Former-commit-id: 2a235faa100107e9d421399ee5ce77b8f972d539
上级
4b178583
变更
8
隐藏空白更改
内联
并排
Showing
8 changed file
with
102 addition
and
62 deletion
+102
-62
cpp/src/core/include/knowhere/index/vector_index/idmap.h
cpp/src/core/include/knowhere/index/vector_index/idmap.h
+12
-0
cpp/src/core/src/knowhere/index/preprocessor/normalize.cpp
cpp/src/core/src/knowhere/index/preprocessor/normalize.cpp
+22
-22
cpp/src/core/src/knowhere/index/vector_index/cpu_kdt_rng.cpp
cpp/src/core/src/knowhere/index/vector_index/cpu_kdt_rng.cpp
+6
-6
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
+4
-14
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
+17
-1
cpp/src/core/test/test_idmap.cpp
cpp/src/core/test/test_idmap.cpp
+22
-13
cpp/src/core/test/test_ivf.cpp
cpp/src/core/test/test_ivf.cpp
+10
-6
cpp/src/core/test/test_nsg/test_nsg.cpp
cpp/src/core/test/test_nsg/test_nsg.cpp
+9
-0
未找到文件。
cpp/src/core/include/knowhere/index/vector_index/idmap.h
浏览文件 @
878b69c6
...
...
@@ -32,6 +32,12 @@ class IDMAP : public VectorIndex, public BasicIndex {
virtual
int64_t
*
GetRawIds
();
protected:
virtual
void
search_impl
(
int64_t
n
,
const
float
*
data
,
int64_t
k
,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
);
std
::
mutex
mutex_
;
};
...
...
@@ -49,6 +55,12 @@ class GPUIDMAP : public IDMAP, public GPUIndex {
VectorIndexPtr
CopyGpuToGpu
(
const
int64_t
&
device_id
,
const
Config
&
config
)
override
;
protected:
void
search_impl
(
int64_t
n
,
const
float
*
data
,
int64_t
k
,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
)
override
;
BinarySet
SerializeImpl
()
override
;
void
LoadImpl
(
const
BinarySet
&
index_binary
)
override
;
};
...
...
cpp/src/core/src/knowhere/index/preprocessor/normalize.cpp
浏览文件 @
878b69c6
...
...
@@ -9,32 +9,32 @@ namespace knowhere {
DatasetPtr
NormalizePreprocessor
::
Preprocess
(
const
DatasetPtr
&
dataset
)
{
// TODO: wrap dataset->tensor
auto
tensor
=
dataset
->
tensor
()[
0
];
auto
p_data
=
(
float
*
)
tensor
->
raw_mutable_data
();
auto
dimension
=
tensor
->
shape
()[
1
];
auto
rows
=
tensor
->
shape
()[
0
];
#pragma omp parallel for
for
(
auto
i
=
0
;
i
<
rows
;
++
i
)
{
Normalize
(
&
(
p_data
[
i
*
dimension
]),
dimension
);
}
//
// TODO: wrap dataset->tensor
//
auto tensor = dataset->tensor()[0];
//
auto p_data = (float *)tensor->raw_mutable_data();
//
auto dimension = tensor->shape()[1];
//
auto rows = tensor->shape()[0];
//
//
#pragma omp parallel for
//
for (auto i = 0; i < rows; ++i) {
//
Normalize(&(p_data[i * dimension]), dimension);
//
}
}
void
NormalizePreprocessor
::
Normalize
(
float
*
arr
,
int64_t
dimension
)
{
double
vector_length
=
0
;
for
(
auto
j
=
0
;
j
<
dimension
;
j
++
)
{
double
val
=
arr
[
j
];
vector_length
+=
val
*
val
;
}
vector_length
=
std
::
sqrt
(
vector_length
);
if
(
vector_length
<
1e-6
)
{
auto
val
=
(
float
)
(
1.0
/
std
::
sqrt
((
double
)
dimension
));
for
(
int
j
=
0
;
j
<
dimension
;
j
++
)
arr
[
j
]
=
val
;
}
else
{
for
(
int
j
=
0
;
j
<
dimension
;
j
++
)
arr
[
j
]
=
(
float
)
(
arr
[
j
]
/
vector_length
);
}
//
double vector_length = 0;
//
for (auto j = 0; j < dimension; j++) {
//
double val = arr[j];
//
vector_length += val * val;
//
}
//
vector_length = std::sqrt(vector_length);
//
if (vector_length < 1e-6) {
//
auto val = (float) (1.0 / std::sqrt((double) dimension));
//
for (int j = 0; j < dimension; j++) arr[j] = val;
//
} else {
//
for (int j = 0; j < dimension; j++) arr[j] = (float) (arr[j] / vector_length);
//
}
}
}
// namespace knowhere
...
...
cpp/src/core/src/knowhere/index/vector_index/cpu_kdt_rng.cpp
浏览文件 @
878b69c6
...
...
@@ -70,10 +70,10 @@ CPUKDTRNG::Train(const DatasetPtr &origin, const Config &train_config) {
SetParameters
(
train_config
);
DatasetPtr
dataset
=
origin
->
Clone
();
if
(
index_ptr_
->
GetDistCalcMethod
()
==
SPTAG
::
DistCalcMethod
::
Cosine
&&
preprocessor_
)
{
//
if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
//
&& preprocessor_) {
preprocessor_
->
Preprocess
(
dataset
);
}
//
}
auto
vectorset
=
ConvertToVectorSet
(
dataset
);
auto
metaset
=
ConvertToMetadataSet
(
dataset
);
...
...
@@ -88,10 +88,10 @@ CPUKDTRNG::Add(const DatasetPtr &origin, const Config &add_config) {
SetParameters
(
add_config
);
DatasetPtr
dataset
=
origin
->
Clone
();
if
(
index_ptr_
->
GetDistCalcMethod
()
==
SPTAG
::
DistCalcMethod
::
Cosine
&&
preprocessor_
)
{
//
if (index_ptr_->GetDistCalcMethod() == SPTAG::DistCalcMethod::Cosine
//
&& preprocessor_) {
preprocessor_
->
Preprocess
(
dataset
);
}
//
}
auto
vectorset
=
ConvertToVectorSet
(
dataset
);
auto
metaset
=
ConvertToMetadataSet
(
dataset
);
...
...
cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
浏览文件 @
878b69c6
...
...
@@ -280,15 +280,15 @@ void FaissGpuResourceMgr::InitResource() {
is_init
=
true
;
std
::
cout
<<
"InitResource"
<<
std
::
endl
;
//
std::cout << "InitResource" << std::endl;
for
(
auto
&
device
:
devices_params_
)
{
auto
&
device_id
=
device
.
first
;
std
::
cout
<<
"Device Id: "
<<
device_id
<<
std
::
endl
;
//
std::cout << "Device Id: " << device_id << std::endl;
auto
&
device_param
=
device
.
second
;
auto
&
bq
=
idle_map
[
device_id
];
for
(
int64_t
i
=
0
;
i
<
device_param
.
resource_num
;
++
i
)
{
std
::
cout
<<
"Resource Id: "
<<
i
<<
std
::
endl
;
//
std::cout << "Resource Id: " << i << std::endl;
auto
raw_resource
=
std
::
make_shared
<
faiss
::
gpu
::
StandardGpuResources
>
();
// TODO(linxj): enable set pinned memory
...
...
@@ -298,7 +298,7 @@ void FaissGpuResourceMgr::InitResource() {
bq
.
Put
(
res_wrapper
);
}
}
std
::
cout
<<
"End initResource"
<<
std
::
endl
;
//
std::cout << "End initResource" << std::endl;
}
ResPtr
FaissGpuResourceMgr
::
GetRes
(
const
int64_t
&
device_id
,
...
...
@@ -315,16 +315,6 @@ ResPtr FaissGpuResourceMgr::GetRes(const int64_t &device_id,
return
nullptr
;
}
//bool FaissGpuResourceMgr::GetRes(const int64_t &device_id,
// ResPtr &res,
// const int64_t &alloc_size) {
// InitResource();
//
// std::lock_guard<std::mutex> lk(res->mutex);
// AllocateTempMem(res, device_id, alloc_size);
// return true;
//}
void
FaissGpuResourceMgr
::
MoveToIdle
(
const
int64_t
&
device_id
,
const
ResPtr
&
res
)
{
auto
finder
=
idle_map
.
find
(
device_id
);
if
(
finder
!=
idle_map
.
end
())
{
...
...
cpp/src/core/src/knowhere/index/vector_index/idmap.cpp
浏览文件 @
878b69c6
...
...
@@ -50,7 +50,7 @@ DatasetPtr IDMAP::Search(const DatasetPtr &dataset, const Config &config) {
auto
res_ids
=
(
int64_t
*
)
malloc
(
sizeof
(
int64_t
)
*
elems
);
auto
res_dis
=
(
float
*
)
malloc
(
sizeof
(
float
)
*
elems
);
index_
->
search
(
rows
,
(
float
*
)
p_data
,
k
,
res_dis
,
res_ids
);
search_impl
(
rows
,
(
float
*
)
p_data
,
k
,
res_dis
,
res_ids
,
Config
()
);
auto
id_buf
=
MakeMutableBufferSmart
((
uint8_t
*
)
res_ids
,
sizeof
(
int64_t
)
*
elems
);
auto
dist_buf
=
MakeMutableBufferSmart
((
uint8_t
*
)
res_dis
,
sizeof
(
float
)
*
elems
);
...
...
@@ -72,6 +72,11 @@ DatasetPtr IDMAP::Search(const DatasetPtr &dataset, const Config &config) {
return
std
::
make_shared
<
Dataset
>
(
array
,
nullptr
);
}
void
IDMAP
::
search_impl
(
int64_t
n
,
const
float
*
data
,
int64_t
k
,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
)
{
index_
->
search
(
n
,
(
float
*
)
data
,
k
,
distances
,
labels
);
}
void
IDMAP
::
Add
(
const
DatasetPtr
&
dataset
,
const
Config
&
config
)
{
if
(
!
index_
)
{
KNOWHERE_THROW_MSG
(
"index not initialize"
);
...
...
@@ -207,6 +212,7 @@ void GPUIDMAP::LoadImpl(const BinarySet &index_binary) {
if
(
auto
res
=
FaissGpuResourceMgr
::
GetInstance
().
GetRes
(
gpu_id_
)
){
ResScope
rs
(
gpu_id_
,
res
);
res_
=
res
;
auto
device_index
=
faiss
::
gpu
::
index_cpu_to_gpu
(
res
->
faiss_res
.
get
(),
gpu_id_
,
index
);
index_
.
reset
(
device_index
);
}
else
{
...
...
@@ -230,5 +236,15 @@ int64_t *GPUIDMAP::GetRawIds() {
KNOWHERE_THROW_MSG
(
"Not support"
);
}
void
GPUIDMAP
::
search_impl
(
int64_t
n
,
const
float
*
data
,
int64_t
k
,
float
*
distances
,
int64_t
*
labels
,
const
Config
&
cfg
)
{
ResScope
rs
(
res_
);
index_
->
search
(
n
,
(
float
*
)
data
,
k
,
distances
,
labels
);
}
}
}
cpp/src/core/test/test_idmap.cpp
浏览文件 @
878b69c6
...
...
@@ -12,6 +12,7 @@
#include "knowhere/index/vector_index/idmap.h"
#include "knowhere/adapter/structure.h"
#include "knowhere/index/vector_index/cloner.h"
#include "knowhere/common/exception.h"
#include "utils.h"
...
...
@@ -65,19 +66,20 @@ void PrintResult(const DatasetPtr &result,
}
TEST_F
(
IDMAPTest
,
idmap_basic
)
{
assert
(
!
xb
.
empty
());
ASSERT_TRUE
(
!
xb
.
empty
());
Config
Default_cfg
;
index_
->
Train
(
Config
::
object
{{
"dim"
,
dim
},
{
"metric_type"
,
"L2"
}});
index_
->
Add
(
base_dataset
,
Default_cfg
);
EXPECT_EQ
(
index_
->
Count
(),
nb
);
EXPECT_EQ
(
index_
->
Dimension
(),
dim
);
assert
(
index_
->
GetRawVectors
()
!=
nullptr
);
assert
(
index_
->
GetRawIds
()
!=
nullptr
);
ASSERT_TRUE
(
index_
->
GetRawVectors
()
!=
nullptr
);
ASSERT_TRUE
(
index_
->
GetRawIds
()
!=
nullptr
);
auto
result
=
index_
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
result
,
nq
,
k
);
PrintResult
(
result
,
nq
,
k
);
index_
->
Seal
();
auto
binaryset
=
index_
->
Serialize
();
auto
new_index
=
std
::
make_shared
<
IDMAP
>
();
new_index
->
Load
(
binaryset
);
...
...
@@ -126,15 +128,15 @@ TEST_F(IDMAPTest, idmap_serialize) {
}
TEST_F
(
IDMAPTest
,
copy_test
)
{
assert
(
!
xb
.
empty
());
ASSERT_TRUE
(
!
xb
.
empty
());
Config
Default_cfg
;
index_
->
Train
(
Config
::
object
{{
"dim"
,
dim
},
{
"metric_type"
,
"L2"
}});
index_
->
Add
(
base_dataset
,
Default_cfg
);
EXPECT_EQ
(
index_
->
Count
(),
nb
);
EXPECT_EQ
(
index_
->
Dimension
(),
dim
);
assert
(
index_
->
GetRawVectors
()
!=
nullptr
);
assert
(
index_
->
GetRawIds
()
!=
nullptr
);
ASSERT_TRUE
(
index_
->
GetRawVectors
()
!=
nullptr
);
ASSERT_TRUE
(
index_
->
GetRawIds
()
!=
nullptr
);
auto
result
=
index_
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
result
,
nq
,
k
);
//PrintResult(result, nq, k);
...
...
@@ -151,8 +153,16 @@ TEST_F(IDMAPTest, copy_test) {
auto
clone_index
=
CopyCpuToGpu
(
index_
,
device_id
,
Config
());
auto
clone_result
=
clone_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
clone_result
,
nq
,
k
);
//assert(std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawVectors() != nullptr);
//assert(std::static_pointer_cast<GPUIDMAP>(clone_index)->GetRawIds() != nullptr);
ASSERT_THROW
({
std
::
static_pointer_cast
<
GPUIDMAP
>
(
clone_index
)
->
GetRawVectors
();
},
zilliz
::
knowhere
::
KnowhereException
);
ASSERT_THROW
({
std
::
static_pointer_cast
<
GPUIDMAP
>
(
clone_index
)
->
GetRawIds
();
},
zilliz
::
knowhere
::
KnowhereException
);
auto
binary
=
clone_index
->
Serialize
();
clone_index
->
Load
(
binary
);
auto
new_result
=
clone_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
new_result
,
nq
,
k
);
auto
clone_gpu_idx
=
clone_index
->
Clone
();
auto
clone_gpu_res
=
clone_gpu_idx
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
clone_gpu_res
,
nq
,
k
);
...
...
@@ -161,14 +171,13 @@ TEST_F(IDMAPTest, copy_test) {
auto
host_index
=
CopyGpuToCpu
(
clone_index
,
Config
());
auto
host_result
=
host_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
host_result
,
nq
,
k
);
assert
(
std
::
static_pointer_cast
<
IDMAP
>
(
host_index
)
->
GetRawVectors
()
!=
nullptr
);
assert
(
std
::
static_pointer_cast
<
IDMAP
>
(
host_index
)
->
GetRawIds
()
!=
nullptr
);
ASSERT_TRUE
(
std
::
static_pointer_cast
<
IDMAP
>
(
host_index
)
->
GetRawVectors
()
!=
nullptr
);
ASSERT_TRUE
(
std
::
static_pointer_cast
<
IDMAP
>
(
host_index
)
->
GetRawIds
()
!=
nullptr
);
// gpu to gpu
auto
device_index
=
CopyCpuToGpu
(
index_
,
device_id
,
Config
());
auto
device_result
=
device_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
auto
new_device_index
=
std
::
static_pointer_cast
<
GPUIDMAP
>
(
device_index
)
->
CopyGpuToGpu
(
device_id
,
Config
());
auto
device_result
=
new_device_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
device_result
,
nq
,
k
);
//assert(std::static_pointer_cast<GPUIDMAP>(device_index)->GetRawVectors() != nullptr);
//assert(std::static_pointer_cast<GPUIDMAP>(device_index)->GetRawIds() != nullptr);
}
}
cpp/src/core/test/test_ivf.cpp
浏览文件 @
878b69c6
...
...
@@ -394,8 +394,11 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
{
index_type
=
"GPUIVF"
;
index_
=
IndexFactory
(
index_type
);
index_
=
std
::
make_shared
<
GPUIVF
>
(
-
1
);
ASSERT_EQ
(
std
::
dynamic_pointer_cast
<
GPUIVF
>
(
index_
)
->
GetGpuDevice
(),
-
1
);
std
::
dynamic_pointer_cast
<
GPUIVF
>
(
index_
)
->
SetGpuDevice
(
device_id
);
ASSERT_EQ
(
std
::
dynamic_pointer_cast
<
GPUIVF
>
(
index_
)
->
GetGpuDevice
(),
device_id
);
auto
preprocessor
=
index_
->
BuildPreprocessor
(
base_dataset
,
preprocess_cfg
);
index_
->
set_preprocessor
(
preprocessor
);
train_cfg
=
Config
::
object
{{
"nlist"
,
1638
},
{
"gpu_id"
,
device_id
},
{
"metric_type"
,
"L2"
}};
...
...
@@ -412,8 +415,9 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSectio
n
(
"search all"
);
tc
.
ElapseFromBegi
n
(
"search all"
);
}
FaissGpuResourceMgr
::
GetInstance
().
Dump
();
{
// IVF-Search
...
...
@@ -430,7 +434,7 @@ TEST_F(GPURESTEST, gpu_ivf_resource_test) {
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSectio
n
(
"search all"
);
tc
.
ElapseFromBegi
n
(
"search all"
);
}
}
...
...
@@ -461,7 +465,7 @@ TEST_F(GPURESTEST, gpuivfsq) {
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSectio
n
(
"search all"
);
tc
.
ElapseFromBegi
n
(
"search all"
);
}
{
...
...
@@ -493,7 +497,7 @@ TEST_F(GPURESTEST, gpuivfsq) {
if
(
i
>
search_count
-
6
||
i
<
5
)
tc
.
RecordSection
(
"search once"
);
}
tc
.
RecordSectio
n
(
"search all"
);
tc
.
ElapseFromBegi
n
(
"search all"
);
delete
cpu_index
;
delete
search_idx
;
}
...
...
cpp/src/core/test/test_nsg/test_nsg.cpp
浏览文件 @
878b69c6
...
...
@@ -7,6 +7,7 @@
#include <gtest/gtest.h>
#include <memory>
#include "knowhere/common/exception.h"
#include "knowhere/index/vector_index/gpu_ivf.h"
#include "knowhere/index/vector_index/nsg_index.h"
#include "knowhere/index/vector_index/nsg/nsg_io.h"
...
...
@@ -71,6 +72,14 @@ TEST_P(NSGInterfaceTest, basic_test) {
auto
new_result
=
new_index
->
Search
(
query_dataset
,
Config
::
object
{{
"k"
,
k
}});
AssertAnns
(
result
,
nq
,
k
);
ASSERT_EQ
(
index_
->
Count
(),
nb
);
ASSERT_EQ
(
index_
->
Dimension
(),
dim
);
ASSERT_THROW
({
index_
->
Clone
();},
zilliz
::
knowhere
::
KnowhereException
);
ASSERT_NO_THROW
({
index_
->
Add
(
base_dataset
,
Config
());
index_
->
Seal
();
});
{
//std::cout << "k = 1" << std::endl;
//new_index->Search(GenQuery(1), Config::object{{"k", 1}});
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录