Python利用Faiss庫實現(xiàn)ANN近鄰搜索的方法詳解

更新時間：2020年08月03日 09:00:02 作者：螞蟻學Python

這篇文章主要介紹了Python利用Faiss庫實現(xiàn)ANN近鄰搜索的方法，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面來一起學習學習吧

Embedding的近鄰搜索是當前圖推薦系統(tǒng)非常重要的一種召回方式，通過item2vec、矩陣分解、雙塔DNN等方式都能夠產出訓練好的user embedding、item embedding，對于embedding的使用非常的靈活：

輸入user embedding，近鄰搜索item embedding，可以給user推薦感興趣的items
輸入user embedding，近鄰搜搜user embedding，可以給user推薦感興趣的user
輸入item embedding，近鄰搜索item embedding，可以給item推薦相關的items

然而有一個工程問題，一旦user embedding、item embedding數(shù)據(jù)量達到一定的程度，對他們的近鄰搜索將會變得非常慢，如果離線階段提前搜索好在高速緩存比如redis存儲好結果當然沒問題，但是這種方式很不實時，如果能在線階段上線幾十MS的搜索當然效果最好。

Faiss是Facebook AI團隊開源的針對聚類和相似性搜索庫，為稠密向量提供高效相似度搜索和聚類，支持十億級別向量的搜索，是目前最為成熟的近似近鄰搜索庫。

接下來通過jupyter notebook的代碼，給大家演示下使用faiss的簡單流程，內容包括：

讀取訓練好的Embedding數(shù)據(jù)
構建faiss索引，將待搜索的Embedding添加進去
取得目標Embedding，實現(xiàn)搜索得到ID列表
根據(jù)ID獲取電影標題，返回結果

對于已經(jīng)訓練好的Embedding怎樣實現(xiàn)高速近鄰搜索是一個工程問題，facebook的faiss庫可以構建多種embedding索引實現(xiàn)目標embedding的高速近鄰搜索，能夠滿足在線使用的需要

安裝命令：

conda install -c pytorch faiss-cpu

提前總結下faiss使用經(jīng)驗：

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

2. embedding數(shù)據(jù)都需要轉換成np.float32，包括索引中的embedding以及待搜索的embedding

3. ids需要轉換成int64類型

1. 準備數(shù)據(jù)

import pandas as pd
import numpy as np

df = pd.read_csv("./datas/movielens_sparkals_item_embedding.csv")
df.head()

	id	features
0	10	[0.25866490602493286, 0.3560594320297241, 0.15…
1	20	[0.12449632585048676, -0.29282501339912415, -0…
2	30	[0.9557555317878723, 0.6764761805534363, 0.114…
3	40	[0.3184879720211029, 0.6365472078323364, 0.596…
4	50	[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

	id	features
0	10	[0.5974288582801819, 0.17486965656280518, 0.04…
1	20	[1.3099910020828247, 0.5037978291511536, 0.260…
2	30	[-1.1886241436004639, -0.13511677086353302, 0….
3	40	[1.0809299945831299, 1.0048035383224487, 0.986…
4	50	[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. 根據(jù)電影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()

	MovieID	Title	Genres
0	1	Toy Story (1995)	Animation\|Children's\|Comedy
1	2	Jumanji (1995)	Adventure\|Children's\|Fantasy
2	3	Grumpier Old Men (1995)	Comedy\|Romance
3	4	Waiting to Exhale (1995)	Comedy\|Drama
4	5	Father of the Bride Part II (1995)	Comedy

df_result = pd.merge(target_ids, df_movie)
df_result.head()

	MovieID	Title	Genres
0	3380	Railroaded! (1947)	Film-Noir
1	2900	Monkey Shines (1988)	Horror\|Sci-Fi
2	1953	French Connection, The (1971)	Action\|Crime\|Drama\|Thriller
3	121	Boys of St. Vincent, The (1993)	Drama
4	3285	Beach, The (2000)	Adventure\|Drama