Faiss 安装:
#pip
pip install faiss
# GPU
pip install faiss-gpu
#conda
#CPU-only
conda install -c pytorch faiss-cpu
#GPU(+CPU) version
conda install -c pytorch faiss-gpu
#for a specific CUDA version
conda install -c pytorch faiss-gpu cudatoolkit=10.2 # for CUDA 10.2
实现如:
import faiss
# modified from https://github.com/facebookresearch/deepcluster
def get_knn(reference_embeddings, test_embeddings, k,
embeddings_come_from_same_source=False):
"""
Finds the k elements in reference_embeddings that are closest to each
element of test_embeddings.
Args:
reference_embeddings: numpy array of size (num_samples, dimensionality).
test_embeddings: numpy array of size (num_samples2, dimensionality).
k: int, number of nearest neighbors to find
embeddings_come_from_same_source: if True, then the nearest neighbor of
each element (which is actually itself)
will be ignored.
"""
d = reference_embeddings.shape[1]
print("[INFO]running k-nn with k=%d"%k)
print("[INFO]embedding dimensionality is %d"%d)
index = faiss.IndexFlatL2(d)
if faiss.get_num_gpus() > 0: #GPU
print("[INFO]faiss gpu")
index = faiss.index_cpu_to_all_gpus(index)
index.add(reference_embeddings) #
_, indices = index.search(test_embeddings, k + 1)
if embeddings_come_from_same_source:
return indices[:, 1:]
return indices[:, :k]
#
query, reference, query_labels, reference_labels = feats, feats, labels, labels
embeddings_come_from_same_source = True
#
knn_indices = get_knn(reference, query, num_k,
embeddings_come_from_same_source)
knn_labels = reference_labels[knn_indices]
print("[INFO]Done.")