Faiss全称(Facebook AI Similarity Search)是Facebook AI团队开源的针对聚类和相似性搜索库,为稠密向量提供高效相似度搜索和聚类,支持十亿级别向量的搜索,是目前较成熟的近似近邻搜索库。
它包含多种搜索任意大小向量集(备注:向量集大小由RAM内存决定)的算法,以及用于算法评估和参数调整的支持代码。
Faiss用C++编写,并提供与Numpy完美衔接的Python接口。除此以外,对一些核心算法提供了GPU实现。
faiss的安装
pip install faiss-cpu==1.7.3
# pip install faiss-gpu==1.7.2
上面的是使用cpu的版本,如果使用GPU可以安装下面语句
faiss的创建步骤
- 获取向量文件或者向量库。
- 创建faiss的索引index,(训练)将向量添加到索引index中。
- 使用faiss进行检索。
代码及目录结构 github
获取向量文件或者向量库
向量文件存储在item_vec.txt中, 格式是(商品ID,向量),具体结构是下图中展示。
读取向量文件,返回向量的matrix,由于faiss必须处理np.array的形式,所以进行转换,商品id的一个列表,商品id和序列号的键值对。
import faiss
import os
import sys
import numpy as np
# 从文件中读取向量
"""
fea_v_matrix: np.array形式的特征向量矩阵
fea_s_list: id的顺序列表
key_d: id和位置的对应dict
"""
def load_fea_from_file(fname):
line_num = 0
fea_s_list = list()
fea_v_list = list()
key_d = dict()
idx = 0
with open(fname, "r", encoding="utf-8", errors="replace") as f:
for row_content in f:
line_num += 1
try:
raw_content_list = row_content.strip().split()
fea_s = raw_content_list[0]
fea_v = np.array(raw_content_list[1:]).astype(np.float32)
except Exception as e:
sys.exit("fail to load file vector")
if fea_s not in key_d:
key_d[fea_s] = idx
idx += 1
fea_s_list.append(fea_s)
fea_v_list.append(fea_v)
fea_v_matrix = np.array(fea_v_list)
return fea_v_matrix, fea_s_list, key_d
创建索引
检索索引使用的方式可以查看文档,此处创建了一个精度索引refine,提高检索的精度。
def build_index_fun(path):
# path : './item_vec.txt'
fv_mat, fea_s_list, key_d = load_fea_from_file(path)
# cluster_num 是聚类的数量
M = fv_mat.shape[1] // 4
cluster_num = int(fv_mat.shape[0]) // 14
# metric_type
metric_type = faiss.METRIC_INNER_PRODUCT
# 创建索引
index = faiss.index_factory(fv_mat.shape[1], f"IVF{cluster_num},PQ{M}x4fs", metric_type)
# 正则化矩阵特征向量
faiss.normalize_L2(fv_mat)
# 训练索引
index.train(fv_mat)
# 索引添加矩阵向量
index.add(fv_mat)
# 创建一个精度索引, 它细化了距离计算,并根据这些距离对结果重新排序
index_refine = faiss.IndexRefineFlat(index, faiss.swig_ptr(fv_mat))
return index, index_refine, fea_s_list, key_d
保存索引到文件
将索引保存到文件中,使用的时候能够快速获取和使用, 键值对能够更快的查找文章来源:https://www.toymoban.com/news/detail-633690.html
# 将faiss的索引index写到文件中
def write_index(path):
# path = './item_vec.txt'
index, index_refine, keys, key_d = build_index_fun(path)
faiss.write_index(index, './vec/index_m.faiss')
faiss.write_index(index_refine, './vec/index_refine_m.faiss')
# 商品和序列的键值对也写到文件中
with open('./vec/index_key.faiss', "w") as fout:
for ky in keys:
fout.write(ky)
fout.write("\n")
return index, index_refine, keys, key_d
write_index('./item_vec.txt')
加载索引和设置查询参数
def load_index(index_dir):
# 加载索引文件
index = faiss.read_index('./vec/index_m.faiss')
index_refine = faiss.read_index('./vec/index_m.faiss')
# 设置基础索引
index_refine.base_index = index
index_refine.k_factor = 10
index.nprobe = 100
# 商品位置键值对
keys = list()
key_d = dict()
with open("./vec/index_key.faiss") as fin:
idx = 0
for line in fin:
line = line.strip("\r\n")
key_d[line] = idx
idx += 1
keys.append(line)
return index, index_refine, keys, key_d
向量检索
向量检索必须是矩阵的形式传入。文章来源地址https://www.toymoban.com/news/detail-633690.html
# 向量检索
def search():
# 推荐的结果集(商品ID,得分)
res_ans = []
# 检索最近邻的数量
k = 10
# 输入一个index库里的向量,进行检索
index, index_refine, keys, key_d = load_index('./item_vec.txt')
# 获取一个商品的id
fs = '653417384213'
i = key_d[fs]
# 根据位置获取
index_refine.make_direct_map()
fv = index_refine.reconstruct(i)
# 检索必须是矩阵的形式
matrix = np.expand_dims(fv, axis=0)
D, I = index_refine.search(matrix, k)
# 获取结果集
res_id_list, res_dis_list = I[0], D[0]
for i in range(len(res_id_list)):
res_ans.append((keys[res_id_list[i]], res_dis_list[i]))
print(res_ans)
# load_fea_from_file('./item_vec.txt')
# build_index_fun('./item_vec.txt')
write_index('./item_vec.txt')
search()
到了这里,关于faiss的简单使用的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!