본문 바로가기

ML with SckitLearn

Faiss를 이용한 K-means구현 [사이킷런에 비해 8X 빠르고, 27X 적은 에러 구현]

반응형

페이스북에서 대용량 벡터(Vectoer)의 유사도(Similarity)검색을 위한 오픈소스 Faiss를 공개하였는데요. SIFT데이터셋을 사용하여 데이터가 10만개일 때는 쿼리 속도가 20ug로 종전의 최고 성능과 유사하지만 데이터가 10억개일 때는 17.7us로 종전 최고 성능보다 8배 정보 빠르게 구현할 수 있다고 합니다. 

 

아쉽게도 현재 Linux와 OSX(Mac)만을 지원하고 있습니다.

 

Faiss을 통해  K-Means를 빠르게 단 몇줄의 코딩으로 구현을 할 수 있게 됩니다. 

faiss library를 설치합니다.  faiss가 포함된 k-means는 일반적인 Scikit-Learn에 비해 매우 빠르게 실행이 됩니다.

import numpy as np
import faiss

class FaissKMeans:
    def __init__(self, n_clusters=10, n_init=100, max_iter=300):
        self.n_clusters = n_clusters
        self.n_init = n_init
        self.max_iter = max_iter
        self.kmeans = None
        self.cluster_centers_ = None
        self.inertia_ = None

    def fit(self, X, y):
        self.kmeans = faiss.Kmeans(d=X.shape[1],
                                   k=self.n_clusters,
                                   niter=self.max_iter,
                                   nredo=self.n_init)
        self.kmeans.train(X.astype(np.float32))
        self.cluster_centers_ = self.kmeans.centroids
        self.inertia_ = self.kmeans.obj[-1]

    def predict(self, X):
        return self.kmeans.index.search(X.astype(np.float32), 1)[1]
  • np.float32를 사용합니다. Faiss는 해당 type에서만 실행이 되니 프로그램시 해당 type만 사용하세요.
import faiss
import numpy as np

D = 128
N = 10000
K = 10  # The number of clusters
X = np.random.random((N, D)).astype(np.float32)

# Setup
kmeans = faiss.Kmeans(d=D, k=K, niter=20, verbose=True)
# For GPU(s), run the following line. This will use all GPUs
# kmeans = faiss.Kmeans(d=D, k=K, niter=20, verbose=True, gpu=True)

# Run clustering
kmeans.train(X)

# Error for each iteration
print(kmeans.obj)  # array with 20 elements

# Centroids after clustering
print(kmeans.centroids.shape)  # (10, 128)

# The assignment for each vector.
dists, ids = kmeans.index.search(X, 1)  # Need to run NN search again
print(ids.shape)  # (10000, 1)

# Params
print("D:", kmeans.d)
print("K:", kmeans.k)
print("niter:", kmeans.cp.niter)

 

반응형