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检索，则代码如下：

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.