はじめに
kNNなどの近傍探索はpythonやnumpyだけだとデータ数に応じて時間がだいぶかかるようになります。
もちろん厳密なNNではなく近似最近傍探索(ANN search)を行うのが中心かと思います。
Pythonから使え、ANNをC等で最適化しているライブラリがあったのでいくつか試してみました。
弊社内ではFLANNやFaissが使われていました。
今回はこれらと追加で、人気のありそうなAnnoy、速そうなNMSLIBを比較してみます。
（社内kibelaでまとめたものを転載してみました）
参考記事
まとめる際のきっかけになった記事です。
ベンチマーク等はこちらが参考になります。
-
ann-benchmarks
- 著者は実はAnnoyの作者なのに、他のいいライブラリを勧めてくれています。
- ブログ版
-
- この図を見るとnmslibとfaissが特に早そう。
-
Approximate Nearest Neighbours for Recommender Systems
- こちらはnmslibのpython bindingを作ってる？人の記事です。
対象ライブラリ
上の記事を参考に候補を選びました。
| ライブラリ | GH star | メモ |
|---|---|---|
| Sci-kit learn | (総合なので省略) | python+numpy。遅そう |
| FLANN | 939 | ブリティッシュコロンビア大学の人製。古め？pyflannがpy3対応してなかった |
| Annoy | 3976 | spotify製。少し前に人気があった。お手軽と評判？ |
| faiss | 4420 | facebook製。一番starが多い。この中だと唯一GPUが使える。 |
| nmslib | 988 | Non-metric space? 上のベンチマーク記事だとCPU環境で一番早い。 |
近似近傍探索のアルゴリズムもライブラリ間で差がありそうですが、今回は精度含め、あまり気にしないでおきます。
上の記事では精度と速度のトレードオフを比較しているので参考にしてください。
色々比較
データ作成
データはfaissのwikiにある作り方で作成しました。
スパースでないただの乱数なので現実のパフォーマンスとは乖離はありそうです。
import numpy as np
d = 64 # dimension
nb = 100000 # database size
nq = 10000 # nb of queries
np.random.seed(1234) # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.
Sci-kit learn
みんな大好きscikit learn。PyData製
インストール
pipでシンプルに可能です。
$ pip install scikit-learn
使い方
from sklearn.neighbors import NearestNeighbors
nn = NearestNeighbors(metric='cosine')
nn.fit(xb)
dists, result = nn.kneighbors(xq, n_neighbors=5)
print (result)
print (dists)
他のモデルと同じようなIFなので簡単ですね。
FLANN(pyflann)
Fast Library for Approximate Nearest Neighborsの略
インストール
$ pip install pyflann
# python3系では動かないのでPRされてるものを入れます。
$ pip install git+ssh://git@github.com/EdsterG/pyflann.git
もちろん本体のコードからbuildしてinstallした方が確実かと思います。
使い方
from pyflann import FLANN, set_distance_type
set_distance_type('manhattan') # defaultはeuclidean
flann = FLANN()
flann.build_index(xb)
result, dists = flann.nn_index(xq, num_neighbors=5)
print (result)
print (dists)
cosine距離とかはサポートしていないですが、L2正規化済みのinputならeuclid距離と共通になりそうです。
Annoy
Approximate Nearest Neighbors Oh Yeahの略。C++製
インストール
pipのみでお手軽に可能
$ pip install annoy
使い方
from annoy import AnnoyIndex
# 最初にvectorサイズを入れる
t = AnnoyIndex(xb.shape[1])
for i, v in enumerate(xb):
# indexつけて一個ずつ入れる必要がある。これが遅そう。。
t.add_item(i, v)
t.build(10) # 10 trees
result = []
dists = []
for i, v in enumerate(xq):
# 出すのも一つずつ。。
r, d = t.get_nns_by_vector(v, 5, include_distances=True)
result.append(r)
dists.append(d)
print(np.array(result))
print(np.array(dists))
コード量もシンプルで扱いやすいが、numpyをそのまま突っ込んだり予測させたりできないのが難点。
この部分でpython loopを使う必要があります。
Faiss
Facebook AI Similarity Searchの略。
名前からすると、近傍探索専用じゃないのかもしれません。
（INSTALL.mdにFair AI Similarity Searchの略みたいに書いてありましたが、FAIR = Facebook AI Researchですね。。）
インストール
pytorchもそうですが、Facebookは過激なconda信奉者なのかわかりませんが、pypiにはあげてなくて、condaからしかインストールできません。
それだけでゲンナリしそうですが、GPU使うためには更にCUDA依存も解決する必要がありそうです。
参考
# CPU version only
conda install faiss-cpu -c pytorch
# Make sure you have CUDA installed before installing faiss-gpu, otherwise it falls back to CPU version
conda install faiss-gpu -c pytorch # [DEFAULT]For CUDA8.0
conda install faiss-gpu cuda90 -c pytorch # For CUDA9.0
conda install faiss-gpu cuda91 -c pytorch # For CUDA9.1
そういう苦難もビルドすれば、もう少しなんとかなるかもしれません。
使い方
import faiss
index = faiss.IndexFlatL2(xb.shape[1]) # build the index
index.add(xb) # add vectors to the index
dists, result = index.search(xq, k=5) # actual search
print(result)
print(dists)
# ボロノイ分割してから探索する速いモードもある。
nlist = 100
quantizer = faiss.IndexFlatL2(data.shape[1]) # the other index
index = faiss.IndexIVFFlat(quantizer, data.shape[1], nlist, faiss.METRIC_L2)
# here we specify METRIC_L2, by default it performs inner-product search
index.train(data) # こちらはtrainで分割してくれる用。
index.add(data) # add may be a bit slower as well
# 分割したクラスタの近接クラスタも探索する場合はprobeを広げる。デフォルト1
# index.nprobe = 10
dists, result = index.search(test, k=5) # actual search
print(result)
print(dists)
# GPU使うとさらに5-10倍早いそう。
使い方はだいぶシンプル。
Indexクラスが何種類もあるのが少しわかりづらいかも。
詳しい使い方は公式やこの記事を見るといいと思います。
NMSLIB
Non-Metric Space Libraryの略
インストール
pipのみでお手軽に可能
$ pip install nmslib
使い方
import nmslib
# Annoy同様にデータを入れてbuildする。Numpy配列で入れられる。
index = nmslib.init(method='hnsw', space='cosinesimil')
index.addDataPointBatch(xb)
index.createIndex({'post': 2}, print_progress=True)
# 基本的にAnnoy同様に一件ずつ検索して、返却される。
ids, distances = index.knnQuery(xq[0], k=5)
print(ids)
print(distances)
# マルチスレッド処理用のメソッドがある。
# ただし、中身は普通のクエリと同様一件ずつなので、まとめる場合は処理が必要。
neighbours = index.knnQueryBatch(xq, k=5, num_threads=4)
result = np.concatenate([i[None,:] for i,d in neighbours])
dists = np.concatenate([d[None,:] for i,d in neighbours])
print(result)
print(dists)
使い方はあまり差がありませんが、インターフェースは少し癖があります。
アルゴリズムもデフォルトでHierarchical Navigable Small World(HNSW)というものを使うようです。
これ自体はfaissにも実装されているみたいですね。
適当な速度比較
試しに、CPU環境で、上のデータで10回繰り返しの速度を測定してみました。
条件や精度は測定せずに、上のコードをそのまま適当に測ってるので、お茶にごし程度と思って下さい。
| ライブラリ | 学習込み | 学習済み | 感想 |
|---|---|---|---|
| Sci-kit learn | 10.3 s ± 20.6 ms | 10.2 s ± 20.2 ms | 劇遅。学習は全く何もしてないので分離しても遅いまま。 |
| FLANN | 213 ms ± 4.64 ms | 89.3 ms ± 241 µs | 普通に速い。学習すればさらに速い。 |
| Annoy | 1.54 s ± 10.5 ms | 157 ms ± 584 µs | 学習が遅い！python loopを含むせいかも。学習させても爆速ではない。 |
| faiss.IndexFlatL2 | 288 ms ± 2.01 ms | 278 ms ± 4.26 ms | 速い。flannの次に速い。でも学習はなにもしてないのか、早くはならない。 |
| faiss.IndexIVFFlat | 63 ms ± 222 µs | 15.5 ms ± 42.9 µs | 学習込みでも爆速！学習なしで最速。精度を犠牲にしているのか？GPU使ったらどうなる？ |
| nmslib | 4.28 s ± 26.6 ms | 43.2 ms ± 2.19 ms | 学習はだいぶ遅い。。ベンチマークはなんだったのか。。使い方間違ってる？学習を除けばだいぶ速いが、最速ではない。。 |
学習込みか無しかは、build系の処理をループ内に含めるかそうでないかです。
速度だけで行ったらfaissがベストな選択だと言えそうです。
とはいえ、速度x精度の比較でnmslibが勝っていたりするらしいので、その辺を調べて検討すべきかもしれません。
おまけ
上で紹介した、NMSLIBメンバーのブログにあった比較グラフ。
annoy/faiss/nmslibでの比較ですがnmslibがトレードオフ上は良いでしょう。
トレーニングが頻発する場合はまた別だと思います。
ただ、GPU使用版のfaissを見ると、
あ、圧倒的じゃないか・・・
それでも、この記事では筆者はGPUのない環境ではNMSLIBをおすすめ！と結論していますね。
感想
一番速度メリットが大きいのはFaissで間違いないかと思います。
開発主体もFacebookですし、安心感もあるかなと思います。
GPU使うということなので、環境周りが若干面倒かもしれません。
個人的感想
- sklearn :
-
環境構築不要（だいたい入ってる） -
遅すぎるのでサンプルにしか使わない
-
-
FLANN
-
ちょっと試すには手軽でいい。
-
論文等でよく使われているので実績も多く安心。
-
機能が少ない。
-
pip installできるpythonバインディングが不安
-
-
Annoy
-
ちょっと試すには手軽でいい
-
本体コードの量が少なくて安心
-
python使用例が多そうで、bindingも使いやすい
-
速度やインターフェースがイマイチ
-
-
Faiss
-
最速。GPU使用可能。
-
使いやすさも悪くない。
-
メジャーでstarも多く、Facebook開発なので(?)安心。
-
installが若干面倒。condaに飼いならされてないので。。
-
python bindingはSWIG。メモリエラーが捕まらない？
-
-
NMSLIB
-
速いらしいが使い方がわかってないせいか、そこまで早くない。
-
ちょっとマイナーで、不安
- 今後に期待
-
直近の用途では複数の探索範囲をサーチする必要があり、メモリ効率も大事だったようで、
FLANN/Faiss/NMSLIBはpythonからIndexモデルを作りすぎるとメモリエラーになってしまいました。
そのため、Annoyでしのいでしまいました。（使い方が悪いだけかと思いますが。。）
探索範囲が固定でindexがbuild済みのケースとはまた違ったかと思います。
ビルド済みindexの読み込み効率なども比較した方がいいかもしれませんね。
参考
- 比較記事
-
ann-benchmarks
- 著者は実はAnnoyの作者なのにそれよりいいライブラリを進めてる感じ。
- ブログ版
- Approximate Nearest Neighbours for Recommender Systems
-
ann-benchmarks
- Faiss
- Annoy