利用faiss进行相似性计算

最新推荐文章于 2025-01-22 00:38:07 发布

陨星落云

最新推荐文章于 2025-01-22 00:38:07 发布

阅读量856

点赞数

CC 4.0 BY-SA版权

本文链接：https://blog.csdn.net/qq_28368377/article/details/119647548

该博客介绍了如何利用Faiss库进行大规模的向量搜索。通过设置维度、向量数、聚类数等参数，训练并建立索引，然后执行搜索任务，找到最接近的邻居。实验结果显示了搜索得到的距离和对应的索引号。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

# -*- coding: utf-8 -*-
"""
Created on Thu Aug 12 14:47:54 2021

@author: xiaohuihui
"""
import imageio
from tqdm import tqdm
import numpy as np
import faiss


dimension = 16 # 维度
n = 200 # 向量数
# n = 4
np.random.seed(1)


db_vectors = np.random.random((n,dimension)).astype('float32')
# db_vectors = imageio.imread(r'D:\ShanDong\high_quality_dataset256\img\1.tif').reshape(256*256,3).astype('float32')
# db_vectors = np.array([[1,2,3],[1,2.5,3],[1,2,4],[1,2,3]]).astype('float32')
nlist = 5 # 聚类数

quantiser = faiss.IndexFlatL2(dimension)
index = faiss.IndexIVFFlat(quantiser,dimension,nlist,faiss.METRIC_L2)

print(index.is_trained) # 判断是否参与训练

index.train(db_vectors) # 用于训练的向量
print(index.ntotal) # 索引总数
index.add(db_vectors) # 增加向量并更新索引
print(index.is_trained)
print(index.ntotal)

n_query = 10
k = 3 # 返回3个邻居

np.random.seed(0)

query_vectors = np.random.random((n_query,dimension)).astype('float32')
# query_vectors = db_vectors
distances,indices = index.search(query_vectors,k)

print('distances:\n',distances)
print('indices:\n',indices)

结果：

False
0
True
200
distances:
 [[0.8966236  0.91420597 1.2566044 ]
 [1.2293743  1.287607   1.289156  ]
 [1.0859215  1.1417681  1.33444   ]
 [1.1885746  1.2039149  1.3462903 ]
 [1.2022905  1.6060164  1.6962758 ]
 [0.653887   1.1338001  1.1754503 ]
 [1.4977793  1.5412362  1.5997455 ]
 [0.8660071  0.90651    1.3378503 ]
 [0.8879874  0.9618496  1.1274483 ]
 [1.1947813  1.3962246  1.4737606 ]]
indices:
 [[ 16  81 185]
 [ 63  21  57]
 [154 136 146]
 [ 43 155 143]
 [ 36 143 120]
 [ 11  83   0]
 [126  46 180]
 [ 81  16 141]
 [ 83  73 119]
 [104 168   5]]