Faiss原始碼剖析：類結構分析

摘要：在下文中，我將嘗試透過Faiss原始碼中各種類結構的設計來梳理Faiss中的各種概念以及它們之間的關係。

本文分享自華為雲社群《》，原文作者：HW007。

Faiss是由Facebook AI Research研發的為稠密向量提供高效相似度搜尋和聚類的框架。透過其官方給出的新手指南，我們可以快速地體驗Faiss的基本功能。但是，相信大多數人看完官方的新手指南後，對Faiss很多的概念還是有點模糊、無法清晰的明確這些概念之間的邊界。比如說在Faiss中，Quantizer是個什麼概念、其與Index之間的聯絡是什麼；還有各種Index之間的關係又是什麼等等。為此，在下文中，我將嘗試透過Faiss原始碼中各種類結構的設計來梳理Faiss中的各種概念以及它們之間的關係。

首先奉上Faiss原始碼的類圖全家福如下，詳細的EA類圖檔案見附件：

圖一：Faiss的類圖全家福

首先，我們來看一下Faiss最主要的功能：相似度搜尋。如下圖所示，以圖片搜尋為例，所謂相似度搜尋，便是在給定的一堆圖片（下圖中左上角的圖集）中，尋找出我指定的目標（下圖中左下角的巴士圖片）最像的K張圖片，也簡稱為KNN（K近鄰）問題。

接下來我們看一下為了解決KNN問題，在工程上我們至少需要做哪些事情。顯然，有兩件事是必須要做的，第一，我們要把上面例子中的那個相簿儲存起來；第二，當使用者指定一種圖片後，我們需要知道怎麼從儲存的相簿中找到最近相似的K張圖片。由此，我們確定了Faiss在其應用場景中至少應該具備的兩個功能：新增功能和搜尋功能。

對於熟悉資料庫的同學來說，應該能在這裡嗅到點“CRUD”的味道。的確，當我們對“圖集”有新增儲存這樣的動作後，修改和刪除等功能也便接踵而來了。由此Faiss本質上就是一個向量資料庫。對於資料庫來說，時空最佳化是兩個永恆的主題，即在儲存上如何以更少的空間來儲存更多的資訊，在搜尋上如何以更快的速度來搜尋出更準確的資訊。如何減少搜尋所需的時間？在資料庫中很最常見的操作便是加各種索引，把各種加速搜尋演算法的功能或空間換時間的策略都封裝成各種各樣的索引，以滿足各種不同的引用場景。

由此，我們便不難理解為什麼Faiss中為什麼會有那麼多的Index了，因為Index這個概念本身就與加速搜尋是綁在一起的。由此也可以看出在Faiss中，如何又快又準地找到相似向量是第一要務。下圖中給出的是Faiss中最重要的兩個基類：Index和IndexBinary。

在上圖中，用白色的箭頭標出了這兩個基類中最重要的三個函式，其中add()和search() 函式便對應了我上文中所提到的Faiss至少應該實現的兩個基本功能：儲存和搜尋。在此順帶提一下，與傳統的資料庫相比，Faiss的Index還包含了資料儲存的功能，如果你一開始就從字面上按照傳統資料庫中索引的概念來理解地話，就會感覺有點怪怪的。接下來，我們重點聊聊Index中的train()函式，我們都知道天上是不會白白掉餡餅的，對於Faiss來說，不管其為了減少儲存空間還是加速搜尋，都需要提前做好一些準備工作，這便是train()函式發揮作用的時候了。

以減少儲存為例子，我們都知道在圖片處理中透過PCA可以將圖片從高維空間（p維）轉換到低維空間（q維, 其中 p > q ），其具體操作便是是將高維空間中的圖片向量（n*p）乘以一個轉換矩陣（p*q），得到一個低維空間中的向量（n*q）。為了使得在整個降維的過程中資訊丟失最少，我們需要對待轉換圖片進行分析計算得到相應的轉換矩陣（p*q）。也就是說這個降維中乘以的轉換矩陣是與待轉換圖片息息相關的。

回到我們的Faiss中來，假設我期望使用PCA預處理來減少Index中的儲存空間，那在整個處理流程中，除了輸入搜尋相簿外，我必須多輸入一個轉換矩陣，但是這個轉換矩陣是與相簿息息相關的，是可以由相簿資料計算出來的。如果把這個轉換矩陣看成一個引數的話，我們可以發現，在Faiss的一些預處理中，我們會引入一些引數，這些引數又無法一開始由人工來指定，只能透過喂樣本來訓練出來，所以Index中需要有這樣的一個train() 函式來為這種引數的訓練提供輸入訓練樣本的介面。由此，我們也可以發現，這些餵給train()函式的樣本資料最好與之後要新增儲存的圖集以及搜尋目標一致比較好，比如說，你先給Index喂一個豬臉資料集訓練出PCA中的轉換矩陣，再給這個Index新增人臉資料集，最後再在這個索引上做人臉識別，這樣肯定比不上一開始就喂人臉資料集得到PCA轉換矩陣的效果好。

由上，我們已經可以從train()、add()和search()三大函式大概地瞭解到Faiss中的Index是個什麼東西了，接下來我們看一下Faiss中有哪些不同的Index。從圖一中的類圖中可以看到，在Faiss中，大多數類基本都繼承或使用了Index介面，他們要麼對Index介面中定義的train、add和search函式進行了自己個性化的實現（如圖一中被淡橙色標註的類），要麼就是對已經實現的三大函式的類進行包裝，提供一些三大函式之外的流程上的加工處理（如圖一中被淡藍色標註的類）。

