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でコンパイルした方がいいかと思います。
## ソースコードのダウンロード
$ git clone https : //github.com/facebookresearch/faiss.git
$ cd faiss
## Linuxの場合は makefile.inc.Linux を利用
$ cp example_makefiles / makefile . inc . Linux . / makefile . inc
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でこつこつと実装されているようです。
## Ubuntuの場合
$ sudo apt - get install libopenblas - dev
## CentOS or Amazon Linuxの場合
$ sudo yum install openblas - devel
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
$ diff makefile . inc example_makefiles / makefile . inc . Linux
63c63
< #BLASLDFLAGS?=/usr/lib64/libopenblas.so.0
-- -
> BLASLDFLAGS ? = / usr / lib64 / libopenblas . so . 0
67c67
< BLASLDFLAGS ? = / usr / lib / libopenblas . so . 0
-- -
> # BLASLDFLAGS?=/usr/lib/libopenblas.so.0
112 , 113c112
< #PYTHONCFLAGS=-I/usr/include/python2.7/ -I/usr/lib64/python2.7/site-packages/numpy/core/include/
< PYTHONCFLAGS = - I / usr / include / python3 . 5m - I / usr / local / lib / python3 . 5 / dist - packages / numpy / core / include
-- -
> PYTHONCFLAGS = - I / usr / include / python2 . 7 / - I / usr / lib64 / python2 . 7 / site - packages / numpy / core / include /
134a134
> - gencode arch = compute_35 , code = "compute_35" \
135a136
> - gencode arch = compute_60 , code = "compute_60" \
GTX 970のCompute Capability(CC)は5.2なのでそれ以外のコード生成オプションは削っておきます。デフォルトのままでも問題ないですけど、無駄を省いてコンパイル時間を短くしたいならハードウェア環境に合わせて設定すると良いです。
PythonインタフェースはSWIG(3.x以上)でビルドされます。依存ライブラリとしてNumPyも必要なので事前にインストールしておきます。
## SWIGのインストール
$ sudo apt - get install swig
## NumPyのインストール
$ sudo pip install numpy
Makefileの準備と依存ライブラリの導入が終わったら後はコンパイルするだけですが、最低限 SSE4.2/POPCOUNT/OpenMP に対応しているCPU/C++コンパイラが必要となっています。AVXは使ってないようなのでSandy Bridge以前のかなり古いモデルでも大丈夫なはずです。
## CPU版 C++ライブラリ
## jオプションはCPUコア数に合わせる
$ make - j4
## デフォルトではstatic libraryが生成される、もしshared libraryが欲しいなら
$ make libfaiss . so
## CPU版 Pythonライブラリ
$ make py
## GPU版 C++ライブラリ
$ cd gpu
$ make - j4
## GPU版 Pythonライブラリ
$ make py
GPU版 C++ライブラリはコンパイルにかなり時間がかかると思います。最近だとこういう手順でプロダクトをインストールするのは逆に新鮮に感じます。リリース初期の頃のCaffeみたいだ。
ちなみに付属のMakefileはinstallターゲットがないので、手動でライブラリとヘッダファイル群を任意の場所に配置します。 ここでは既にパスの通ったディレクトリに置きました。
$ ldconfig - p | grep faiss
libfaiss . so ( libc6 , x86 - 64 ) = > / usr / local / lib / libfaiss . so
$ ls / usr / local / include / faiss
AutoTune . h DeviceUtils . h GpuIndexFlat . h GpuIndicesOptions . h IndexFlat . h IndexPQ . h MetaIndexes . h StackDeviceMemory . h Timer . h hamming . h
AuxIndexStructures . h FaissAssert . h GpuIndexIVF . h GpuResources . h IndexIVF . h IndexProxy . h PolysemousTraining . h StandardGpuResources . h VectorTransform . h index_io . h
Clustering . h GpuAutoTune . h GpuIndexIVFFlat . h Heap . h IndexIVFPQ . h IndexWrapper - inl . h ProductQuantizer . h StaticUtils . h WorkerThread . h utils . h
DeviceMemory . h GpuIndex . h GpuIndexIVFPQ . h Index . h IndexLSH . h IndexWrapper . h RemapIndices . h TestUtils . h faiss . h
動作確認
CPU版 C++ライブラリのビルド時にテストプログラムもコンパイルされて実行可能バイナリが生成されているのでそのまま実行できます。GPU版も用意されているのですが、これは2017/05現在、バグがあって特定のGPUだとSEGVしてしまうので完走できません。どうやらGeForce系GPU(GTX 970や1080)でのクラッシュ報告が多いようです。EC2 p2-xlargeインスタンスのTesla-K80なら大丈夫でした。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
## CPU版
$ . / tests / demo_ivfpq _ indexing
[ 0.001 s ] Generating 100000 vectors in 128D for training
[ 0.168 s ] Training the index
Training IVF quantizer on 100000 vectors in 128D
Training IVF residual
Input training set too big ( max size is 65536 ) , sampling 65536 / 100000 vectors
computing residuals
training 4x256 product quantizer on 65536 vectors in 128D
Training PQ slice 0 / 4
Clustering 65536 points in 32D to 256 clusters , redo 1 times , 25 iterations
Preprocessing in 0.00 s
Iteration 24 ( 2.32 s , search 2.10 s ) : objective = 113338 imbalance = 1.004 nsplit = 0
Training PQ slice 1 / 4
Clustering 65536 points in 32D to 256 clusters , redo 1 times , 25 iterations
Preprocessing in 0.00 s
Iteration 24 ( 2.81 s , search 2.55 s ) : objective = 113197 imbalance = 1.004 nsplit = 0
Training PQ slice 2 / 4
Clustering 65536 points in 32D to 256 clusters , redo 1 times , 25 iterations
Preprocessing in 0.00 s
Iteration 24 ( 3.83 s , search 3.48 s ) : objective = 113227 imbalance = 1.009 nsplit = 0
Training PQ slice 3 / 4
Clustering 65536 points in 32D to 256 clusters , redo 1 times , 25 iterations
Preprocessing in 0.01 s
Iteration 24 ( 5.52 s , search 5.06 s ) : objective = 113061 imbalance = 1.004 nsplit = 0
[ 21.231 s ] storing the pre - trained index to / tmp / index_trained . faissindex
[ 21.233 s ] Building a dataset of 200000 vectors to index
[ 21.705 s ] Adding the vectors to the index
add_core times : 2477.588 806.523 13.101
[ 25.002 s ] imbalance factor : 1.23894
[ 25.002 s ] Searching the 5 nearest neighbors of 9 vectors in the index
[ 25.006 s ] Query results ( vector ids , then distances ) :
query 0 : 1234 11667 163213 83335 13439
dis : 7.09815 9.59665 9.70155 9.80191 9.8625
query 1 : 1235 76478 56663 28798 117320
dis : 7.79175 10.1169 10.4342 10.4568 10.6911
query 2 : 1236 189820 80604 185497 81842
dis : 7.60428 11.0645 11.1066 11.1707 11.2591
query 3 : 1237 116226 85618 36842 187787
dis : 8.03953 10.4461 10.6549 10.7921 10.7933
query 4 : 1238 91514 86365 52306 69284
dis : 7.4044 10.0915 10.1178 10.1657 10.2154
query 5 : 1239 143840 70772 2444 96181
dis : 7.18461 9.72012 10.2507 10.364 10.5746
query 6 : 1240 72767 170354 122756 61561
dis : 7.31866 9.74185 9.94351 9.97688 10.203
query 7 : 1241 30037 196333 162441 194151
dis : 8.19873 11.4148 11.6176 11.6699 11.6904
query 8 : 1242 66419 10135 189152 10030
dis : 8.2098 10.8027 11.1994 11.2308 11.2371
note that the nearest neighbor is not at distance 0 due to quantization errors
テストプログラム内部で行われていることはチュートリアルを進めていけばわかるようになるのでここでは省略します。公式のチュートリアルよりもコードは綺麗な気がしますけど。。
チュートリアル
公式のチュートリアルに補足を加えながら順番に動作確認していきます。
Faissが提供するクラスや各機能が扱うデータの型を意識しながら学んだ方がより理解しやすいと思うので、最初はC++インタフェースで試します。Pythonインタフェースは前述の通りSWIGで生成されているため、C++インタフェースが使えるようになればPythonインタフェースを使うための学習コストは少なくなります。
## g++によるコンパイルオプション例(shared library利用)
$ g ++ - Wall - O2 - Wl , -- no - as - needed - I / usr / include - L / usr / lib - lfaiss sample . cpp - o sample
データの準備
Faissが扱うベクトルデータは基本的に単精度(32bits)浮動小数点値を要素とした配列を用います。C++ならfloat型配列 、Python(NumPy)だと numpy.ndarray
オブジェクトとなり、1つのベクトルを1行としたrow-majorな行列データとして扱うことになります。
int d = 64 ; // 次元数
int nb = 100000 ; // データベースのサイズ(格納データ数)
int nq = 10000 ; // 検索クエリのサイズ(クエリ数)
float * xb = new float [ d * nb ] ; // スタックオーバーフローを防ぐためヒープ上に作成
float * xq = new float [ d * nq ] ;
for ( int i = 0 ; i < nb ; i ++ ) {
for ( int j = 0 ; j < d ; j ++ ) xb [ d * i + j ] = drand48 ( ) ; // drand48() [0.0, 1.0) で一様分布する非負の倍精度浮動小数点実数値
xb [ d * i ] += i / 1000. ;
}
for ( int i = 0 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < d ; j ++ ) xq [ d * i + j ] = drand48 ( ) ;
xq [ d * i ] += i / 1000. ;
}
64次元ベクトルデータをデータベース登録用に10万個、検索クエリ用に1万個それぞれ作成しています。各ベクトルの1次元目の値だけ修正してますが、検索時にその意図がわかるのでここでは気にしなくても良いです。このチュートリアルでは通常のfloat型配列で入力データを管理していますが、Faissの実装を読んでみたところ、インデックス作成時に std::vector
のコンテナにデータを詰め替えています。もしOpenCVと併用するなら cv::Mat
に変換するアダプタを作っておきたいですね。
インデックスの作成
データベースのインデックスを作成します。ここからはFaissのAPIを使っていきますが、インデックスの種類はたくさんあるのでここでは基本的なものをピックアップして紹介していこうと思います。
faiss::IndexFlatL2
まずは基本となる距離尺度にL2ノルム(L2距離)を利用したインデックスの作り方を確認します。
faiss :: IndexFlatL2 index ( d ) ; // L2ノルムインデックスのインスタンス生成
printf ( "is_trained = %s\n" , index . is _ trained ? "true" : "false" ) ; // is_trained = true
// 全ベクトルデータのインデックス登録
// faiss::IndexFlatL2::add(idx_t n, const float *x) , idx_t: typedef long
index . add ( nb , xb ) ;
printf ( "ntotal = %ld\n" , index . ntotal ) ; // ntotal = 100000
IndexFlatL2
はL2ノルムでの全探索用(Brute-force search)インデックスとなっており、入力ベクトルデータは未加工のまま全て登録します。publicなメンバとして is_trained
と ntotal
を参照していますが、それぞれインデックスの学習済フラグ(データ分布等を事前解析、IndexFlatL2では利用されない)、登録データ数を示しています。
検索
データの登録が終わったら検索してみます。前手順で作成した検索クエリ用データを利用してk-近傍探索を行います。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
int k = 4 ;
{ // sanity check: search 5 first vectors of xb (データ確認用の簡易チェック)
long * I = new long [ k * 5 ] ; // 検索結果のインデックスを格納
float * D = new float [ k * 5 ] ; // 検索結果の距離(ここではL2距離)の値を格納
// faiss::IndexFlatL2::search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels) const
// 入力データ(つまりデータベースにインデックス済みの既知のデータ)5件をクエリにして検索
// k=4なのでクエリ毎に距離が近い順に4件まで検索
index . search ( 5 , xb , k , D , I ) ;
// print results
printf ( "I=\n" ) ;
for ( int i = 0 ; i < 5 ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
printf ( "D=\n" ) ;
for ( int i = 0 ; i < 5 ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%7g " , D [ i * k + j ] ) ;
printf ( "\n" ) ;
}
delete [ ] I ;
delete [ ] D ;
}
{ // search xq (xq:事前に作成した検索クエリ用ダミーデータ)
long * I = new long [ k * nq ] ;
float * D = new float [ k * nq ] ;
index . search ( nq , xq , k , D , I ) ;
// print results
printf ( "I (5 first results)=\n" ) ;
for ( int i = 0 ; i < 5 ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
printf ( "I (5 last results)=\n" ) ;
for ( int i = nq - 5 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
delete [ ] I ;
delete [ ] D ;
}
delete [ ] xb ;
delete [ ] xq ;
* 出力結果例
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
I = // 検索結果のインデックス (k=4なので距離が近い順に4件、5クエリ分)
0 371 187 551 // インデックス済データなので上位1件目は当然検索クエリ自身となる
1 324 126 279
2 213 260 487
3 199 402 626
4 388 839 222
D = // 検索結果の距離の値 (昇順ソート済)
0 7.52398 7.5694 7.79739
0 7.15704 7.19868 7.25345
0 6.61889 6.96309 7.00845
0 7.001 7.11882 7.48969
0 7.14342 7.89587 7.9118
I ( 5 first results ) = // ダミーの検索クエリ(未知データ)で同様に検索
486 287 1111 482
320 410 37 460
940 189 822 193
807 1002 388 505
657 156 95 1016
I ( 5 last results ) =
10842 10827 9938 10004
9403 10267 10880 10330
9896 10146 10093 10361
8603 10523 10582 9895
11460 10123 11099 10876
データベースに登録したデータ、検索クエリ用のデータも全てランダムに生成したにもかかわらず、類似検索結果のインデックスに偏りがあることがわかります。ここで、データ生成時に各ベクトルの1次元目のデータだけ少し加工したことを思い出しましょう。データ配列のインデックス値に依存する数値が加算されており、検索結果にもその影響が表れていることが確認できます。
for ( int i = 0 ; i < nb ; i ++ ) {
for ( int j = 0 ; j < d ; j ++ ) xb [ d * i + j ] = drand48 ( ) ;
xb [ d * i ] += i / 1000. ;
}
for ( int i = 0 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < d ; j ++ ) xq [ d * i + j ] = drand48 ( ) ;
xq [ d * i ] += i / 1000. ;
}
もし、データを加工しない場合は以下のような検索結果になります。
## 例
I ( 5 first results ) =
55040 31506 25128 42339
52766 50610 38693 41096
40905 55042 24436 4212
64192 15007 95432 32551
45146 62032 44525 82328
I ( 5 last results ) =
40461 79526 12960 79960
88504 16638 53654 1626
59759 33664 35785 835
67728 38098 19827 30085
81127 43351 79032 65174
2017/05現在のCPU版の実装では、次元数が4の倍数かつ同時検索クエリ数が20件以下の場合はSSEを使って距離計算をしてくれるようです。つまり、上記コード中の最初のsanity checkではSSEを利用、検索クエリデータ xq での検索ではインストール時に採用したBLAS実装(ここではOpenBLAS)に計算処理が委譲されます。また、マルチスレッドで処理される箇所はOpenMPを使って実装されているので大変読みやすくて助かります。
faiss::IndexIVFFlat
IndexFlatL2
をそのまま利用すると、いくらSSEやOpemMPを駆使して計算をしたとしても検索速度が遅くなってしまうため、Billion-scaleの大規模データに対応することは困難です。ここでは、入力ベクトルデータをクラスタリングしておき、検索時にはクエリデータに対応するセルおよび近隣のセルをいくつか選び出して、それらのセルに属するデータのみを最近傍点の候補として距離計算を行うことで検索時間を削減します。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
#include <memory>
. . . 省略
auto xb = std :: make_unique < float [ ] > ( d * nb ) ; // delete [] を排除したいのでチュートリアルを少しだけ改変
auto xq = std :: make_unique < float [ ] > ( d * nq ) ;
. . . 省略
faiss :: IndexFlatL2 quantizer ( d ) ;
int nlist = 100 ;
int k = 4 ;
// 転置インデックスクラス
// IndexIVFFlat(Index* quantizer, size_t d, size_t nlist_, MetricType = METRIC_INNER_PRODUCT)
// 量子化対象のインデックス、データの次元数、転置インデックスのリストサイズ、距離尺度
faiss :: IndexIVFFlat index ( &quantizer , d , nlist , faiss :: METRIC_L2 ) ;
assert ( ! index . is_trained ) ; // is_trained: false
index . train ( nb , xb . get ( ) ) ; // ベクトル量子化、転置インデックス作成 (データベース登録対象データを利用)
assert ( index . is_trained ) ; // is_trained: true
index . add ( nb , xb . get ( ) ) ; // データ登録
{ // search xq
auto I = std :: make_unique < long [ ] > ( k * nq ) ;
auto D = std :: make_unique < float [ ] > ( k * nq ) ;
index . search ( nq , xq . get ( ) , k , D . get ( ) , I . get ( ) ) ;
printf ( "I=\n" ) ; // print neighbors of 5 last queries
for ( int i = nq - 5 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
index . nprobe = 10 ; // default nprobe is 1, try a few more
index . search ( nq , xq . get ( ) , k , D . get ( ) , I . get ( ) ) ;
printf ( "I=\n" ) ;
for ( int i = nq - 5 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
}
* 出力結果例
I =
10049 10147 10188 10184
9403 9750 9812 10346
10361 10184 9920 10333
9895 9946 9335 9677
10876 9647 9756 11103
I =
10842 10827 9938 10004
9403 10267 10880 10330
9896 10146 10093 10361
8603 10523 10582 9895
11460 10123 11099 10876
IndexIVFFlat
は各セルの代表ベクトル(母点)を保存しておくための転置ファイル(転置インデックス)用のクラスとなっており、処理手順としてはベクトル量子化処理の学習ステップが一つ増えています(IndexIVFFlat::train
)。ここではデータベース登録対象データをそのまま学習に利用していますが、別途用意したデータで学習しても良いです。publicなメンバである nprobe
で検索対象のセル数を管理しており、この値を増やすと検索精度の向上と引き替えに検索時間が増えます。上記チュートリアルでは全探索と比較してnprobe=1で15 ~ 20倍、nprobe=10で3 ~ 5倍くらい検索速度が向上しました。
faiss::IndexIVFPQ
IndexFlatL2
および IndexIVFFlat
を併用した場合でも、入力ベクトルデータは全て未加工でインデックス登録していました。これでは数億オーダーのデータを全てメモリ上に載せるのは現実問題として困難です(NVMe等ストレージハードウェアも急激に進歩しているのですがここでは考えないこととします)。次は、入力ベクトルデータをProduct Quantization(PQ: 直積量子化)により圧縮することでインデックスするデータ量を削減します。PQはハッシュベースの手法に替わってここ数年の論文等でよく見かけますね。Product Quantizationに関する日本語の概要説明については以下の映像奮闘記さんのエントリーが読みやすくて参考になるかと思います。
クエリベクトルをM分割し、N/M次元のサブベクトルに分解する
クエリのサブベクトルと、それぞれのコードブックに記された代表ベクトル間の距離を予め計算しておく
それぞれのエントリがどのインデックスに位置しているのか確認し、2で計算した距離を参照して足し算を行う
最小距離をもつエントリを候補として提出する
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
. . . 省略
faiss :: IndexFlatL2 quantizer ( d ) ; // the other index
int nlist = 100 ;
int k = 4 ;
int m = 8 ;
// Product Quantizationを伴う転置インデックス
// IndexIVFPQ(Index * quantizer, size_t d, size_t nlist, size_t M, size_t nbits_per_idx)
// 量子化対象のインデックス、データの次元数、転置インデックスのリストサイズ、サブベクトル数、サブベクトルの量子化ビット数
faiss :: IndexIVFPQ index ( &quantizer , d , nlist , m , 8 ) ; // 8 specifies that each sub-vector is encoded as 8 bits
// 転置インデックス作成のための事前学習(ベクトル量子化処理)
// ここではデータベース登録対象データを利用して学習しているが、別途用意した学習用データを用いて良い
index . train ( nb , xb . get ( ) ) ;
index . add ( nb , xb . get ( ) ) ; // データ登録
{ // sanity check
auto I = make_unique < long [ ] > ( k * 5 ) ;
auto D = make_unique < float [ ] > ( k * 5 ) ;
index . search ( 5 , xb . get ( ) , k , D . get ( ) , I . get ( ) ) ;
printf ( "I=\n" ) ;
for ( int i = 0 ; i < 5 ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
printf ( "D=\n" ) ;
for ( int i = 0 ; i < 5 ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%7g " , D [ i * k + j ] ) ;
printf ( "\n" ) ;
}
}
{ // search xq
auto I = make_unique < long [ ] > ( k * nq ) ;
auto D = make_unique < float [ ] > ( k * nq ) ;
index . nprobe = 10 ;
index . search ( nq , xq . get ( ) , k , D . get ( ) , I . get ( ) ) ;
printf ( "I=\n" ) ;
for ( int i = nq - 5 ; i < nq ; i ++ ) {
for ( int j = 0 ; j < k ; j ++ ) printf ( "%5ld " , I [ i * k + j ] ) ;
printf ( "\n" ) ;
}
}
* 出力結果例
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
I =
0 66 283 647
1 658 421 8
2 349 1111 37
3 626 139 199
4 38 276 717
## sanity checkではインデックス済みデータをクエリとして検索しているが、検索結果1位の距離は0にはならない
D =
1.68216 6.97229 7.07727 7.20292
1.33991 6.47805 6.49344 6.8318
1.21912 6.37759 6.40442 6.42391
1.51774 5.83502 6.08639 6.62265
1.4905 6.57351 6.84026 6.85133
I =
10049 10053 9608 9938
10880 10392 10135 9812
10815 10660 10361 9747
9946 9921 9511 9246
10876 9250 10123 10869
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アプリケーション開発の現場だと触れる機会がないですから、この分野の技術はプライベートの時間を使ってじっくり学んでいきたいです。
参考
Tags: algorithm , c/c++ , hardware