はじめに

この記事はHUITアドベントカレンダー2018 7日目の記事になります。

Python版のfaiss-gpuについての解説記事がなかなか見当たらず苦労したので、使用方法をまとめました。

Faissとは

Facebook Resarchが提供する近傍探索ライブラリ。
Github：https://github.com/facebookresearch/faiss

特徴

C++, Python
GPUが使えるので爆速→参考
k-NN, k-Mean, PCAなどができる

今回やること

Install
k-NNの基本実装
k-NNの高速化

Githubにjupyter形式でコードをあげています。こちらでは手順の説明を最小限にして、パラメータについてより詳細な解説を入れています。

1. Install

cudaのバージョンチェック

ls /usr/local/cuda

anacondaでインストール

ここの通りに

# 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 cuda92 -c pytorch # For CUDA9.2
# cuda90/cuda91 shown above is a feature, it doesn't install CUDA for you.

Colaboratory上でインストール

Google Colaboratoryでインストールする場合、Driveをマウントしたあと、

ダウンロード（一度ダウンロードすればこの2行は実行しなくて良い）

# cudaバージョンにあったものを下記のリンク先から選択してダウンロード
!wget https://anaconda.org/pytorch/faiss-gpu/1.4.0/download/linux-64/faiss-gpu-1.4.0-py36_cuda9.2.148_1.tar.bz2
!tar xvjf faiss-gpu-1.4.0-py36_cuda9.2.148_1.tar.bz2

wgetのリンク：https://anaconda.org/pytorch/faiss-gpu/filesから選択

インストール（Colab起動毎に必要）

