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.