GPU対応の類似検索(最近傍探索)ライブラリ Faissの紹介 part2 GPUを利用した検索

Facebook AI Research (FAIR)が開発したGPU対応の類似検索ライブラリ Faiss に関する2回目の紹介エントリーとなります。前回は、FaissのインストールとC++チュートリアルの説明を行いました。今回はGPUを利用した検索処理を試してみます。

Faissのインストール方法については前回のエントリーを参照してください。

環境

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

今回からはプログラムのベンチマークも行っていくので自宅のGPUサーバをメインに使って検証します。EC2のp2.xlargeをがっつり使うと請求が怖いし、そもそもレイテンシが大きくてストレスが溜まるので。

前回の導入部分の補足としてCUDAコード部分のコンパイルにはGCCではなくNVCC(CUDAコンパイラ)を利用するので事前にインストール済みであることを確認しておきます。また、GPU製品のCompute Capabilityは3.5以降が必要となっています。

GPUでk-meansクラスタリング

まずは、公式ドキュメントにサンプルとして掲載されているGPUを利用したk-meansクラスタリングプログラムの動作確認をします。日本語での補足も簡単に入れておきます。

faiss::gpuに属するクラスについては後述しますが、faiss::gpu::GpuIndexFlatL2 は前回のエントリーで紹介した faiss::IndexFlatL2 のGPU版となっています。faiss::Clustering::train 実装内部ではイテレーション毎に現在のセントロイドをインデックスしておき、全データをクエリとして検索しています。

* コンパイル

* 実行結果例

GPUを利用したk-meansクラスタリングは過去にもいろいろライブラリ等が出ていますし、サンプルとしてはよく見るものです。ここでの動作確認では5,000,000個のデータを20,000クラスタに分けるのに4分弱で処理できました。いまいちよく分からないのでCPU版と比較してみます。

上記プログラムをCPUでのクラスタリング処理に変更するには、インデックスオブジェクトである faiss::gpu::GpuIndexFlatL2faiss::IndexFlatL2 に変えるだけです。ただし上記プログラムの設定だとCPU版の処理が一向に終わらなかったので以下のようにテストの規模を縮小することで調整しました。

  • データ(128次元): 500,000
  • クラス(分割数): 2,000
  • イテレーション: 20

500,000個のデータを2,000クラスタに分けます。それぞれ5回計測した平均処理時間は以下の通りです。また、参考までに各製品のリリース時期も併記しています。

CPU
Intel Core i5-2500K
(2011/1)
CPU
Intel Xeon E5-2686
(2016/6)
GPU
NVIDIA GM204 [GeForce GTX 970]
(2014/9)
GPU
NVIDIA GK210GL [Tesla K80]
(2014/11)
511.15s 114.10s 5.70s 4.85s

前回のエントリーにも書いた通り、CPU版はOpenMPやSIMDを駆使して実装されており十分に高速です。しかしGPU版はそれと比較しても遙かに高速に動作します。また、CPUからGPUへのデータ転送は必要に応じて都度行われるため(計算に必要なデータだけGPU上に載せる)、VRAM容量を超えるような巨大なデータセットでもクラスタリング可能となっています。

さらに、FP16(半精度浮動小数点値での演算)を有効にしてGTX 970上で比較してみました。こちらも同様に5回計測した平均処理時間を載せます。

  • データ(128次元): 5,000,000
  • クラス(分割数): 20,000
  • イテレーション: 20
FP16無効 FP16有効
228.47s 191.76s

なかなか興味深い結果となりました。確かPascal世代からFP16の性能が格段に上がっていたかと思うので(1SPでFP16二つを同時演算したり)、PascalのGPUでベンチマークしてみたいところです。Facebook ResearchのレポートだとTesla P100でだいたい100秒ほどで完走するらしいです。
* GPU k means example

GPUを利用した検索

次は前回紹介したCPU版チュートリアルにある各種検索プログラムをGPU対応していきます。

GPUリソースの確保・初期化

GPUで処理するために予めGPUリソースの事前確保作業が必要ですが、生のCUDAプログラムと比較すれば十分抽象化されているので特に難しいところは無いかと思います。GPU版のヘッダファイルは faiss/gpu 以下にあります。

  • faiss::gpu::StandardGpuResources
  • faiss::gpu::GpuIndexFlatConfig

前述のk-meansのサンプルコードと同じ設定でだいたい大丈夫です。内部ではcuBLASのハンドル初期化やCUDAのストリーム生成処理等を行っています。デフォルトだとストリーム(cudaStream_t)は2本準備しているようです。

