faiss_svc.py

#!/usr/bin/env python

"""
    faiss_svc.py
"""

import faiss
import numpy as np
from time import time
from sklearn.svm import LinearSVC

class FaissLinearSVC:
    def __init__(self, model):
        self.coef_      = np.ascontiguousarray(model.coef_.astype(np.float32))
        self.intercept_ = model.intercept_[0]
    
    def decision_function(self, findex, k):
        scores, idx = findex.search(self.coef_, k=k)
        return scores.squeeze() + self.intercept_, idx.squeeze()


n_obs   = 800000
dim     = 128
n_model = 100
k       = 32

x = np.random.uniform(0, 1, (n_obs, dim))
x = np.ascontiguousarray(x).astype(np.float32)

model = LinearSVC(fit_intercept=True).fit(x[:n_model], np.random.choice((0, 1), n_model))

fmodel = FaissLinearSVC(model)
findex = faiss.IndexFlatIP(x.shape[1])
findex.add(x)

t = time()
top_faiss_scores, top_faiss_idx = fmodel.decision_function(findex, k)
top_faiss_scores, top_faiss_idx
faiss_time = time() - t

t = time()
svc_scores  = model.decision_function(x)
top_svc_idx = np.argsort(-svc_scores)[:k]
top_svc_scores = svc_scores[top_svc_idx]
top_svc_scores, top_svc_idx
svc_time = time() - t

print('faiss_time=%f | svc_time=%f' % (faiss_time, svc_time))
assert (top_faiss_idx == top_svc_idx).all()
assert np.allclose(top_faiss_scores, top_svc_scores)