Faiss 튜토리얼 : Similarity Search와 Clustering을 빠르게!

Faiss

Similarity Search와 Clustering을 할 때, query/database vector의 length와 dimension이 큰 경우가 많다.

이런 경우 기존의 sklearn등을 활용한 방식으로는 시간이 너무 오래 걸린다.

이에 Facebook에서 공개한 효율적인 라이브러리가 있는데, 그것이 faiss이다.

오늘은 faiss의 기본적인 Python Interface를 살펴보도록 하겠다.

K-nearest-neighbor search

Faiss에서 Neighbor Search는 두가지 단계로 나뉜다.

먼저 database vector를 바탕으로 index를 만든다. Index는 database vector를 wrapping 하며 효율적인 search를 위한 전처리를 한다고 한다.

import faiss                   

db_vector = ...
query_vector = ...

index = faiss.IndexFlatL2(db_vector)   # build the index

이제 query vector를 index에 add한다.

index.add(query_vector)           

마지막으로 k-nearest-neighbor search를 진행한다.

D, I = index.search(query_vector, k)     # actual search

위에서 나오는 D, I는 (len(query_vecotr), len(db_vector)) 형태의 array인데, D는 찾아낸 nearest-neighbor와 query간의 거리이고, I는 찾아낸 nearest-neighbor의 database index이다.

K-means clustering

ncentroids = 1024
niter = 20
verbose = True
d = x.shape[1]
kmeans = faiss.Kmeans(d, ncentroids, niter=niter, verbose=verbose, nredo=10, spherical=False, gpu=True)
kmeans.train(x)

위에서 d는 vector의 dimension, ncentroids는 centroid의 개수, niter는 iteration 숫자, nredo는 random한 centroid initialization으로 몇번 반복할 것인지, spherical은 cosine distance을 쓸 것인지의 여부이다.

kmeans train이 끝나면 아래와 같이 최종 centroids를 얻어낼 수 있다. 새로운 데이터의 레이블을 알고 싶으면 아래 centroids를 database, 새로운 데이터를 query로 해서 위에서 설명한 k-nearest-neighbor search를 진행하면 끝이다.

kmeans.centroids 

여담: Cosine Similarity

위의 k nearest neighbors에서 cosine similarity로 하고 싶다면 (기본은 euclidean), 주어진 벡터를 l2 normalize한후 inner product로 search하면 된다.

import faiss                   

db_vector = ...
query_vector = ...

faiss.normalize_L2(db_vector)
faiss.normalize_L2(query_vector)

index = faiss.IndexFlatIP(db_vector)   # build the inner product index