GPU版インデックスクラス

チュートリアルで紹介した各種CPU版インデックス用クラスのGPU対応は簡単なのでまとめて紹介します。

CPU GPU
faiss::IndexFlatL2 faiss::gpu::GpuIndexFlatL2
faiss::IndexIVFFlat faiss::gpu::GpuIndexIVFFlat
faiss::IndexIVFPQ faiss::gpu::GpuIndexIVFPQ

2017/05現在、faiss::gpu::GpuIndexIVFPQ にはGPU版インデックスクラスのポインタを直接受け取るコンストラクタは未実装のようですので別のコンストラクタを利用しています。この場合は引数で渡す faiss::MetricType に応じたGPU版インデックスクラスのインスタンスを内部で生成/保持してくれます。前回紹介していませんでしたが距離尺度は faiss::MetricType 定数で指定するようになっており、チュートリアルではL2ノルムの紹介のみでしたが、もちろん内積(Inner Product)にも対応しています。

また、faiss::gpu::IndicesOptions は以下のように定義されています。これはCPU<->GPU間でのインデックス情報の転送・保存をどのような形式で行うかを指定します。

CPU版とは異なりGPU版の各種インデックスクラスのメンバ変数はなぜかカプセル化されているみたいなので、提供されているgetter/setter経由でデータ操作を行います。GPU対応する際には忘れないようにしましょう。

ベンチマーク

次は最近傍探索処理の測定を行います。チュートリアルのプログラムをベースにパラメータ等を調整して比較しました。論文中ではDeep1BYFCC100Mなどの巨大なデータセットでベンチマークをしていて真似したかったんですけど、自宅のGPU開発環境だと空冷が適当なので kidle_inject が発生し、温度上昇防止のためCPU使用に制限をかけてしまいます。。ここではANN_SIFT1Mと同じ小規模データでベンチマークを行いました。

  • 次元数: 128
  • データ数: 1,000,000 (100万データ)
  • 同時実行クエリ数: 10,000 (1万クエリ)
  • インデックス: faiss::gpu::IndexFlatL2 (全探索) / faiss::gpu::IndexIVFPQ (PQ + 転置インデックス)
  • 転置インデックスのリストサイズ 100、サブベクトル数 8、量子化ビット数 8、FP16有効
  • 検索時のkの値: 1 ~ 1024 の範囲で変更

ベンチマーク用プログラムはここに記載するには少し長くなるのでGistにアップロードしておきます。

僕の個人開発環境だと1万クエリの上位20件を検索するのに、全探索(IndexFlatL2)の場合は約1.1秒、PQ + 転置インデックス(IVFPQ)の場合は約20ミリ秒でした。ただし、Faissの真価はbillion-scaleでの検索処理において発揮されるので、ここでのベンチマークは参考程度ということで。

また、Facebook AI Research公式の詳細なベンチマーク結果は以下を参照してください。

注意点など

GPUでの検索処理にはいくつか制限が設けられています。2017/06現在の仕様なので変更される可能性があります。

GPUテンポラリメモリ

cudaMalloc/cudaFree はコストが高い処理なので事前にある程度のメモリ領域が自動的に確保されるのですが、そのサイズをユーザー側が指定することができます。StandardGpuResources のAPIで指定できます。

テンポラリメモリが足りないと以下のような警告が出力されます。cudaMalloc が新たに必要になってしまうから遅くなるよというヒントであって処理が継続できないわけではありません。

PQ(Product Quantization)のために内部ルックアップテーブルを準備する際にはそこそこ十分な容量を確保しておいた方が良いかと思います。

kおよびnprobe値の上限

kとnprobe値は1024が上限となっていますが、1024もあれば実用上は十分かなと思います。

Product Quantizationのパラメータ

PQの実装上、量子化ビット数(サブベクトル当たりのデータサイズ)は8ビット以下、データの次元数をサブベクトル数で割った余りが0になるように調整する必要があります。今回紹介したベンチマークプログラムではデータは128次元、サブベクトル数は8、量子化ビット数は8で設定しています。

おわりに

今回はFaiss GPU版のインタフェースを紹介しました。CPU版から移行修正しやすいように工夫されているので特に難しいところはなかったかなと思います。GPU版の実装は論文と併せて読み進めていますが、k-selection(選択アルゴリズム)のGPU実装の仕組みが難しくて難航しています。。理論面を調べつつ、そろそろPythonインタフェースを用いたアプリケーション寄りの開発もやってみたいと思っています。

あわせて読む:

コメントを残す

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