選自arXiv.org機器之心編譯作者:吳攀
相似性搜索的規模和速度一直是研究者努力想要克服的難題。近日,Facebook 人工智慧研究團隊在 arXiv 發布的新論文《Billion-scale similarity search with GPUs》宣稱在這一問題上取得了重大進展,在 GPU 上實現了十億規模級的相似性搜索。該團隊已經將相關實現的代碼進行了開源。機器之心在此對該研究論文及其代碼項目進行了簡單介紹。
論文地址:https://arxiv.org/abs/1702.08734開源地址:https://github.com/facebookresearch/faiss
摘要
相似性搜索(similarity search)可以在用於圖像或視頻等複雜數據處理的專用資料庫系統中尋找應用,這些複雜數據通常是用高維特徵表示的,而且往往需要特定的索引結構(indexing structure)。本論文解決了讓這種任務更好地利用 GPU 的問題。儘管 GPU 擅長數據並行任務,但之前的方法會在並行性不高的算法(如 k-min selection)上遇到瓶頸或不能有效利用內存的層次結構。我們提出了一種用於 k-selection 的設計,其可以以高達理論峰值性能 55% 的速度進行運算,從而實現了比之前最佳的 GPU 方法快 8.5 倍的最近鄰搜索。我們通過為基於積量化(product quantization)的暴力計算、近似和壓縮域搜索(compressed-domain search)提出優化過的設計,從而將其應用到了不同的相似性搜索場景中。在所有這些場景中,我們的方法的表現都遠遠超越了之前的最佳表現。我們的實現可以在 35 分鐘內從 Yfcc100M 數據集的 9500 萬張圖像上構建一個高準確度的 k-NN 圖(graph),也可以在 12 個小時內在 4 個 Maxwell Titan X GPU 上構建一個連接了 10 億個向量的圖。我們已經開源了我們的方法,以便人們進行比較和重複利用。
WarpSelect 簡介
我們的 k-selection 實現名叫 WarpSelect,其完全在寄存器(register)中維持狀態,且僅需要在數據上進行單次通過,從而避免了 cross-warp synchronization。其使用了 merge-odd 和 sort-odd 作為原語(primitive)。因為寄存器文件能比共享內存提供遠遠更多的存儲,所以它支持 k ≤ 1024。每一個 warp 都專用於 k-selection 的 n 個數組 [ai] 中的單一一個。如果 n 足夠大,那麼每個 [ai] 一個 warp 會導致 GPU 被完全佔用。如果 l 是預先已知的,每個 warp 的大 l 會通過遞歸分解(recursive decomposition)進行處理。
圖 2:WarpSelect 概覽。輸入值從左邊流入,右邊的 warp queue 保存其輸出結果
算法 3:WarpSelect 的偽代碼
以下是對 FAIR 開源的相關項目 Faiss 的 README.md 文件的編譯介紹:
項目地址:https://github.com/facebookresearch/faiss
Faiss 是一個用於有效的相似性搜索和密集向量聚類的庫。其包含了在任何大小(甚至可以大到裝不進 RAM)的向量集中進行搜索的算法。其也包含用於評估和參數調整的支持性代碼。Faiss 是用 C++ 編寫的,且帶有用於 Python/numpy 的完整封裝器(wrapper)。其中一些最有用的算法是在 GPU 上實現的。Faiss 由 Facebook AI Research 開發。
介紹
Faiss 包含幾種用於相似性搜索的方法。其假定其實例可以被表示成向量且可以通過整數識別,然後這些向量可以與 L2 實例或點積進行比較。類似於一個查詢向量(query vector)的向量是那些有最低 L2 實例或最高點積的帶有查詢向量的向量。其也支持餘弦相似性(cosine similarity),因為這是在歸一化的向量上的點積。
這裡的大多數方法都和那些基於二元向量和緊湊的量化代碼類似,僅使用了一種對向量的壓縮過的表示,也不要求保持原始向量。這通常會導致搜索準確度下降,但這些方法可以在單個向量上在主內存中擴展到數十億個向量。
該 GPU 實現可以接受來自 CPU 或 GPU 內存的輸入。在一個帶有 GPU 的伺服器上,其 GPU 索引可以被用作其 CPU 索引的插入替換(比如,使用 GpuIndexFlatL2 替代 IndexFlatL2),且來自或發往 GPU 內存的副本可以被自動地處理。
構建
該庫基本上是通過 C++ 實現的,帶有可選的通過 CUDA 的 GPU 支持和一個可選的 Python 接口。它使用了一個 Makefile 進行編譯。參見 INSTALL 了解詳情:https://github.com/facebookresearch/faiss/blob/master/INSTALL
Faiss 的工作方式
Faiss 是圍繞一種存儲了一個向量集的索引類型(index type)而構建的,並且提供了一個使用 L2 和/或點積向量比較在其中進行搜索的函數。一些索引類型是簡單的基線,比如精準搜索。大部分可用的索引結構都對應了與以下方面的權衡:
搜索時間搜索質量每個索引向量所用的內存訓練時間無監督學習時對外部數據的需求
Faiss 的完整文檔
以下連結可訪問文檔:
完整文檔(包括一個教程)請參閱這個 wiki 頁面:http://github.com/facebookresearch/faiss/wikidoxygen 文檔給出了每個類的信息:http://rawgithub.com/facebookresearch/faiss/master/docs/html/annotated.html要重現這篇研究論文的結果,請參考基準 README:https://github.com/facebookresearch/faiss/blob/master/benchs/README.md
本文為機器之心編譯,轉載請聯繫本公眾號獲得授權。