鲲鹏社区首页
中文
注册
开发者
支持
积分兑换 NEW
我要评分
文档获取效率
文档正确性
内容完整性
文档易理解
提交
在线提单
论坛求助

使用示例

下方使用示例以使用sift-128-euclidean.hdf5数据集，线程数80为例，数据集可通过以下方式获取：

1
wget http://ann-benchmarks.com/sift-128-euclidean.hdf5 --no-check-certificate

假设程序运行的目录为“/path/to/kbest_test”，完整的目录结构应如下所示：

1
2
3
4
5
6
7
8
├── graph_indices                                           // 存放构建好的图索引，运行时（对应数据集配置文件"save_types"为save_graph）会自动创建
      └── sift-128-euclidean_KGN-RNN_R_50_L_100.kgn         // 构建好的图索引，运行时（对应数据集配置文件"save_types"为save_graph）自动生成
├── searcher_indices                                        // 存放构建好的检索器，运行时（对应数据集配置文件"save_types"为save_searcher）会自动创建
      └── sift-128-euclidean_KGN-RNN_R_50_L_100.kgn         // 构建好的检索器，运行时（对应数据集配置文件"save_types"为save_searcher）自动生成
├── datasets                                                // 存放数据集
      └── sift-128-euclidean.hdf5
├── main.py                                                 // 包含运行函数的文件
└── sift_99.json                                            // 对应数据集配置文件

运行步骤如下：

  1. 假设程序运行的目录为“/path/to/kbest_test”，检查目录下是否存在datasets/sift-128-euclidean.hdf5，main.py，sift_99.json。其中，main.py，sift_99.json将在下方提供。
  2. 确保sift_99.json文件中的“num_numa_nodes”为实际运行时的NUMA数量。
  3. 安装相关依赖。
    1
    pip install scikit-learn h5py psutil numpy==1.24.2
  4. 运行main.py。
    1
    python main.py 80 -1 sift_99.json

    测试指令参数与解释如下所示。

    python main.py <threads> <batch_size> <json_name>
    • “threads”表示实际运行时的线程数。
    • “batch_size”表示batch查询模式下一次查询的query数量，设置为“-1”时表示一次查询数据集中所有的查询query。
    • “json_name”表示测试数据集对应的配置文件名称。

    执行结果如下。

main.py内容如下：
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
import os
import sys
import json
from time import time

import numpy as np
import h5py
import psutil
from sklearn import preprocessing

from kbest import KBest

sys.path.append("../")


class Dataset:                                                                                                          # 数据集基类
    def __init__(self):
        self.name = ""
        self.metric = "L2"
        self.d = -1
        self.nb = -1
        self.nq = -1
        self.base = None
        self.query = None
        self.gt = None
        self.file = None

    def evaluate(self, pred, k=None):
        nq, topk = pred.shape
        if k is not None:
            topk = k
        gt = self.get_groundtruth(topk)
        cnt = 0
        for i in range(nq):
            cnt += np.intersect1d(pred[i], gt[i]).size
        return cnt / nq / k

    def get_base(self):
        ret = np.array(self.file['train'])
        if self.metric == "IP":
            ret = preprocessing.normalize(ret)
        return ret

    def get_queries(self):
        ret = np.array(self.file['test'])
        if self.metric == "IP":
            ret = preprocessing.normalize(ret)
        return ret

    def get_groundtruth(self, k):
        ret = np.array(self.file['neighbors'])
        return ret[:, :k]

    def get_fname(self, dir):
        if dir is None:
            dir = "datasets"
        if not os.path.exists(dir):
            os.mkdir(dir)
        return f'{dir}/{self.name}.hdf5'


class DatasetCustom(Dataset):
    name = ""
    metric = ""

    def __init__(self, path=None):
        self.name = getname(path)
        if "euclidean" in self.name:
            self.metric = "L2"
        elif "L2" in self.name:
            self.metric = "L2"
        elif "angular" in self.name:
            self.metric = "IP"
        elif "IP" in self.name:
            self.metric = "IP"
        else:
            print("[ERROR] only support IP or L2 datasets.")
            exit(-1)
        if not os.path.exists(path):
            print("[ERROR] dataset {} not found.".format(path))
            exit(-1)            
        self.file = h5py.File(path)


