ナード戦隊データマン

データサイエンスを用いて悪と戦うぞ

Facebook ResearchのfaissでKmeansを実行 (python)

faissとは、Facebook Researchによる高速な類似性検索・クラスタリングのためのライブラリです。以前は、類似性検索を試したので今回はクラスタリングを試します。

jupyterで実行

前回はテキストをベクトル化しましたが、Kmeansではデータ数がある程度必要なので、numpyで生成します。

In[1]:

import numpy as np
d = 64                           
nb = 3000                   
np.random.seed(1234)            
x = np.random.random((nb, d)).astype('float32')

次に、Kmeansの訓練をします。

In[2]:

ncentroids = 10
niter = 10
verbose = True
d = x.shape[1]
kmeans = faiss.Kmeans(d, ncentroids, niter, verbose)
kmeans.train(x)

教師なし学習なので、事前に生成したxがどう分類されたのかをそのまま見ることができます。

In[3]:

D, I = kmeans.index.search(x, 1)
I[:10]

Out[3]:

array([[1],
       [6],
       [6],
       [5],
       [5],
       [1],
       [1],
       [3],
       [0],
       [9]])

訓練したオブジェクトkmeansにcentroidsの情報があるので、それを用いればインデクスからcentroidsに最も近い順にIDを取得できます。

In[4]:

index = faiss.IndexFlatL2 (d)
index.add (x)
D, I = index.search (kmeans.centroids, 15)
centroid_id = 3
I[3]

Out[4]:

array([2912, 2440, 1273, 1682,  917,  173, 1905, 2480, 2992, 1001, 2406,
       2298,  364,  889, 2078])

モジュール化

インデクスに対するページング機能を実装したかったため、簡易的にモジュール化しました。

def search_page_by_faiss(index, kmeans, centroid_id, page=1, size=10):
    D, I = index.search(kmeans.centroids[centroid_id:centroid_id+1], index.ntotal)
    return I[0][size*(page-1):size*page]

ただし、実際にはfaissのIDと、何らかのドキュメント(あるいは画像等)のIDを対応させるためのマッピングを作成する必要があります。マッピングを用いる場合は、このコードは以下のように書き換えられます。

def search_page_by_faiss(index, kmeans, indices2id, centroid_id, page=1, size=10):
    D, I = index.search(kmeans.centroids[centroid_id:centroid_id+1], index.ntotal)
    return [indices2id[i] for i in I[0][size*(page-1):size*page]]

ちなみに、faissについてまだそこまで詳しく調べてないので、「上位k件を取得する」以外の方法でページングができるのか不明です。ただ、上位k件の取得はとても高速なので、インデクスのサイズが膨大にならない限りはこのモジュールでもまあ足りると思います。

おわりに

Kmeansを訓練し、さらにそのcentroidsを用いてインデクスから近傍探索する方法を書きました。もし、faissで効率よくページング機能を実装する方法があれば教えてください。

参考

[1] https://github.com/facebookresearch/faiss/wiki/Getting-started [2] https://github.com/facebookresearch/faiss/wiki/Faiss-building-blocks:-clustering,-PCA,-quantization