faiss的簡單使用

簡介

faiss是爲稠密向量提供高效類似度搜索和聚類的框架。由Facebook AI Research研發。 具備如下特性。

• 一、提供多種檢索方法
• 二、速度快
• 三、可存在內存和磁盤中
• 四、C++實現，提供Python封裝調用。
• 五、大部分算法支持GPU實現

下面給出一些快速連接方便查找更多內容。

github
官方文檔
c++類信息
Troubleshooting
官方安裝文檔

安裝

文檔中給出來編譯安裝，conda等安裝方式。由於公司服務器編譯安裝須要權限，全部咱們通常使用conda的方式安裝python Module。

# 更新conda
conda update conda
# 先安裝mkl
conda install mkl
# faiss提供gpu和cpu版，根據服務選擇
conda install faiss-cpu -c pytorch # cpu
conda install faiss-gpu -c pytorch # gpu
# 校驗是否安裝成功
python -c "import faiss"

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

Quick Start

這裏先給出官方提供的demo來感覺一下faiss的使用。

首先構建訓練數據和測試數據

import numpy as np
d = 64                           # dimension
nb = 100000                      # database size
nq = 10000                       # nb of queries
np.random.seed(1234)             # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9

上面咱們構建了shape爲[100000,64]的訓練數據xb，和shape爲[10000,64]的查詢數據xq。
而後建立索引(Index)。faiss建立索引對向量預處理，提升查詢效率。
faiss提供多種索引方法，這裏選擇最簡單的暴力檢索L2距離的索引：IndexFlatL2。
建立索引時必須指定向量的維度d。大部分索引須要訓練的步驟。IndexFlatL2跳過這一步。
當索引建立好並訓練(若是須要)以後，咱們就能夠執行add和search方法了。add方法通常添加訓練時的樣本，search就是尋找類似類似向量了。
一些索引能夠保存整型的ID，每一個向量能夠指定一個ID，當查詢類似向量時，會返回類似向量的ID及類似度(或距離)。若是不指定，將按照添加的順序從0開始累加。其中IndexFlatL2不支持指定ID。

import faiss                   # make faiss available
index = faiss.IndexFlatL2(d)   # build the index
print(index.is_trained)
index.add(xb)                  # add vectors to the index
print(index.ntotal)

 

• 1
• 2
• 3
• 4
• 5

咱們有了包含向量的索引後，就能夠傳入搜索向量查找類似向量了。

k = 4                          # we want to see 4 nearest neighbors
D, I = index.search(xq, k)     # actual search
print(I[:5])                   # neighbors of the 5 first queries
print(D[-5:])                  # neighbors of the 5 last queries

 

• 1
• 2
• 3
• 4

上面代碼中，咱們定義返回每一個須要查詢向量的最近4個向量。查詢返回兩個numpy array對象D和I。D表示與類似向量的距離(distance)，維度，I表示類似用戶的ID。

咱們能夠獲得相似於下面的結果

[[  0 393 363  78]
 [  1 555 277 364]
 [  2 304 101  13]
 [  3 173  18 182]
 [  4 288 370 531]]

[[ 0.          7.17517328  7.2076292   7.25116253]
 [ 0.          6.32356453  6.6845808   6.79994535]
 [ 0.          5.79640865  6.39173603  7.28151226]
 [ 0.          7.27790546  7.52798653  7.66284657]
 [ 0.          6.76380348  7.29512024  7.36881447]]

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11

加速搜索

若是須要存儲的向量太多，經過暴力搜索索引IndexFlatL2速度很慢，這裏介紹一種加速搜索的方法的索引IndexIVFFlat。翻譯過來叫倒排文件，實際上是使用K-means創建聚類中心，而後經過查詢最近的聚類中心，而後比較聚類中的全部向量獲得類似的向量。

建立IndexIVFFlat時須要指定一個其餘的索引做爲量化器(quantizer)來計算距離或類似度。

這裏同使用IndexFlatL2對比，在add方法以前須要先訓練。

下面簡述示例中的幾個參數。

faiss.METRIC_L2: faiss定義了兩種衡量類似度的方法(metrics)，分別爲faiss.METRIC_L2、faiss.METRIC_INNER_PRODUCT。一個是歐式距離，一個是向量內積。

nlist：聚類中心的個數

k：查找最類似的k個向量

index.nprobe：查找聚類中心的個數，默認爲1個。

代碼示例以下

