faiss库中ivf-sq(ScalarQuantizer,标量量化)代码解读-3

发布于:2024-11-29 ⋅ 阅读:(26) ⋅ 点赞:(0)

测试3-IVFPQ.cpp

想要测试IVFSQ的命令,需要在faiss里面进行修改来帮助我理解IVFSQ的流程,目的是去理解IVFSQ如何构建索引、训练索引、添加数据,以及执行搜索操作。

编辑命令

修改了3-IVFPQ.cpp文件里面的内容,但是不需要动/tutorial/cpp/CMakeList.txt里面的内容,只需要转到faiss-1.7.4目录下,运行如下命令:

make -C build 3-IVFPQ && ./build/tutorial/cpp/3-IVFPQ

会在/faiss-1.7.4/build/tutorial/cpp/下生成3-IVFPQ,之后直接运行gdb ./3-IVFPQ
如果在/tutorial/cpp/添加了3-IVFPQ-l2.cpp和修改了/tutorial/cpp/CMakeList.txt,那么需要运行下面的命令:

faiss@node68:~/faiss-1.7.4$ make -C build test
faiss@node68:~/faiss-1.7.4$ make -C build 3-IVFPQ-l2 && ./build/demos/3-IVFPQ-l2

查看3-IVFPQ.cpp原有命令

查看3-IVFPQ.cpp原有命令,然后在此基础上进行修改,使其可以用来测试IVFSQ。最后使用

make -C build 3-IVFPQ && ./build/tutorial/cpp/3-IVFPQ

来生成可执行文件/build/tutorial/cpp/3-IVFPQ,并运行gdb ./3-IVFPQ
接下来查看相关的代码内容:

#include <cstdio>
#include <cstdlib>
#include <random>

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

using idx_t = faiss::idx_t; //将faiss::idx_t起个别名为idx_,本质上FAISS 定义 idx_t 为 64 位整数

