Python利用Faiss庫實現ANN近鄰搜索的方法詳解
Embedding的近鄰搜索是當前圖推薦系統非常重要的一種召回方式，通過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數據量達到一定的程度，對他們的近鄰搜索將會變得非常慢，如果離線階段提前搜索好在高速緩存比如redis存儲好結果當然沒問題，但是這種方式很不實時，如果能在線階段上線幾十MS的搜索當然效果最好。
Faiss是Facebook AI團隊開源的針對聚類和相似性搜索庫，為稠密向量提供高效相似度搜索和聚類，支持十億級別向量的搜索，是目前最為成熟的近似近鄰搜索庫。
接下來通過jupyter notebook的代碼，給大家演示下使用faiss的簡單流程，內容包括：
讀取訓練好的Embedding數據 構建faiss索引，將待搜索的Embedding添加進去 取得目標Embedding，實現搜索得到ID列表 根據ID獲取電影標題，返回結果對于已經訓練好的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 pdimport 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.dtypedtype(’int64’)ids_size = ids.shape[0]ids_size3706
構建datas
import jsonimport numpy as npdatas = []for x in df['features']: datas.append(json.loads(x))datas = np.array(datas).astype(np.float32)datas.dtypedtype(’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]dimension10
2. 建立索引
import faissindex = faiss.IndexFlatL2(dimension)index2 = faiss.IndexIDMap(index)ids.dtypedtype(’int64’)index2.add_with_ids(datas, ids)index.ntotal3706
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_embeddingarray([[ 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.dtypedtype(’float32’)topk = 30D, I = index.search(user_embedding, topk) # actual searchI.shape(1, 30)Iarray([[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 33801 29002 19533 1214 3285Name: MovieID, dtype: int64df_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
總結
到此這篇關于Python利用Faiss庫實現ANN近鄰搜索的文章就介紹到這了,更多相關Python用Faiss庫ANN近鄰搜索內容請搜索好吧啦網以前的文章或繼續瀏覽下面的相關文章希望大家以后多多支持好吧啦網！
相關文章：
網公網安備