faiss的简单使用
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 = 0fea_s_list = list()fea_v_list = list()key_d = dict()idx = 0with open(fname, "r", encoding="utf-8", errors="replace") as f:for row_content in f:line_num += 1try: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] = idxidx += 1fea_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] // 4cluster_num = int(fv_mat.shape[0]) // 14# metric_typemetric_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保存索引到文件
将索引保存到文件中，使用的时候能够快速获取和使用， 键值对能够更快的查找
# 将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_dwrite_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 = indexindex_refine.k_factor = 10index.nprobe = 100# 商品位置键值对keys = list()key_d = dict()with open("./vec/index_key.faiss") as fin:idx = 0for line in fin:line = line.strip("\r\n")key_d[line] = idxidx += 1keys.append(line)return index, index_refine, keys, key_d向量检索
向量检索必须是矩阵的形式传入。
# 向量检索
def search():# 推荐的结果集（商品ID，得分）res_ans = []# 检索最近邻的数量k = 10# 输入一个index库里的向量，进行检索index, index_refine, keys, key_d = load_index('./item_vec.txt')# 获取一个商品的idfs = '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()本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！