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.")
Last modification:May 6th, 2021 at 02:31 pm