從圖一中我們可以看到這些被淡藍色標註的偏包裝的Index子類，他們與Index基類之間既有“is a”又有“hold a”關係，在類結構上出現這種關係的時候，設計者要麼是在設計一個樹或連結串列的節點，要麼是在設計一個包裝類。顯然在Faiss中更偏向於後者。一方面，淡藍色的Index子類藉助其所“hold”的Index來提供基本的train、add和search功能，使其自身符合Index介面的定義標準，成為一種Index，為之後的層層巢狀包裝提供支援。另一方面，他又對其所“hold”的Index類進行了一些通用的功能擴充套件。如下圖的IndexPreTransform類所示，Faiss將對待儲存圖集的預處理，如歸一化、PCA降維等功能抽象成一個VectorTransform介面，讓IndexPreTransform使用它來為其所“hold”的Index新增預處理功能，這種預處理功能是與其所“hold”的是什麼Index沒有任何關係，因此我更偏向於將這種功能歸結為Index之外的流程上的包裝功能。如IndexPreTransform類提供了資料預處理功能、IndexIDMap類提供了自定義ID功能、IndexShards類為Index的平行計算提供了相關的支援等。

接下來我們來看一下圖一中被淡橙色標註的Index子類，如IndexLSH、IndexPQ、IndexIVFPQ等，從名字中我們可以大概瞭解到這些類都是基於一些不同的演算法實現的不同索引，他們的train、add和search方法各有差異。但在整體上還是能找到一些其他結構上的共性。在上文中，我們知道Index具有儲存的功能，這些被淡橙色標註的Index子類在資料儲存方式上基本可以劃分為兩大類，一類是統一存到一個容器中，如在IndexLSH、IndexPQ等中我們都可以看到一個命名為codes的vector容器。另一類是分桶儲存到多個容器中，這主要為索引後續的非精確分桶區域性搜尋提供支援，為此，Faiss特地抽象出InvertedLists介面，需要支援分桶區域性搜尋的Index子類均會有hold一個實現了InvertedLists介面（淡紫色標註）的例項來儲存其資料。如下圖所示，Faiss為InvertedLists介面提供了陣列、連結串列和磁碟檔案等三種不同的實現。

在圖一中還有兩個被標記為淡綠色的類ProductQuantizer和ScalarQuantizer值得大家關注下，在結構上，這兩個類均沒有派生的子類，並且所有其他的類與他們的關係均為“hold a”關係，很純粹的工具類。從其命名中的Quantizer（量化器）字尾可知，這兩個工具類的作用是將“連續或稠密”的資料進行“離散或稀疏化”，簡單來說就是進行聚類的操作，就像我們把18歲以下的稱為少年，18~50歲的稱為中年一樣，我們把具體年齡量化成年齡段的過程就是一個聚類的過程。從圖一中還可以看到，帶有Quantizer字尾的類還有四個：MultiIndexQuantizer、MultiIndexQuantizer2、IndexScalarQuantizer和Level1Quantizer。其中前三個均是透過對ProductQuantizer或ScalarQuantizer的包裝來實現Quantizer的功能，沒什麼稀奇的地方，但最後一個Level1Quantizer類竟然是包裝了兩個Index類，而且其中一個Index類的屬性名還是quantizer，如下圖所示。

難道Index也是一種Quantizer？的確，對於Index來說，我們更熟悉的是其將資料集儲存起來，再尋找某個資料在該資料集中的K個最近鄰點的功能。但如果Index中儲存的是資料分類後各個類的中心點呢，那麼對於某個資料，我們便可以在該Index上透過KNN來求得其K（此時K=1）個最近鄰點，這些求出來的中心點所代表的類便是該資料在聚類中該歸屬的類。由此我們可以看到Index是可用來聚類，將資料量化成類的中心點的。因此，Index可以被包裝成一個Quantizer也便不足為奇了。其實Index的這種聚類功能在Faiss的設計中是很常見的，除了上面所說的用來做Quantizer外，還可以用來輔助實現K-means演算法，這也是為什麼Level1Quantizer類中除quantizer外還存在一個名為clustering_index的Index型別屬性的原因。透過上面的分析，我們還可以知道，在Faiss的Quantizer類中，或明或暗都應該有個地方來儲存用來輔助量化的“centroids”，即類中心點，它們在大多數場景中都是經過資料訓練出來的（如對資料進行K-means聚類），在少數場景中也可以直接人為設定。

讓我們最後來關注下IndexIVF類（上圖中被圈出來的淡紫色類）。也許在上文介紹淡紫色的InvertedLists類簇時，有人會有疑問，InvertedLists類及其派生子類在Faiss中主要為Index提供非精確的分桶區域性搜尋功能，這種功能與Index的種類毫無關係，按上文對Index派生的子類的分類標準來看，IndexIVF類應該是一個偏包裝的Index子類，應該被標註為淡藍色才對。的確，如上圖所示，雖然IndexIVF類沒有直接“hold a”Index類，但其透過繼承Level1Quantizer類間接“hold a”Index類，確實也是一個偏包裝的Index派生子類。圖一的顏色標註只是為了突出擁有IVF功能的Index類，透過顏色來輔助各個功能類簇在視覺上的區分度而已，不必深究。

透過上文，我們可以發現，Faiss的整個類結構設計是非常清晰簡潔的，其首先將KNN問題的解決過程切分成train、add和search三個步驟並抽象出Index基類。接著從這些基類派生出各種偏功能實現或者偏流程包裝的Index子類。此外還為Index提供了兩種的儲存方式：集中和分桶（IVF）。最後還提供了SQ和PQ兩種量化編碼工具以及將這些編碼工具或其他的Index包裝成Quantizer的類。