Python如何利用Faiss库实现ANN近邻搜索

发布时间:2020-08-03 09:39:29 作者:小猪
来源:亿速云 阅读:296

这篇文章主要讲解了Python如何利用Faiss库实现ANN近邻搜索,内容清晰明了,对此有兴趣的小伙伴可以学习一下,相信大家阅读完之后会有帮助。

Embedding的近邻搜索是当前图推荐系统非常重要的一种召回方式,通过item2vec、矩阵分解、双塔DNN等方式都能够产出训练好的user embedding、item embedding,对于embedding的使用非常的灵活:

然而有一个工程问题,一旦user embedding、item embedding数据量达到一定的程度,对他们的近邻搜索将会变得非常慢,如果离线阶段提前搜索好在高速缓存比如redis存储好结果当然没问题,但是这种方式很不实时,如果能在线阶段上线几十MS的搜索当然效果最好。

Faiss是Facebook AI团队开源的针对聚类和相似性搜索库,为稠密向量提供高效相似度搜索和聚类,支持十亿级别向量的搜索,是目前最为成熟的近似近邻搜索库。

接下来通过jupyter notebook的代码,给大家演示下使用faiss的简单流程,内容包括:

对于已经训练好的Embedding怎样实现高速近邻搜索是一个工程问题,facebook的faiss库可以构建多种embedding索引实现目标embedding的高速近邻搜索,能够满足在线使用的需要

安装命令:

conda install -c pytorch faiss-cpu 

提前总结下faiss使用经验:

1. 为了支持自己的ID,可以用faiss.IndexIDMap包裹faiss.IndexFlatL2即可

2. embedding数据都需要转换成np.float32,包括索引中的embedding以及待搜索的embedding

3. ids需要转换成int64类型

1. 准备数据

import pandas as pd
import numpy as np
df = pd.read_csv("./datas/movielens_sparkals_item_embedding.csv")
df.head()
idfeatures
010[0.25866490602493286, 0.3560594320297241, 0.15…
120[0.12449632585048676, -0.29282501339912415, -0…
230[0.9557555317878723, 0.6764761805534363, 0.114…
340[0.3184879720211029, 0.6365472078323364, 0.596…
450[0.45523127913475037, 0.34402626752853394, -0….

构建ids

ids = df["id"].values.astype(np.int64)
type(ids), ids.shape
(numpy.ndarray, (3706,))
ids.dtype
dtype('int64')
ids_size = ids.shape[0]
ids_size
3706

构建datas

import json
import numpy as np
datas = []
for x in df["features"]:
 datas.append(json.loads(x))
datas = np.array(datas).astype(np.float32)
datas.dtype
dtype('float32')
datas.shape
(3706, 10)
datas[0]
array([ 0.2586649 , 0.35605943, 0.15589039, -0.7067125 , -0.07414215,
 -0.62500805, -0.0573845 , 0.4533663 , 0.26074877, -0.60799956],
 dtype=float32)
# 维度
dimension = datas.shape[1]
dimension
10

2. 建立索引

import faiss
index = faiss.IndexFlatL2(dimension)
index2 = faiss.IndexIDMap(index)
ids.dtype
dtype('int64')
index2.add_with_ids(datas, ids)
index.ntotal
3706

4. 搜索近邻ID列表

df_user = pd.read_csv("./datas/movielens_sparkals_user_embedding.csv")
df_user.head()
id features
idfeatures
010[0.5974288582801819, 0.17486965656280518, 0.04…
120[1.3099910020828247, 0.5037978291511536, 0.260…
230[-1.1886241436004639, -0.13511677086353302, 0….
340[1.0809299945831299, 1.0048035383224487, 0.986…
450[0.42388680577278137, 0.5294889807701111, -0.6…
user_embedding = np.array(json.loads(df_user[df_user["id"] == 10]["features"].iloc[0]))
user_embedding = np.expand_dims(user_embedding, axis=0).astype(np.float32)
user_embedding
array([[ 0.59742886, 0.17486966, 0.04345559, -1.3193961 , 0.5313592 ,
 -0.6052168 , -0.19088413, 1.5307966 , 0.09310367, -2.7573566 ]],
 dtype=float32)
user_embedding.shape
(1, 10)
user_embedding.dtype
dtype('float32')
topk = 30
D, I = index.search(user_embedding, topk) # actual search
I.shape
(1, 30)
I
array([[3380, 2900, 1953, 121, 3285, 999, 617, 747, 2351, 601, 2347,
 42, 2383, 538, 1774, 980, 2165, 3049, 2664, 367, 3289, 2866,
 2452, 547, 1072, 2055, 3660, 3343, 3390, 3590]])

5. 根据电影ID取出电影信息

target_ids = pd.Series(I[0], name="MovieID")
target_ids.head()
0 3380
1 2900
2 1953
3 121
4 3285
Name: MovieID, dtype: int64
df_movie = pd.read_csv("./datas/ml-1m/movies.dat",
  sep="::", header=None, engine="python",
  names = "MovieID::Title::Genres".split("::"))
df_movie.head()
MovieIDTitleGenres
01Toy Story (1995)Animation|Children's|Comedy
12Jumanji (1995)Adventure|Children's|Fantasy
23Grumpier Old Men (1995)Comedy|Romance
34Waiting to Exhale (1995)Comedy|Drama
45Father of the Bride Part II (1995)Comedy
df_result = pd.merge(target_ids, df_movie)
df_result.head()
MovieIDTitleGenres
03380Railroaded! (1947)Film-Noir
12900Monkey Shines (1988)Horror|Sci-Fi
21953French Connection, The (1971)Action|Crime|Drama|Thriller
3121Boys of St. Vincent, The (1993)Drama
43285Beach, The (2000)Adventure|Drama

看完上述内容,是不是对Python如何利用Faiss库实现ANN近邻搜索有进一步的了解,如果还想学习更多内容,欢迎关注亿速云行业资讯频道。

推荐阅读:
  1. 如何在python中实现K近邻算法
  2. python实现K最近邻算法

免责声明:本站发布的内容(图片、视频和文字)以原创、转载和分享为主,文章观点不代表本网站立场,如果涉及侵权请联系站长邮箱:is@yisu.com进行举报,并提供相关证据,一经查实,将立刻删除涉嫌侵权内容。

python faiss

上一篇:如何使用es6数组的flat(),flatMap()函数

下一篇:找不到php.ini怎么办

相关阅读

您好,登录后才能下订单哦!

密码登录
登录注册
其他方式登录
点击 登录注册 即表示同意《亿速云用户服务条款》