GPU対応の類似検索(最近傍探索)ライブラリ Faissの紹介 part1 導入/チュートリアル

Facebook AI Research (FAIR)が開発したGPU対応の類似検索ライブラリ Faiss を紹介します。

[06/25追記]
Faiss GPU版の検索についてエントリーを書きました。

論文は以下で公開されています。

論文タイトルの通り、10億スケールの大規模データに対してGPUを駆使した効率的な最近傍探索アルゴリズムの研究事例となっています。分量が多くなりそうなので複数の記事に分けて紹介しようと思います。

環境

Amazon Linux AMI release 2017.03
Intel Xeon CPU E5-2686 v4 @ 2.30GHz
NVIDIA GK210GL [Tesla K80]
Ubuntu 16.04 LTS (x86_64)
Intel Core i5-2500K CPU @ 3.30GHz
NVIDIA GM204 [GeForce GTX 970]
CUDA Toolkit 8.0
GCC 4.8.3 / 5.4.0
Python 3.5.2
NumPy 1.12

最初はTesla系GPUを使いたかったのでEC2のp2.xlargeインスタンス上で動作確認をしてましたけど、やっぱりオレゴンは遠くてターミナルがもっさりするので途中から自宅サーバに切り替えて作業しました。快適です。ちょっと前にGTX 1070を買いましたので、お下がりのGTX 970をLinuxサーバに挿しています。GTX 760は使い途がないのでメルカリかヤフオク行きですかね。

導入

※ GPU環境がある場合はNVIDIA DriverとCUDA Toolkit 8.0以降の事前導入が必要です

ここでは通常のC++ライブラリ(GPU非対応)、GPU対応のC++ライブラリ、Python用のライブラリをそれぞれソースコードからコンパイル/インストールします。プラットフォーム毎にMakefileから読み込む設定(makefile.inc.xx)が用意されているのでコピー・編集して使います。2017/05現在、CMake対応も進められているようですけれどまだ作りかけなので普通にmakeでコンパイルした方がいいかと思います。

CPU版C++/Pythonライブラリはビルド時にBLASライブラリをリンクしますが、ここではBLAS実装としてOpenBLASを利用しました。Intel MKLやAtlasでもOK、macOSならAccelerateフレームワークにBLAS実装が含まれているのでそれも利用できます。対応ライブラリの中ではIntel MKLが最速ですので、もしプロダクション環境等で可能な限り高速化したい場合はMKLをリンクさせましょう。GPU版はcuBLASが利用されます。こちらはCUDA Toolkitに含まれているので別途インストールする必要はありません。最近のCUDAだとBLAS drop-inライブラリとしてNVBLASが付属しているのでそちらをリンクするプロダクトも増えてきたように思いますが、FaissではcuBLASでこつこつと実装されているようです。

GTX 970のCompute Capability(CC)は5.2なのでそれ以外のコード生成オプションは削っておきます。デフォルトのままでも問題ないですけど、無駄を省いてコンパイル時間を短くしたいならハードウェア環境に合わせて設定すると良いです。

PythonインタフェースはSWIG(3.x以上)でビルドされます。依存ライブラリとしてNumPyも必要なので事前にインストールしておきます。

Makefileの準備と依存ライブラリの導入が終わったら後はコンパイルするだけですが、最低限 SSE4.2/POPCOUNT/OpenMP に対応しているCPU/C++コンパイラが必要となっています。AVXは使ってないようなのでSandy Bridge以前のかなり古いモデルでも大丈夫なはずです。

GPU版 C++ライブラリはコンパイルにかなり時間がかかると思います。最近だとこういう手順でプロダクトをインストールするのは逆に新鮮に感じます。リリース初期の頃のCaffeみたいだ。

ちなみに付属のMakefileはinstallターゲットがないので、手動でライブラリとヘッダファイル群を任意の場所に配置します。ここでは既にパスの通ったディレクトリに置きました。

動作確認

CPU版 C++ライブラリのビルド時にテストプログラムもコンパイルされて実行可能バイナリが生成されているのでそのまま実行できます。GPU版も用意されているのですが、これは2017/05現在、バグがあって特定のGPUだとSEGVしてしまうので完走できません。どうやらGeForce系GPU(GTX 970や1080)でのクラッシュ報告が多いようです。EC2 p2-xlargeインスタンスのTesla-K80なら大丈夫でした。

テストプログラム内部で行われていることはチュートリアルを進めていけばわかるようになるのでここでは省略します。公式のチュートリアルよりもコードは綺麗な気がしますけど。。

チュートリアル

公式のチュートリアルに補足を加えながら順番に動作確認していきます。

Faissが提供するクラスや各機能が扱うデータの型を意識しながら学んだ方がより理解しやすいと思うので、最初はC++インタフェースで試します。Pythonインタフェースは前述の通りSWIGで生成されているため、C++インタフェースが使えるようになればPythonインタフェースを使うための学習コストは少なくなります。

データの準備

Faissが扱うベクトルデータは基本的に単精度(32bits)浮動小数点値を要素とした配列を用います。C++ならfloat型配列 、Python(NumPy)だと numpy.ndarray オブジェクトとなり、1つのベクトルを1行としたrow-majorな行列データとして扱うことになります。

64次元ベクトルデータをデータベース登録用に10万個、検索クエリ用に1万個それぞれ作成しています。各ベクトルの1次元目の値だけ修正してますが、検索時にその意図がわかるのでここでは気にしなくても良いです。このチュートリアルでは通常のfloat型配列で入力データを管理していますが、Faissの実装を読んでみたところ、インデックス作成時に std::vector のコンテナにデータを詰め替えています。もしOpenCVと併用するなら cv::Mat に変換するアダプタを作っておきたいですね。

