Faiss和Rapidsai_Raft使用记录

背景

最近在做基于图的近似向量检索的实验,需要用到Faiss库和Rapids系列的Raft库,同时由于要统计一些算法内部的数据,因此不能直接使用它们预编译的Python库,而要手动从源码编译并通过C++调用,这里记录一下编译运行时遇到的一些问题和技巧,其中Raft的坑尤其多。

Faiss

编译

Faiss库编译比较简单,按照官网的教程即可。不过使用默认cmake开关编译出来的效率比较低,如果要和预编译的Python库对齐的话需要手动指定cmake开关。首先需要安装Intel的矩阵库MKL,按照这个网站的操作来即可,因为系统是Ubuntu,所以我选择的是APT安装方式,按照网站添加了源然后apt install intel-oneapi-mkl-devel。然后我使用的cmake开关是:

cmake -DFAISS_ENABLE_GPU=OFF -DFAISS_ENABLE_PYTHON=OFF -DBUILD_TESTING=OFF -DBUILD_SHARED_LIBS=OFF -DCMAKE_BUILD_TYPE=Release -DFAISS_OPT_LEVEL=avx512 -DBLA_VENDOR=Intel10_64_dyn -DMKL_LIBRARIES=/path/to/mkl/libs

前四个和效率没什么关系,看你需不需要编译GPU、Python库、测试代码和动态库。-DFAISS_OPT_LEVEL=avx512用来指定距离计算时用什么级别的向量化指令,需要根据你的CPU架构而定,我用的服务器支持AVX512,可以通过cat /proc/cpuinfo查看你的CPU支持情况,一般AVX2都是支持的,所以可以设置为``-DFAISS_OPT_LEVEL=avx2`。之后正常make就可以了,没什么坑。

使用

使用的时候,使用如下的编译指令:

g++ run.cpp -O3 \
-I ../faiss -L ../faiss/build/faiss -fopenmp -lfaiss_avx512 \
-m64  -lmkl_intel_lp64 -lmkl_gnu_thread -lmkl_core -lgomp -lpthread -lm -ldl \
-o run

其中,-I跟的是faiss源码所在目录,-L跟的是编译出来的libfaiss.a所在目录,如果你前面make install了就不用加这两个参数。-lfaiss_avx512和上面编译faiss时的cmake开关是对应的,如果选择了AVX512,这里就可以链接libfaiss_avx512.a,如果选择了AVX2,则得链接libfaiss_avx2.a,即使用参数-lfaiss_avx2。当然,不管选择哪个cmake开关,都会默认编译出一个没用AVX的版本libfaiss.a,所以使用-lfaiss也可以,尽管效率会比较低。

第三行的是Intel MKL的编译参数,以上只是我个人的选项,建议通过这个网站自动生成。注意两点,一个是Select interface layer这一栏,一定要选C API with 32-bit integer,不然可能会报错;另一个是Select OpenMP library根据你用的OpenMP库来选,如果你装过libiomp就可以选Intel的OpenMP库,不过我比较过它和GNU的OpenMP在我用的Faiss算法上效率差不多,所以用哪个应该都无所谓。生成完了用Use this link line里的内容替换上面的第三行。

技巧

这里介绍三个技巧,前两个针对预编译的Python库,第三个针对C++库:

  1. Python可以使用contrib子目录下的内容,里面包含了一些实用函数,比如vecs_io.py里的函数可以很方便的读写SIFT1M/GIST1M数据集的文件格式;torch_utils.py则可以自动包装faiss里的常见索引让它们能直接接受pytorch张量作为参数,GPU张量也可以直接传进去,节省了传到Host端的numpy数组再传到显存的开销。

  2. Faiss内有的属性或者方法返回值没有提供Python包装,比如返回一个vector或者指针,这就可以用faiss/python文件下的几个文件解决,具体而言,可以使用array_conversions.py提供的vector_to_arraycopy_array_to_vector进行转换;而swigfaiss.swig则定义了不少更底层的操作,比如swig_ptr可以将numpy数组转换成C指针,memcpy可以在Python代码里对不同C指针指向的数据进行倒腾,omp_set_num_threads可以设置faiss内部的线程数。这些在Python代码里都用faiss.XXX()调用即可。

  3. 有时需要统计Faiss里图算法的距离计算次数,直接改库的代码比较麻烦,可以通过继承距离计算类和Flat类,然后传给图算法的构造函数,距离计算次数用全局变量来存储,具体代码如下(以HNSW为例):

#include <faiss/IndexFlat.h>
#include <faiss/IndexHNSW.h>
#include <faiss/utils/distances.h>

long long total_dist;

// 和faiss/IndexFlat.cpp里的实现基本一致,仅仅删去了一些没用到的代码并添加了total_dist的统计
struct FlatL2Dis : FlatCodesDistanceComputer {
    size_t d;
    idx_t nb;
    const float* q;
    const float* b;

    float distance_to_code(const uint8_t* code) final {
        total_dist++;
        return fvec_L2sqr(q, (float*)code, d);
    }
    float symmetric_dis(idx_t i, idx_t j) override {
        total_dist++;
        return fvec_L2sqr(b + j * d, b + i * d, d);
    }
    explicit FlatL2Dis(const IndexFlat& storage, const float* q = nullptr):
        FlatCodesDistanceComputer(storage.codes.data(), storage.code_size),
        d(storage.d), nb(storage.ntotal), q(q), b(storage.get_xb()) {}
    void set_query(const float* x) override {
        q = x;
    }
};

// 因为IndexFlat定义在faiss/IndexFlat.h头文件里,所以可以直接继承IndexFlat,只重载get_FlatCodesDistanceComputer
struct IndexFlatMy: IndexFlat {
    explicit IndexFlatMy(idx_t d): IndexFlat(d, METRIC_L2) {}
    IndexFlatMy() {}
    FlatCodesDistanceComputer* get_FlatCodesDistanceComputer() const override {
        return new FlatL2Dis(*this);
    }
};

// ...

IndexFlatMy indexflatmy(d);
IndexHNSW index(&indexflatmy, 32);

total_dist = 0;
index.add(myn, traindata.data());
printf("%lld\n", total_dist);

Raft

Raft这个库,主要是为了使用其中GNND和CAGRA算法,坑多得不可理喻,我尝试了好几次才最终成功。

编译

编译就是一个大坑,因为直接按文档里的指示是编译不出来的😅。按照文档的指示,Raft库有三种用法:

  • 第一种是Python预编译库,这个没啥问题,但用不了GNND算法,CAGRA算法也必须整体使用,无法分别统计两阶段的时间,也无法调整更细致的参数;
  • 第二种是C++预编译库,通过Conda安装,然后去include Raft的Runtime API(cpp/include/raft_runtime下那些),但问题是这些Runtime API也要引用Raft本体API啊,所以我是理解不了这一种方案相比第三种方案到底有什么意义;
  • 第三种就是直接引用源码,由于这个库是header-only template library,所以理论上直接在代码里include想要的头文件然后想办法编译成功就行了。我就是这样想的,然后就掉进重重深坑里摔得狗血临头。关键在于这个库依赖非常多Rapids家族其他的库,而这些库又没有上APT源,因此很难解决盘根错结的依赖关系以及各种环境变量、路径的设置,然后它还用的是自己的CMake插件,下载了啥更是一团浆糊。在经历了Spdlog版本不对,未引用libcudacxx等错误之后,编译器最终死在了层层叠叠不知道是C++标准还是NVCC兼容性导致的一堆语法错误中🤪。

正确方法藏在这里。具体操作步骤如下:

  1. 强烈建议连外网,因为它会用CMake插件自动从Github抓取依赖库,Github以及它的子域名能否访问是个玄学,至少我这里不连外网总会卡在某个地方,手动改成镜像网址也不现实,至少我是看不懂它乱七八糟的配置路径,而且每次编译的时候都要保持联网并且连外网状态(挺离谱的)。
  2. 把整个项目克隆下来,然后将cpp/template这个文件夹单独复制出来,之后的操作都在这个文件夹下操作,所以克隆的项目可以删掉了。
  3. 按照你的需要修改src文件夹下的代码,修改CMakeLists.txt以添加链接库、增删要编译的源文件,然后./build.sh就可以编译了,编译时间比较长,因为需要抓取依赖库。编译后的程序在build文件夹下。

运行

运行也是一个大坑,这里直接给出我执行CAGRA两阶段的代码:

#include <vector>
#include <raft/core/copy.cuh>
#include <raft/core/mdspan.hpp>
#include <raft/neighbors/cagra.cuh>

using namespace raft::neighbors;

// 读取SIFT1M数据文件的函数
template<typename T> int vecs_read(const std::string &filename, std::vector<T> &out, long long cnt) {
    std::ifstream f(filename, std::ios::binary); if (!f) return -1;
    int dim; f.read((char *)&dim, 4); out.resize(cnt * dim); f.seekg(0, std::ios::beg);
    for (long long i = 0; i < cnt; i++) {
        f.seekg(4, std::ios::cur); f.read((char *)(out.data() + i * dim), dim * sizeof(T));
    }
    return dim;
}

int main(int argc, char** argv) {

    int myn = 1000000;
    std::vector<float> traindata;
    int d = vecs_read(argv[1], traindata, myn); printf("%d\n", d);
    raft::device_resources dev_resources;

    // 定义host矩阵,然后将读取的向量数据库传过去,再定义device矩阵,最后将host矩阵的内容复制到device矩阵
    auto host_dataset = raft::make_host_matrix<float, int64_t>(dev_resources, myn, d);
    for (int i = 0; i < myn; i++)
        for (int j = 0; j < d; j++)
            host_dataset(i, j) = traindata[i * d + j];
    auto dataset = raft::make_device_matrix<float, int64_t>(dev_resources, myn, d);
    raft::copy(dev_resources, dataset.view(), host_dataset.view());

    // 我用的是使用GNND构建CAGRA一阶段KNN图的方式
    experimental::nn_descent::index_params build_params;
    build_params.graph_degree = 64;
    auto knn_graph = raft::make_host_matrix<int64_t, int64_t>(myn, 64);
    cagra::build_knn_graph(dev_resources, raft::make_const_mdspan(dataset.view()), knn_graph.view(), build_params);

    // 二阶段优化
    auto optimized_graph = raft::make_host_matrix<int64_t, int64_t>(myn, 32);
    cagra::optimize(dev_resources, knn_graph.view(), optimized_graph.view());

    return 0;
}

有两个坑点,第一个是如何将读取的向量数据库传到设备中,我查了偌大的文档居然没一处提到。demo再这么说也得用一个比较实际的向量数据库作示例吧,竟然用随机生成的数据库。这个库还专门为随机生成向量数据库提供了一套API,而没有提供和vector还是传统数组交互的API,真的是无语😅😅😅……虽然也有可能是我才疏学浅没找到API,但就像我说的,要是真有这种API,正常就应该在最明显的地方(比如README)给个示例,最好是用SIFT1M这种常用的benchmark,然后把读取、训练、查询明明白白的展示出来。最后没办法,仿照这个写了个暴力复制。个人猜测host_matrixdevice_matrix应该是有API能够暴露指针的,然后通过memcpy或者cudaMemcpy倒腾数据,因为Python库可以直接用pytorch或者cupy的张量作接口,但我确实是没在C++这边找到。(更新:看了库里几个算法的实现,这个方法应该是.data_handle(),返回数据模板类型的指针,但我还没有测试)

第二个坑点是cagra::build_knn_graphcagra::optimize的模板类型需求,稍有不慎编译就过不了,通过cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh可以看到两个函数分别限制了knn_graphoptimized_graph的第二个模板参数必须是int64_t,同时限制了它们的第一个模板参数必须一致。

总结

以Faiss为代表的大部分库编译运行基本都没太大问题,cmake ..,缺啥补啥,然后make -j,使用时-I-L-l指定好路径和链接库就完事了;麻烦的还是想Raft这种的,集齐了header-only、cuda、cmake插件等诸多坑点,加上文档写的太垃圾,才会这么恶心。遇到这种麻烦库,绝招就是去找它目录下的testing、demo、template这种文件夹或者同开发者以这些名字命名的仓库,一般这种地方提供的cmakelists/makefile总是可用的。

posted @ 2024-06-30 12:58  YuanZiming  阅读(60)  评论(0编辑  收藏  举报