int main() {
    int d = 64;      // dimension
    int nb = 100000; // database size
    int nq = 10000;  // nb of queries

    std::mt19937 rng;
    std::uniform_real_distribution<> distrib;

    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] = distrib(rng);
        xb[d * i] += i / 1000.;
    }

    for (int i = 0; i < nq; i++) {
        for (int j = 0; j < d; j++)
            xq[d * i + j] = distrib(rng);
        xq[d * i] += i / 1000.;
    }

    int nlist = 100;
    int k = 4;
    int m = 8;                       // bytes per vector
    faiss::IndexFlatL2 quantizer(d); // the other index
    faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8);

    index.train(nb, xb);
    index.add(nb, xb);

    { // sanity check
        idx_t* I = new idx_t[k * 5];
        float* D = new float[k * 5];

        index.search(5, xb, k, D, I);

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

        printf("D=\n");
        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < k; j++)
                printf("%7g ", D[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

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

        index.nprobe = 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("%5zd ", I[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

    delete[] xb;
    delete[] xq;

    return 0;
}

代码解析:

  1. 初始化数据库和查询数据
int d = 64;      // dimension
int nb = 100000; // database size
int nq = 10000;  // nb of queries
  • d:向量的维度,这里每个向量是一个 64 维的浮点向量。
  • nb:数据库大小,有 100,000 个数据向量。
  • nq:查询向量的数量,有 10,000 个查询向量。
  1. 生成随机向量
std::mt19937 rng;
std::uniform_real_distribution<> distrib;
  • 使用 C++ 的随机数生成器创建数据库向量和查询向量。
    数据库向量生成
float* xb = new float[d * nb];
for (int i = 0; i < nb; i++) {
    for (int j = 0; j < d; j++)
        xb[d * i + j] = distrib(rng);
    xb[d * i] += i / 1000.;
}
  • 通过 xb 数组存储所有数据库向量。
  • 向量的每个维度随机生成一个值,并对第一个维度稍作偏移,增加一定规律性(xb[d * i] += i / 1000.)。
    查询向量生成
float* xq = new float[d * nq];
for (int i = 0; i < nq; i++) {
    for (int j = 0; j < d; j++)
        xq[d * i + j] = distrib(rng);
    xq[d * i] += i / 1000.;
}
  • 通过 xq 数组存储所有查询向量。
  • 与数据库向量类似,也对第一个维度稍作偏移,增加一定规律性。
  1. 初始化和构建索引
int nlist = 100;
int k = 4;
int m = 8;                       // bytes per vector
faiss::IndexFlatL2 quantizer(d); // the other index
faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8);

IVFPQ 索引构建

  • nlist:倒排表(Inverted List)的数量,分桶数设置为 100。
  • m:Product Quantization (PQ) 的分块数,每个向量分为 m=8 个子向量。
  • faiss::IndexFlatL2:
    • 用于量化(quantization)的粗排索引,这里使用的是 L2 距离,实际上就是提供一个计算两个d维向量计算距离的公式
    • IndexFlatL2 是一个简单的暴力检索索引。
  • faiss::IndexIVFPQ:
    • 使用倒排文件和 PQ 量化的索引,创建一个数据结构IndexIVFPQ名为index的变量。
    • 参数 8 表示每个子向量用 8 位进行编码。
  1. 训练和添加数据
index.train(nb, xb);
index.add(nb, xb);
  • train:对数据进行训练,这一步会学习数据库向量的分布以构建量化器(quantizer)。
  • add:将训练后的向量添加到索引中进行存储。
  1. 检索示例
    (1)对数据库自身进行检索
idx_t* I = new idx_t[k * 5];
float* D = new float[k * 5];

index.search(5, xb, k, D, I);
  • 检索前 5 个数据库向量(xb)。
  • 每次检索返回 k=4 个最近邻结果。
  • 结果输出:
    • I:存储检索到的向量索引。
    • D:存储检索到的向量距离。
  • 打印结果,显示 5 个查询向量的最近邻索引和对应的距离。
    (2)对查询向量进行检索
idx_t* I = new idx_t[k * nq];
float* D = new float[k * nq];

index.nprobe = 10;
index.search(nq, xq, k, D, I);
  • 检索所有查询向量(xq)。
  • 每次检索返回 k=4 个最近邻结果。
  • 设置 index.nprobe = 10:
    • nprobe 表示倒排列表的探查数目(即搜索时访问的桶的数量)。
    • 增大 nprobe 会提升召回率,但检索速度会降低。
  • 打印最后 5 个查询向量的检索结果。
  1. 释放内存
delete[] xb;
delete[] xq;
delete[] I;
delete[] D;
  • 程序结束前清理动态分配的内存,避免内存泄漏。

总结

这段代码完整地演示了一个基于 FAISS 和 IVFPQ 的向量检索流程:

  1. 生成数据库向量和查询向量。
  2. 使用 FAISS 创建基于倒排文件和产品量化的索引。
  3. 训练索引并将数据库向量添加到索引中。
  4. 执行检索并输出结果。

使用该代码测试ivfSQ

更改该代码,使得该代码可以用来测试ivfSQ。
前面的生成数据的代码不用考虑只需要去考虑如下代码该如何替换:

int nlist = 100;
int k = 4;
int m = 8;                       // bytes per vector
faiss::IndexFlatL2 quantizer(d); // the other index
faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8);

index.train(nb, xb);
index.add(nb, xb);
  • 对于距离索引IndexFlatL2仍然选择使用quantizer,不过这里我修改名称为quantizer1.
  • 对于后面创建IndexIVFPQ数据结构,我们要改为IVFSQ的内容,也就是把这里修改为IndexIVFScalarQuantizer数据结构,并且把这个变量命名为index1,查看IndexIVFScalarQuantizer数据结构内容:
IndexIVFScalarQuantizer(
            Index* quantizer,
            size_t d,
            size_t nlist,
            ScalarQuantizer::QuantizerType qtype,
            MetricType metric = METRIC_L2,
            bool by_residual = true);

把相关内容赋值上去,这行代码改为:

faiss::IndexIVFScalarQuantizer index1(
            &quantizer1, d, nlist, faiss::ScalarQuantizer::QT_8bit);
  • 对于ivfSQ需要选择量化的范围,那么这里的范围参考范围选择
    代码为: index1.sq.rangestat = faiss::ScalarQuantizer::RS_meanstd;选择为平均值计算。
  • index1.train:标量量化(Scalar Quantization, SQ)准备数据的范围。相对应的add也是修改相对应的变量名。

总的修改的内容后面添加注释,表明进行了修改如下:

#include <cstdio>
#include <cstdlib>
#include <random>

#include <faiss/IndexFlat.h>
#include <faiss/ScalarQuantizer.h>  //修改

using idx_t = faiss::idx_t;

int main() {
    int d = 64;      // dimension
    int nb = 100000; // database size
    int nq = 10000;  // nb of queries

    std::mt19937 rng;
    std::uniform_real_distribution<> distrib;

    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] = distrib(rng);
        xb[d * i] += i / 1000.;
    }

    for (int i = 0; i < nq; i++) {
        for (int j = 0; j < d; j++)
            xq[d * i + j] = distrib(rng);
        xq[d * i] += i / 1000.;
    }

    int nlist = 100;
    int k = 4;
    int m = 8;                       // bytes per vector
    faiss::IndexFlatL2 quantizer1(d); // the other index //修改
    faiss::IndexIVFScalarQuantizer index1(               //修改
            &quantizer1, d, nlist, faiss::ScalarQuantizer::QT_8bit);  //修改
    index1.sq.rangestat = faiss::ScalarQuantizer::RS_meanstd;    //修改
    index1.train(nb, xb);                               //修改
    index1.add(nb, xb);                                 //修改

    { // sanity check
        idx_t* I = new idx_t[k * 5];
        float* D = new float[k * 5];

        index1.search(5, xb, k, D, I);                   //修改

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

        printf("D=\n");
        for (int i = 0; i < 5; i++) {
            for (int j = 0; j < k; j++)
                printf("%7g ", D[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

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

        //index.nprobe = 10;                   //修改
        index1.search(nq, xq, k, D, I);        //修改

        printf("ivfsq8 I=\n");                 //修改
        for (int i = nq - 5; i < nq; i++) {
            for (int j = 0; j < k; j++)
                printf("%5zd ", I[i * k + j]);
            printf("\n");
        }

        delete[] I;
        delete[] D;
    }

    delete[] xb;
    delete[] xq;

    return 0;
}

之后在faiss1.7.4的目录下,使用

make -C build 3-IVFPQ && ./build/tutorial/cpp/3-IVFPQ

会在/faiss-1.7.4/build/tutorial/cpp/下生成3-IVFPQ,之后直接运行gdb ./3-IVFPQ

然后在gdb下进行调试,设置断点:
在这里插入图片描述
后面继续调试,查看后面的相对应的代码。