Examples

This section uses the sift-128-euclidean.hdf5 dataset with 80 threads as an example. Run the following command to obtain the dataset:

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

Assume that the directory where the program runs is /path/to/kbest_test. The complete directory structure is as follows:

├── graph_indices                                           // Store the built graph index, which is automatically created during run time. (In the corresponding dataset configuration file, save_types is set to save_graph.)
      └── sift-128-euclidean_KGN-RNN_R_50_L_100.kgn         // The built graph index, which is automatically generated during run time. (In the corresponding dataset configuration file, save_types is set to save_graph.)
├── searcher_indices                                        // Store the built searcher, which is automatically created during run time. (In the corresponding dataset configuration file, save_types is set to save_searcher.)
      └── sift-128-euclidean_KGN-RNN_R_50_L_100.kgn         // The built searcher, which is automatically generated during run time. (In the corresponding dataset configuration file, save_types is set to save_searcher.)
├── datasets                                                // Store the dataset.
      └── sift-128-euclidean.hdf5
├── main.py                                                 // The file that contains the running functions.
└── sift_99.json                                            // The corresponding dataset configuration file.

Procedure:

Assume that the program runs in the /path/to/kbest_test directory. Check whether the datasets/sift-128-euclidean.hdf5, main.py and sift_99.json files exist in the directory. main.py and sift_99.json are provided at the end of this section.
Ensure that the num_numa_nodes in the sift_99.json file is set to the actual number of NUMA nodes during run time.

Install related dependencies.

pip install scikit-learn h5py psutil numpy==1.24.2

Run main.py.
1
python main.py 80 -1 sift_99.json
The test command parameters are described as follows:
```
python main.py <threads> <batch_size> <json_name>
```
- threads indicates the number of running threads.
- batch_size indicates the number of queries to be executed at a time in batch query mode. If batch_size is set to -1, all queries in the dataset are executed at a time.
- json_name indicates the name of the configuration file corresponding to the test dataset.
The command output is as follows:

The content of main.py is as follows:

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:                                                                                                          # Dataset base class.
    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):                                                                                           # Subclass of the sift-128-euclidean.hdf5 dataset used in the test case.
    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,               # Initialization.
                               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)                                                            # Build a graph index.
            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()                                                                              # Build a searcher.
                print(f"[INFO] Done build searcher, now save to {index_save_path}")
                self.index.save(index_save_path)                                                                        #  Save the searcher.
            elif save_types == "save_graph":
                self.index.saveGraph(index_save_path)                                                                   # Save the graph index.
            else:
                self.index.buildSearcher()
                print("[INFO] serialize.")
                data_arr = self.index.serialize()                                                                       # Serialization.

        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)                                                              # Load the searcher.
        elif save_types == "save_graph":
            self.index = KBest(self.numa_enabled, self.num_numa_nodes)
            load_result = self.index.loadGraph(index_load_path)                                                         # Load the graph index.
        else:
            print("Index deserialize.")
            self.index = KBest(self.numa_enabled, self.num_numa_nodes)
            load_result = self.index.deserialize(data_arr)                                                              # Deserialization.
        if load_result == -1:
            exit(-1)

    def set_query_arguments(self, ef):
        self.index.setEf(ef)                                                                                            # Set the size of the candidate node list for search.
        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)                             # Search.

    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()

The content of sift_99.json is as follows:

{
  "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"
}

Parent topic: Python