!cp -r lib/python3.6/site-packages/* /usr/local/lib/python3.6/dist-packages/
!pip install mkl

チェック

import faiss # -> Errorが出なければOK!

2. k-NNの基本実装

k-NN（k-Nearest Neighbors）概要

$d$次元ベクトル空間内におけるある未知のデータ点$\bf q$に対して、既知のデータ集合${\bf X} = \{ {\bf x}_0, \ldots, {\bf x}_{N-1} \}$の中から、指定した距離*において近傍の$k$個の点を選出する。
*例えばユークリッド距離

データの用意

まず、ベクトル空間の次元dを定義し、既知のデータ点を格納したデータベース配列と検索対象のデータ点を格納したクエリ配列を用意します。

import time
import numpy as np
import faiss
d = 64                         # ベクトルの次元(dimension)
nb = 1000000                    # データベースのサイズ(database size)
nq = 100000                     # クエリベクトルの数(nb of queries)
np.random.seed(1234)
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.

インデックス配列の用意

データを用意したら、データベースに対応するインデックス配列を用意します。インデックス配列では探索アルゴリズムを指定できます。

インデックス配列の例

IndexFlatL2 → ここで使用
- L2ノルムを指標として総当たりで完全な探索を行う（時間がかかる）←今回はこれを使う
IndexFlatIP
- コサイン類似度を指標として完全な探索を行う（時間がかかる）
IndexIVFFlat → 高速化で使用
- 大規模データでの探索も可能な高速化アルゴリズムを使用して探索を行う

その他のインデックス配列を使うにはここを参照。

初期化

インデックス配列の初期化を行います。GPU使用時はベクトル空間の次元dとresourceとflat_configを以下のように設定する必要があります。

print("BUILD THE INDEX")
res = faiss.StandardGpuResources()
flat_config = faiss.GpuIndexFlatConfig()
index = faiss.GpuIndexFlatL2(res, d, flat_config)

データを追加

次にインデックス配列にデータを追加します。これによってデータベースのデータ点を探索することができるようになります。

index.add(xb)
print(index.ntotal) # データ数nbを返す

近傍探索実行

インデックス配列のインスタンス.searchを用いて近傍探索を行います。.search(x, k)はクエリ配列xと求める近傍点の数kを引数に取り、クエリ配列内の各ベクトルに対する$k$個の近傍点までの距離を格納した配列Dと近傍点のインデックス値を格納した配列Iを返します。

s = time.time()
D, I = index.search(xq, k)
e = time.time()
print(I[-5:])
print(D[-5:])
print("time: {}".format(e-s)) # 実行時間計測

3. k-NNの高速化

クラスタリングによる高速化

大規模なデータセットでの検索となると、GPUを使ったとしても上記手法では遅すぎることがあります。これは、IndexFlatL2の探索アルゴリズムが総当たりで完全な結果を求めているからです。具体的には計算量が$n^2$オーダーであるため、あっという間に現実的な時間では解決できなくなってしまいます。
そこで、近傍探索の前にクラスタリングを行なって、高精度で近似的な探索をすることで探索を高速化することを試みた探索アルゴリズムがFaissでは複数用意されてます。今回は上述したIndexIVFFlatを使用して高速化を行っていきます。

ボロノイ領域

ボロノイ図（[Wikipedia](https://ja.wikipedia.org/?curid=91418)より）

クラスタリングによってデータ空間はボロノイ図状に分割されます。この区切られた空間一つ一つをボロノイ領域といい、各領域は中心点を持ちます。近傍探索ではクエリに対し、その周辺領域のみで探索を行わせることで、総当りの手法を避け、高速化を図っています。

IndexIVFFlat

IndexIVFFlatではd, resource, configに加え、ボロノイ領域の個数nlistと、距離基準METRIC_*を指定する必要があります。また、configにはIndexIVFFlat用のものを使用します。

nlist = 100 #ボロノイ領域の個数
res = faiss.StandardGpuResources()
ivfflat_config = faiss.GpuIndexIVFFlatConfig()
index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, ivfflat_config)
# faiss.METRIC_L2によって距離基準をL2ノルムに

ボロノイ領域の学習

IndexIVFFlatでは転置インデックスによって中心点を学習させていき、ボロノイ領域を学習します。詳しいアルゴリズムはわかりません…。

s = time.time()
index.train(xb)
e = time.time()
print("train time: {}".format(e-s)) # 学習時間を計測

近傍探索実行

決定したボロノイ領域を使用して高速探索を行います。この際、index.setNumProbes(nprobe)という関数を使用して探索する領域数を決める変数nprobeを設定します。これによって高速な近傍探索ができます！

index.add(xb)
index.setNumProbes(1)  # 1領域のみを探索
s = time.time()
D1, I1 = index.search(xq, k)
e = time.time()
print(I1[-5:])                  # neighbors of the 5 last queries
print("search time: {}".format(e-s))
print("accuracy: {:.1f}%".format(np.sum(I==I1)/(k*nq)*100))

s = time.time()
index.setNumProbes(10)         # default nprobe is 1, try a few more
D2, I2 = index.search(xq, k)
e = time.time()
print(I2[-5:])                  # neighbors of the 5 last queries
print("search time: {}".format(e-s))
print("accuracy: {:.1f}%".format(np.sum(I==I2)/(k*nq)*100))

変数と精度・高速化効果の関係

nprobe

ボロノイ図について、例えば上図のような二次元平面のある場所で赤点の領域と青点の領域が分けられていて、赤点に属するクエリ（緑）を与えた場合を考えましょう。もちろん、クエリの最近傍点は緑線で結ばれた青点です。
ですが、nprobe=1の時は赤点の領域しか探索しないため、最近傍点である青点の探索が行われません。一方で、nprobe=2とすれば、青点の領域も探索することができるため、最近傍点を正しく探索することができます。
このように、nprobeによって精度が上げられますが、一方で探索領域が広がることになるので高速化の効果は下がってしまいます。

nlist

ボロノイ領域の個数nlistが増加した時を考えます。
ボロノイ領域が増加するということは、それだけ細かく空間を分割しなければならないため、訓練時間がかかります。
ですが、訓練を終えてしまえば、1領域内に存在するデータ点の平均個数が低い状態での探索が行えるため、探索時間は短縮することができます。その一方で、境界線の数が増えるため、上記のような問題がより多く発生して精度が下がってしまいます。

トレードオフの関係

以上の関係をまとめると以下のようになります。

変数	精度	高速化
`nprobe`	比例	反比例
`nlist`	反比例	比例（訓練時は反比例）

よって、自分の行いたいタスクに求められる実行時間や精度からこれらのパラメータを適切に設定することが求められます。

次回

次回はkmeansの実装方法についてまとめます！

参考

本家Github
- Faiss
比較
- Approximate Nearest Neighbours for Recommender Systems
- 近似最近傍探索ライブラリ比較
Colaboratoryへのインストール方法
- Installing faiss on Google Colaboratory
基本実装
- Getting started（公式チュートリアル）
高速化
- Faster search（公式）
- ボロノイ領域 - Wikipedia