インデックスの作成

データベースのインデックスを作成します。ここからはFaissのAPIを使っていきますが、インデックスの種類はたくさんあるのでここでは基本的なものをピックアップして紹介していこうと思います。

faiss::IndexFlatL2

まずは基本となる距離尺度にL2ノルム(L2距離)を利用したインデックスの作り方を確認します。

IndexFlatL2 はL2ノルムでの全探索用(Brute-force search)インデックスとなっており、入力ベクトルデータは未加工のまま全て登録します。publicなメンバとして is_trainedntotal を参照していますが、それぞれインデックスの学習済フラグ(データ分布等を事前解析、IndexFlatL2では利用されない)、登録データ数を示しています。

検索

データの登録が終わったら検索してみます。前手順で作成した検索クエリ用データを利用してk-近傍探索を行います。

* 出力結果例

データベースに登録したデータ、検索クエリ用のデータも全てランダムに生成したにもかかわらず、類似検索結果のインデックスに偏りがあることがわかります。ここで、データ生成時に各ベクトルの1次元目のデータだけ少し加工したことを思い出しましょう。データ配列のインデックス値に依存する数値が加算されており、検索結果にもその影響が表れていることが確認できます。

もし、データを加工しない場合は以下のような検索結果になります。

2017/05現在のCPU版の実装では、次元数が4の倍数かつ同時検索クエリ数が20件以下の場合はSSEを使って距離計算をしてくれるようです。つまり、上記コード中の最初のsanity checkではSSEを利用、検索クエリデータ xq での検索ではインストール時に採用したBLAS実装(ここではOpenBLAS)に計算処理が委譲されます。また、マルチスレッドで処理される箇所はOpenMPを使って実装されているので大変読みやすくて助かります。

faiss::IndexIVFFlat

IndexFlatL2 をそのまま利用すると、いくらSSEやOpemMPを駆使して計算をしたとしても検索速度が遅くなってしまうため、Billion-scaleの大規模データに対応することは困難です。ここでは、入力ベクトルデータをクラスタリングしておき、検索時にはクエリデータに対応するセルおよび近隣のセルをいくつか選び出して、それらのセルに属するデータのみを最近傍点の候補として距離計算を行うことで検索時間を削減します。

* 出力結果例

IndexIVFFlat は各セルの代表ベクトル(母点)を保存しておくための転置ファイル(転置インデックス)用のクラスとなっており、処理手順としてはベクトル量子化処理の学習ステップが一つ増えています(IndexIVFFlat::train)。ここではデータベース登録対象データをそのまま学習に利用していますが、別途用意したデータで学習しても良いです。publicなメンバである nprobe で検索対象のセル数を管理しており、この値を増やすと検索精度の向上と引き替えに検索時間が増えます。上記チュートリアルでは全探索と比較してnprobe=1で15 ~ 20倍、nprobe=10で3 ~ 5倍くらい検索速度が向上しました。

faiss::IndexIVFPQ

IndexFlatL2 および IndexIVFFlat を併用した場合でも、入力ベクトルデータは全て未加工でインデックス登録していました。これでは数億オーダーのデータを全てメモリ上に載せるのは現実問題として困難です(NVMe等ストレージハードウェアも急激に進歩しているのですがここでは考えないこととします)。次は、入力ベクトルデータをProduct Quantization(PQ: 直積量子化)により圧縮することでインデックスするデータ量を削減します。PQはハッシュベースの手法に替わってここ数年の論文等でよく見かけますね。Product Quantizationに関する日本語の概要説明については以下の映像奮闘記さんのエントリーが読みやすくて参考になるかと思います。

  1. クエリベクトルをM分割し、N/M次元のサブベクトルに分解する
  2. クエリのサブベクトルと、それぞれのコードブックに記された代表ベクトル間の距離を予め計算しておく
  3. それぞれのエントリがどのインデックスに位置しているのか確認し、2で計算した距離を参照して足し算を行う
  4. 最小距離をもつエントリを候補として提出する

* 出力結果例

IndexIVFPQ はPQを伴う転置ファイル用クラスです。上記チュートリアルではサブベクトル数を8、各サブベクトルは8bitsにベクトル量子化されます。つまり、元データは64(次元数)*32(float)=2048bits、PQすると8*8=64bitsとなり32分の1のデータサイズとなります。また、各サブベクトル毎の距離計算は独立しているため並列処理とも相性が良さそうです。IndexFlatL2 による全探索(線形探索)結果を比較すると量子化の影響は見られますが、高い精度を維持できていることが確認できます。

注意点として、C++版チュートリアルのコードはGPUではなくCPUで動作します。GPU版実装の詳細は次回パート以降に紹介する予定です。

おわりに

今回はチュートリアルレベルの内容でしたので特に目新しい部分は無かったかと思います。Faissの論文は一通り読みましたが、主にGPU上で効率的に計算するアルゴリズムの記述が主でしたので、次パート以降はその理論や実装面についても深掘りしていきたいと思います。

SSEやOpenMP、CUDAを活用した処理は学生時代にけっこう勉強したんですけど、すっかり知識が抜けていたので焦りました。普段のWebアプリケーション開発の現場だと触れる機会がないですから、この分野の技術はプライベートの時間を使ってじっくり学んでいきたいです。

参考

あわせて読む:

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です