# KNN ## 背景 GNN的很大一部分应用是把产出的embedding结果用于向量召回,即把生成的embedding作为检索库,从中查找给定输入向量的TopK个匹配程度最高的向量,也就是K最近邻检索。 另外一方面,KNN也是构图的一种手段。在已知向量的前提下,通过计算向量间的相似程度,相似程度较高的两个向量之间被认为有隐含的某种联系,即建立一条Graph的边,再通过GNN算法去挖掘潜在意义。 出于这些考虑,GL支持大规模分布式离线向量检索(KNN),作为GNN生态的一部分。GL提供KNN相关的API,返回结果为numpy格式,可与GL的其他功能结合,也可与TF训练结合,给开发者一个完整的体验。 ## 用法 在GL的语义里,一切数据皆为图的元素,KNN的数据源也不例外。在GL里,我们把Graph的Node作为KNN数据的来源,即支持把Node数据构建成KNN检索库,Node需要有float类型的属性,这些属性将被视作向量。由于GL天然支持异构图,KNN检索库也可存在多个,不同顶点对应的KNN检索库是不同的。 例如,Graph具有item类型的顶点,每个顶点都具有21个float属性,如果要建立KNN检索,则代码如下: ```python import graphlearn as gl import tensorflow as tf import numpy as np flags = tf.app.flags FLAGS = flags.FLAGS flags.DEFINE_integer("task_index", None, "Task index") flags.DEFINE_string("job_name", None, "worker or ps") flags.DEFINE_string("worker_hosts", "", "worker hosts") flags.DEFINE_string("tables", "", "odps table name") if __name__ == "__main__": worker_count = len(FLAGS.worker_hosts.split(',')) item_path = FLAGS.tables.split(',') gl.set_knn_metric(metric) # 0 is l2 distance(default), 1 is inner product. option = gl.IndexOption() option.name = "knn" option.index_type = "ivfflat" option.nlist = 5 option.nprobe = 2 g = gl.Graph() \ .node(item_path, "item", decoder=gl.Decoder(attr_types=["float"] * 21), option=option) \ .init(task_index=FLAGS.task_index, task_count=worker_count) inputs = np.array([[5.0] * 21, [31.0] * 21], dtype=np.float32) ids, distances = g.search("item", inputs, gl.KnnOption(k=3)) print(ids, distances) g.close() ``` 这个示例中,21~23行添加图的顶点数据和初始化,细节请参照图对象一章。其中对于 item类型的Node,增加了一个`option` 选项,该选项表示对Node进行的额外操作,这里我们传入一个 `gl.IndexOption()` 对象,表示对Node建立索引,由16~21行所示,索引名称为 `knn`(必须为knn),建立knn的索引类型为 `ivfflat`。第27行是构造的输入数据,2个21维的float32向量,注意必须为float32,且维度与检索库的float属性个数一致。输入支持batch,即shape可以是[N, 21]的numpy array。 第29行是检索操作,`g.search()`的第一个参数为顶点类型,第二个参数为输入数据,第三个参数为KnnOption,这里只需指定k的值即可。返回结果为ids和distances两个二维向量,shape为[N, k],ids为与inputs向量距离最近的前k个id,int64类型,distances为对应的距离,float32类型。 关于index_type,除了`ivfflat`,还支持`flat`、`ivfpq`,GPU版本的还支持`gpu_ivfflat`、`gpu_flat`、`gpu_ivfpq`。其中,ivfflat为聚类后建立索引,平衡精度与搜索代价,最为常用;flat为严格搜索,保证正确但耗时;ivfpq会压缩精度,例如把32维压缩到8维,降低存储但损失精度。`option.nlist` 和`option.nprobe` 为聚类的簇个数和检索时的个数,例如共聚5类,检索时只考虑距离与输入向量最近的2类。这两个选项在index_type为flat时不起作用。当index_type为ivfpq时,`option.m` 参数指定压缩的维度,注意option.m必须为float属性个数的因子。 `gl.set_knn_metric` 用来设置knn距离计算方式,0是l2距离,1是内积距离,默认0. 使用KNN评估召回结果的示列见 examples/eval.