openmp官方源码_Faiss 源码解析

Faiss 源码解析

faiss 是 facebook 开源的一个专门用于做高维向量的相似性搜索的库，有 c++ 和 python 的接口；目前项目地址在 https://github.com/facebookresearch/faiss。本文主要结合 faiss 的官方示例，介绍如何使用 faiss 以及 暴力/IVF/IVFPQ 检索算法在 faiss 的具体实现。

检索算法介绍

检索算法的介绍可以参考 科普，本文主要关注3种检索算法：

1. 暴力搜索：顾名思义，query 和 base 一一比对，选择最近的
2. IVF：首先在具有代表性的数据上训练聚类中心，然后将 base 加入到最近的聚类中心的桶里，在 search 的时候，query 先和聚类中心比对，再在一定数目的桶里做暴力搜索
3. IVFPQ：在 IVF 的基础上，将 base 做 PQ 量化，加速比对

faiss 的编译与安装

可以参考官方给出的编译方法，这里我没有安装 cuda，所以采用的命令是

./configure --without-cuda && make

在编译完 faiss 之后，我们对官方提供的示例也进行编译，路径在 ./tutorial/cpp 下，cd到目录下直接 make 就可以了

如何使用 faiss

官方总共提供了五个示例，其中有两个是 gpu 版本的，三个是 cpu 版本的，我们这里主要关注 cpu 的，分别是 1-Flat.cpp，2-IVFFLAT.cpp，3-IVFPQ.cpp，分别对应着暴力算法检索，IVF 算法检索，IVFPQ 算法检索。不同的算法在用户侧代码基本一致，我们选取 IVFPQ 做简单介绍。

#include <cstdio>
#include <cstdlib>

#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQ.h>


int main() {
    
    int d = 64;                            // 特征维度
    int nb = 100000;                       // base 样本数量
    int nq = 10000;                        // query 样本数量

    float *xb = new float[d * nb];
    float *xq = new float[d * nq];

    for(int i = 0; i < nb; i++) {
    
        for(int j = 0; j < d; j++)
            xb[d * i + j] = drand48();
        xb[d * i] += i / 1000.;
    } // 随机初始化 base 数据

    for(int i = 0; i < nq; i++) {
    
        for(int j = 0; j < d; j++)
            xq[d * i + j] = drand48();
        xq[d * i] += i / 1000.;
    }    // 随机初始化 query 数据


    int nlist = 100;  // 聚类中心个数
    int k = 4;
    int m = 8;                             // bytes per vector
    faiss::IndexFlatL2 quantizer(d);       // 初始化用 L2 暴力 search 的 index
    faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8); // 初始化 ivfpq 的 index，用 L2 暴力 search 的 index 初始化
    index.train(nb, xb); // 训练 index
    index.add(nb, xb); // 将 base 数据加入到 index 中，用于之后的搜索

    {
           // search xq
        long *I = new long[k * nq];
        float *D = new float[k * nq];

        index.nprobe = 10; // 搜索 10 个中心点
        index.search(nq, xq, k, D, I);

        printf("I=n");
        for(int i = nq - 5; i < nq; i++) {
    
            for(int j = 0; j < k; j++)
                printf("%5ld ", I[i * k + j]);
            printf("n");
        }

        delete [] I;
        delete [] D;
    }



    delete [] xb;
    delete [] xq;

    return 0;
}

这段代码主要包括了四个部分，分别是

1. 初始化 base/query 数据和 index
2. 训练 index
3. 加入base 到 index
4. 对 query 做 search

其中，使用 faiss 主要包含了三步。初始化数据准备不用多说，faiss 中要求的数据格式都是 n * d 的矩阵格式，然后被展平到一维 float 数组中。剩下的两步，都是对 index 进行操