Faiss学习笔记

Faiss

Faiss是Facebook的AI Research开发的Similarity Search工具包.

具体地就是对稠密向量进行快速的相似度检索. 这里很容易就能想到推荐系统中的召回过程. 召回就是将所有的物品向量存储在一个数据库中, 然后使用用户向量去检索其中欧氏距离最近的几个向量作为召回的物品列表.

一. Faiss的使用步骤

1. 构建embedding矩阵, 具体地类似于加载预训练embedding, 其排序是从0排序到N-1的. 维度就是(N, v_dim)
2. 选择合适的index, 并将embedding添加到index中
3. 给定index一个query和数量k, 得到按照相似度排序的列表

二. 具体使用方法

import faiss

embedding_dim = 8
index = faiss.IndexFlatIP(embedding_dim)	# 创建index对象

index.add(item_embedding)	# 给index对象添加embedding

D, I = index.search(user_embedding, 10)	# 给定一个或多个向量进行搜索

1. IndexFlatIP

根据相似度的计算方法和内存准确度的tradeoff来选择合适的index对象

embedding_dim就是要创建的index中存储的embedding的维度大小

2. add

使用add方法给index对象加入被搜索的embedding列表, 一般添加的是numpy.

item_embeddings = item_embedding.weight.cpu().detach().numpy().astype('float32')

注意之所以要astype, 因为faiss使用的数据类型默认的使用float32.

TypeError: in method 'IndexFlat_add', argument 3 of type 'float const *'

如果数据类型没有切换对, 就会报上面的这个错.

3. search

输入待搜索的向量, 这个向量可以是一个, 也可以是多个, 所以其维度为: shape(n, embedding_dim)

⭐️ n就表示了可以一次性的进行批量搜索.

得到结果的维度是: shape(n, k).

具体地, D为搜索到的向量矩阵, 而I为搜索到的索引矩阵.

Reference

参考资料1

参考资料2