Skip to content

[Python] 使用 Faiss 相似度搜索快速匹配向量

介紹

Faiss 是 Facebook AI Research 團隊所開發的一款高效向量匹配函式庫,只要輸入向量後,就能夠在大量的高維度向量中迅速匹配最相像的 top k 個向量。

Faiss 主要是 C++ 所開發的,但是也提供了 Python 接口,所以可以輕鬆地透過 Python 調用;另外,Faiss 也支援 GPU 平行化計算,更進一步地處理大量高維度向量資料。

本篇文章將介紹 Faiss 的安裝步驟及使用方法說明。


如何安裝 Faiss

Faiss 可以直接透過 pip 進行安裝。需要注意的是,有分成 CPU 及 GPU 版本。

pip install faiss-cpu


如果要安裝 GPU 版本,則需要事先安裝好 CUDA 及 cuDNN 的相關套件。


使用 Faiss 匹配向量

準備資料

首先,我們先來模擬一些高維度的向量;這些向量可能是我們知識庫中的資料,之後我們要做向量檢索時,就從這些知識庫中的向量進行匹配。

在以下程式碼中,N 代表著我們資料的筆數、D 代表著我們一筆資料向量的維度。然後我們隨機地產生 100,000 筆資料。

import numpy as np

N = 100000
D = 128
data = np.random.random((N, D)).astype(np.float32)


建立索引

在這裡,我們需要建立一個 Faiss 的索引物件,這個索引物件需要知道接下來要查詢向量的維度,也就是我們上面的參數 D

在這裡,我們使用 IndexFlatL2(),預設以歐式距離來計算。

import faiss

index = faiss.IndexFlatL2(D)


如果想要以餘弦相似度(cosine similarity)來計算,則需要使用資料的向量歸一化並轉成內積空間。可以參考這裡的討論


將資料加入索引

我們剛剛已經建立好了資料 data,這就是我們要儲存在知識庫中等待被查詢的向量資料。現在就要將其加入索引。

index.add(data)


進行查詢

假設我們有一組等待被查詢的向量,要找其最像的 k 個向量距離和其索引。

k = 10

query_vector = np.random.random((1, D)).astype(np.float32)
distances, indices = index.search(query_vector, k)


其中 distancesindices 都是形狀為 (1, k) 的 NumPy 陣列。


References


Read More

Tags:

Leave a Reply取消回覆

Click to Copy
Exit mobile version