nlist = 100                       #聚類中心的個數
k = 4
quantizer = faiss.IndexFlatL2(d)  # the other index
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
       # here we specify METRIC_L2, by default it performs inner-product search
assert not index.is_trained
index.train(xb)
assert index.is_trained

index.add(xb)                  # add may be a bit slower as well
D, I = index.search(xq, k)     # actual search
print(I[-5:])                  # neighbors of the 5 last queries
index.nprobe = 10              # default nprobe is 1, try a few more
D, I = index.search(xq, k)
print(I[-5:])                  # neighbors of the 5 last queries

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• 12
• 13
• 14
• 15

減小內存

2018-02-22以後版本添加了磁盤存儲inverted indexes的方式，使用可參考demo.

上面咱們看到的索引IndexFlatL2和IndexIVFFlat都會全量存儲全部的向量在內存中，爲知足大的數據量的需求，faiss提供一種基於Product Quantizer(乘積量化)的壓縮算法編碼向量大小到指定的字節數。此時，存儲的向量時壓縮過的，查詢的距離也是近似的。關於乘積量化的算法可自行搜索。

下面給出demo。相似IndexIVFFlat，這裏使用的是IndexIVFPQ

nlist = 100
m = 8                             # number of bytes per vector
k = 4
quantizer = faiss.IndexFlatL2(d)  # this remains the same
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8)
                                    # 8 specifies that each sub-vector is encoded as 8 bits
index.train(xb)
index.add(xb)
D, I = index.search(xb[:5], k) # sanity check
print(I)
print(D)
index.nprobe = 10              # make comparable with experiment above
D, I = index.search(xq, k)     # search
print(I[-5:])

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11
• 12
• 13
• 14

以前咱們定義的維度爲d = 64，向量的數據類型爲float32。這裏壓縮成了8個字節。因此壓縮比率爲 (64*32/8) / 8 = 32

返回的結果以下，第一個向量同本身的距離爲1.40704751，不是0。由於如上所述返回的是近似距離，可是總體上返回的最類似的top k的向量ID沒有變化。

[[   0  608  220  228]
 [   1 1063  277  617]
 [   2   46  114  304]
 [   3  791  527  316]
 [   4  159  288  393]]

[[ 1.40704751  6.19361687  6.34912491  6.35771513]
 [ 1.49901485  5.66632462  5.94188499  6.29570007]
 [ 1.63260388  6.04126883  6.18447495  6.26815748]
 [ 1.5356375   6.33165455  6.64519501  6.86594009]
 [ 1.46203303  6.5022912   6.62621975  6.63154221]]

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10
• 11

簡化索引的表達

經過上面IndexIVFFlat和IndexIVFPQ咱們能夠看到，他們的構造須要先提供另一個index。相似的，faiss還提供pca、lsh等方法，有時候他們會組合使用。這樣組合的對構造索引會比較麻煩，faiss提供了經過字符串表達的方式構造索引。
如，下面表達式就能表示上面的建立IndexIVFPQ的實例。

index = faiss.index_factory(d, "IVF100,PQ8")

 

• 1

這裏有一點文檔中沒有提到的，經過查看c++代碼，index_factory方法還有第三個參數，就是上面說的metric。可傳入的就上面兩種。

Index *index_factory (int d, const char *description_in, MetricType metric)

 

• 1

更多的組合實例能夠看demo

每類索引的簡寫可查詢Basic indexes

GPU使用

注意有些索引不支持GPU，哪些支持哪些不支持可查詢Basic indexes

可經過faiss.get_num_gpus()查詢有多少個gpu

ngpus = faiss.get_num_gpus()
print("number of GPUs:", ngpus)

 

• 1
• 2

使用gpu的完整示例。

一、使用一塊gpu

# build a flat (CPU) index
index_flat = faiss.IndexFlatL2(d)
# make it into a gpu index
gpu_index_flat = faiss.index_cpu_to_gpu(res, 0, index_flat)

 

• 1
• 2
• 3
• 4

二、使用所有gpu

cpu_index = faiss.IndexFlatL2(d)
gpu_index = faiss.index_cpu_to_all_gpus(cpu_index) # build the index

gpu_index.add(xb)              # add vectors to the index
print(gpu_index.ntotal)

k = 4                          # we want to see 4 nearest neighbors
D, I = gpu_index.search(xq, k) # actual search
print(I[:5])                   # neighbors of the 5 first queries
print(I[-5:])                  # neighbors of the 5 last queries

 

• 1
• 2
• 3
• 4
• 5
• 6
• 7
• 8
• 9
• 10