Faiss源碼剖析:類結構分析

摘要:在下文中,我將嘗試通過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的類。

 

點擊關注,第一時間瞭解華爲雲新鮮技術~

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章