class DatasetSIFT1M(Dataset):                                                                                           # 用例所使用的sift-128-euclidean.hdf5数据集子类
    name = "sift-128-euclidean"
    metric = "L2"

    def __init__(self, dir=None):
        path = self.get_fname(dir)
        if not os.path.exists(path):
            os.system(f'wget --no-check-certificate --output-document  {path} {download(self.name)}')
        self.file = h5py.File(path)
        self.nb = self.file['train'].shape[0]
        self.nq = self.file['test'].shape[0]
        self.d = self.file['test'].shape[1]


class Best:
    def __init__(self, name, level, metric, method_param):
        self.metric = metric
        self.R = method_param['R']
        self.L = method_param['L']
        self.A = method_param['A']
        self.index_type = method_param['index_type']
        self.optimize = method_param['optimize']
        self.batch = method_param['batch']
        self.numa_enabled = method_param['numa_enabled']
        self.num_numa_nodes = method_param['num_numa_nodes']
        self.name = 'kgn_(%s)' % (method_param)
        self.dir = 'graph_indices'
        self.path = f'{name}_{self.index_type}_R_{self.R}_L_{self.L}.kgn'
        self.level = level

    def fit_with_seri(self, X):
        pass

    def fit_with_graph(self, X):
        pass

    def fit(self, X, save_types):
        print(save_types)
        self.data_num = X.shape[0]
        self.d = X.shape[1]
        self.index_build_type = ""
        if self.index_type == "KGN-RNN":
            self.index_build_type = "RNNDescent"
        elif self.index_type == "KGN":
            self.index_build_type = "NNDescent"
        else:
            print(f"[ERROR] index build type {self.index_type} do not support!")
            exit(-1)

        if not os.path.exists(self.dir):
            os.mkdir(self.dir)

        output_file = "indices/info.txt"
        if self.path not in os.listdir(self.dir) or save_types == "serialize":
            self.index = KBest(self.d, self.R, self.L, self.A, self.metric, self.index_build_type, False,               # 初始化
                               self.num_numa_nodes)
            print(f"[INFO] {self.path} not found, start to build Index.")
            self.index.add(self.data_num, X, 20, self.level)                                                            # 构建图索引
            print("[INFO] Done add data, now build searcher.")
            index_save_path = os.path.join(self.dir, self.path)   
            if save_types == "save_searcher":
                self.dir = 'searcher_indices'
                self.path = f'{name}_{self.index_type}_R_{self.R}_L_{self.L}.ksn'
                self.index.buildSearcher()                                                                              # 构建检索器
                print(f"[INFO] Done build searcher, now save to {index_save_path}")
                self.index.save(index_save_path)                                                                        # 保存检索器
            elif save_types == "save_graph":
                self.index.saveGraph(index_save_path)                                                                   # 保存图索引
            else:
                self.index.buildSearcher()
                print("[INFO] serialize.")
                data_arr = self.index.serialize()                                                                       # 序列化

        print(f"[INFO] Load Index from {self.path}")
        if save_types == "save_searcher" or save_types == "save_graph":
            if not os.path.exists(self.dir) or self.path not in os.listdir(self.dir):
                print("[ERROR] Saved index not found. Check the save location and make sure you have write permissions.")
                exit(-1)


        index_load_path = os.path.join(self.dir, self.path)

        if save_types == "save_searcher":
            self.index = KBest(self.numa_enabled, self.num_numa_nodes)
            load_result = self.index.load(index_load_path)                                                              # 加载检索器
        elif save_types == "save_graph":
            self.index = KBest(self.numa_enabled, self.num_numa_nodes)
            load_result = self.index.loadGraph(index_load_path)                                                         # 加载图索引
        else:
            print("Index deserialize.")
            self.index = KBest(self.numa_enabled, self.num_numa_nodes)
            load_result = self.index.deserialize(data_arr)                                                              # 反序列化
        if load_result == -1:
            exit(-1)

    def set_query_arguments(self, ef):
        self.index.setEf(ef)                                                                                            # 设置检索时的候选节点列表大小
        self.ef = ef

    def prepare_batch_query(self, queries, topk):
        self.queries = queries
        self.topk = topk
        self.nq = len(queries)
        self.labels = np.empty(self.nq * self.topk,dtype=np.int64)
        self.dis = np.empty(self.nq * self.topk,dtype=np.float32)

    def run_batch_query(self, threads, level):
        self.index.search(self.nq, self.queries, self.topk, self.dis, self.labels, threads)                             # 检索

    def get_batch_results(self):
        return self.labels.reshape(self.nq, -1)

    def freeIndex(self):
        del self.index


