添加自研算法KScaNN

ann-benchmarks中自带了很多算法，也可以使用华为自研算法KScaNN对数据集进行搜索。以下为添加KScaNN算法的具体步骤。

添加KScaNN算法实现。

打开module.py文件。

1	vim /data/ann-benchmarks-main/ann_benchmarks/algorithms/milvus/module.py

按“i”进入编辑模式，在文件末尾添加下面代码。

class MilvusKSCANN(Milvus):
    def __init__(self, metric, dim, index_param):
        super().__init__(metric, dim, index_param)
        self._index_n_leaves = index_param.get("n_leaves", None)
        self._index_dims_per_block = index_param.get("dims_per_block", None)
        self._index_avq_threshold = index_param.get("avq_threshold", None)
        self._index_soar_lambda = index_param.get("soar_lambda", None)
        self._index_overretrieve_factor = index_param.get("overretrieve_factor", None)
        self._index_train_thread = index_param.get("train_thread", None)

    def get_index_param(self):
        return {
            "index_type": "KSCANN",
            "params": {
                "n_leaves": self._index_n_leaves,
                "dims_per_block": self._index_dims_per_block,
                "avq_threshold": self._index_avq_threshold,
                "soar_lambda": self._index_soar_lambda,
                "overretrieve_factor": self._index_overretrieve_factor,
                "train_thread": self._index_train_thread
            },
            "metric_type": self._metric_type
        }


    def set_query_arguments(self, query_args):
        nprobe, reorder, adp_threshold, adp_refined, num_thread, batch_size = query_args
        self.search_params = {
            "metric_type": self._metric_type,
            "params": {
                "nprobe": nprobe,
                "reorder": reorder,
                "adp_threshold": adp_threshold,
                "adp_refined": adp_refined,
                "num_thread": num_thread,
                "batch_size": batch_size,
                "k": 10
            }
        }
        self.name = f"MilvusKScaNN metric:{self._metric}, index_n_leaves:{self._index_n_leaves}, index_dims_per_block:{self._index_dims_per_block}, index_avq_threshold:{self._index_avq_threshold}, index_soar_lambda:{self._index_soar_lambda}, index_overretrieve_factor:{self._index_overretrieve_factor}, index_train_thread:{self._index_train_thread}, search_nprobe:{nprobe}, search_reorder:{reorder}, search_adp_threshold={adp_threshold}, search_adp_refined={adp_refined}, search_num_thread={num_thread}, search_batch_size={batch_size}"

按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

添加KScaNN算法配置。

打开config.yml文件。

1	vim ann_benchmarks/algorithms/milvus/config.yml

按“i”进入编辑模式，在文件末尾添加下面代码。

- base_args: ["@metric", "@dimension"]
  constructor: MilvusKSCANN
  disabled: false
  docker_tag: ann-benchmarks-milvus
  module: ann_benchmarks.algorithms.milvus
  name: milvus-kscann
  run_groups:
    KScaNN:
      args:
        n_leaves: [2000]
        dims_per_block: [4]
        avq_threshold:
        soar_lambda: [-1]
        overretrieve_factor: [-1]
        train_thread: [16]
      query_args:
        # [
        #   [210],  # nprobe
        #   [900],  # reorder
        #   [0.2],  # adp_threshold
        #   [0],  # adp_refined
        #   [1], # num_thread
        #   [1] # batch_size
        # ]
        [[[10, 100, 0, 0, 1, 1], [15, 140, 0, 0, 1, 1], [25, 160, 0, 0, 1, 1], [35, 190, 0, 0, 1, 1], [40, 200, 0, 0, 1, 1], [45, 220, 0, 0, 1, 1], [50, 240, 0, 0, 1, 1], [60, 250, 0, 0, 1, 1], [70, 300, 0, 0, 1, 1], [80, 400, 0, 0, 1, 1], [100, 500, 0, 0, 1, 1], [120, 600, 0, 0, 1, 1], [150, 800, 0, 0, 1, 1], [200, 900, 0, 0, 1, 1], [250, 900, 0, 0, 1, 1]]]

按“Esc”键，输入:wq!，按“Enter”保存并退出编辑。

KScaNN参数说明见表1。配置参考是兼顾查询结果准确性、内存消耗、时间消耗的综合考虑，测试时可根据实际情况自行选取。

表1 **KScaNN参数**说明
参数名称	参数说明	类型范围	配置参考	配置原则
n_leaves	叶子节点数量	int	[2000]	该参数影响图构建耗时和最终索引质量，过大可能会导致构建耗时过长以及搜索性能下降，过小则会影响检索精度。
dims_per_block	构图时的PQ阶段，多少维构成一个子向量块	int	[4]	一般推荐使用4，根据实际情况自行调整。
avq_threshold	构图时的avq阈值	float	None	该参数影响剪枝策略。对于IP数据集，一般使用0.2；L2数据集一般设置为空。
soar_lambda	控制正交性的参数，仅对IP数据集生效。	float，大于0	[-1]	-1表示不使用，IP数据集根据实际情况自行调整范围。
overretrieve_factor	指定过检索因子，与soar_lambda一起使用，仅对IP数据集生效。	float，[1,2]	[-1]	-1表示不使用，IP数据集根据实际情况自行调整范围。
train_thread	构图时的线程数。	int	[CPU核数]	无特殊情况，设置成CPU核数大小。
nprobe	查询时的复杂query所采用的子空间个数。	int，[1,n_leaves]	[200]	根据实际情况自行调整。
reorder	查询时在重排前保存结果的数量。	int，[k, 底库数量]	[900]	k是最终返回结果数量，reorder需根据实际情况自行调整。
adp_threshold	查询时的智能截断阈值。暂未使用。	float，[0,0.8]	[0]	根据实际情况自行调整。
adp_refined	查询时简单query所采用的子空间个数。暂未使用。	int，[0,nprobe]	[0]	典型值为0，但是搜推设置范围为[1,nprobe]，这里设置为[0,probe]。根据实际情况自行调整。
num_thread	查询时的线程数。	int，大于等于1	[1]	无特殊情况设置为1。
batch_size	并行自动分批时优先批次大小。	int，大于等于1	[1]	无特殊情况设置为1。

父主题： 修改算法参数