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的创建步骤

1. 获取向量文件或者向量库。
2. 创建faiss的索引index，（训练）将向量添加到索引index中。
3. 使用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

保存索引到文件

将索引保存到文件中，使用的时候能够快速获取和使用， 键值对能够更快的查找

# 将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

向量检索

向量检索必须是矩阵的形式传入。

# 向量检索
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()