if __name__ == '__main__':
    dataset_dict = {'sift-128-euclidean': DatasetSIFT1M}
    threads = int(sys.argv[1])
    query_batch_size = int(sys.argv[2])
    json_path = sys.argv[3]

    with open(json_path, 'r') as f:
        config = json.load(f)

    index_types = config['index_types']
    levels = config['levels']
    efs = config['efs']
    R = config['R']
    L = config['L']
    A = config['A']
    topk = config['topk']
    runs = config['runs']
    batch = config['batch']
    optimize = config['optimize']
    dataset_names = config['datasets']
    numa_enabled = config['numa_enabled']
    num_numa_nodes = config['num_numa_nodes']
    save_types = config['save_types']

    results_dir = "results"
    if not os.path.exists(results_dir):
        os.mkdir(results_dir)

    for dataset_name in dataset_names:
        if ".hdf5" in dataset_name:
            dataset = DatasetCustom(dataset_name)
        else:
            dataset = dataset_dict[dataset_name]()
        base = dataset.get_base()
        query = dataset.get_queries()
        gt = dataset.get_groundtruth(topk)
        name = dataset.name
        metric = dataset.metric
        print(name)
        print(metric)
        print(vars(dataset))
        print(dataset.__dict__)
        nq = len(query)

        for it, index_type in enumerate(index_types):
            for it2, level in enumerate(levels):

                p = Best(name, level, metric, {
                    'index_type': index_type, 'R': R, 'L': L, 'A': A, 'optimize': optimize, 'batch': batch,
                    'numa_enabled': numa_enabled, 'num_numa_nodes': num_numa_nodes})

                t = time()
                p.fit(base, save_types)
                ela = time() - t
                print(f"Building time of index: {ela}s")
                qpss = []
                recalls = []
                print(
                    f"dataset: {name}, index: {index_type}, level: {level}")
                for ef in efs:
                    print(f"  ef: {ef}")
                    p.set_query_arguments(ef)
                    mx_qps = 0.0
                    qps_collection = []
                    recall_collect = 0.0
                    for run in range(runs):
                        T = 0
                        res = np.zeros_like(gt)
                        batch_query_time = []
                        if query_batch_size == -1:
                            p.prepare_batch_query(query, topk)
                            t = time()
                            p.run_batch_query(threads, level)
                            T = time() - t
                            batch_query_time.append(T)
                            res = p.get_batch_results()
                        else:
                            T = 0
                            num_batch = (nq + query_batch_size - 1) // query_batch_size

                            for i in range(num_batch):
                                st = i * query_batch_size
                                en = min(st + query_batch_size, nq)
                                p.prepare_batch_query(query[st:en], topk)
                                t = time()
                                p.run_batch_query(threads, level)
                                T = time() - t
                                batch_query_time.append(T)
                                res[st:en] = p.get_batch_results()

                        cnt = 0
                        for i in range(nq):
                            cnt += np.intersect1d(res[i], gt[i]).size
                        recall = cnt / nq / 10
                        qps = nq / sum(batch_query_time) 
                        print(
                            f"    runs [{run + 1}/{runs}], recall: {recall:.4f}, qps: {qps:.2f}")
                        mx_qps = max(mx_qps, qps)
                        qps_collection.append(qps)
                        recall_collect = recall
                    print("level: {} candListSize: {}".format(level, ef))
                    print("recall: {} qps: {}".format(recall_collect, np.mean(qps_collection)))
                    print(qps_collection)
                    qpss.append(mx_qps)
                    recalls.append(recall)
                recalls = np.array(recalls)
                qpss = np.array(qpss)

                print('index_types', index_types)
                print('efs', efs)
                print('recall', recalls)
                print('qps', qpss)

                p.freeIndex()
sift_99.json内容如下：
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
{
  "datasets": [
    "sift-128-euclidean"
  ],
  "index_types": [     
    "KGN-RNN"
  ],
  "R": 50,
  "L": 100,
  "A": 60,
  "levels": [
    2
  ],
  "topk": 10,
  "efs": [
   72
  ],
  "runs": 10,
  "batch": true,
  "optimize": true,
  "plot": true,
  "numa_enabled": false,
  "num_numa_nodes": 4,
  "save_types": "save_graph"
}
父主题： Python