Bootstrap

faiss库中ivf-sq（ScalarQuantizer，标量量化）代码解读-3

测试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下进行调试，设置断点：
在这里插入图片描述
后面继续调试，查看后面的相对应的代码。

悦读

道可道，非常道；名可名，非常名。 无名，天地之始，有名，万物之母。 故常无欲，以观其妙，常有欲，以观其徼。 此两者，同出而异名，同谓之玄，玄之又玄，众妙之门。

最新收录
springboot集成工作流(成熟完整代码demo，可直接使用，运用在多个实际项目案例)
STM32之PWM
字符串常用的方法
数字IC/FPGA——复位篇
linux通过tftp下载的文件大小为0,linux 通过 tftp下载文件
【数据集9】DMSP/OLS夜间灯光数据 原创
单点登录
[日更-2019.4.15] -Ubuntu18编译内核源码并安装
2023-09-19-小米软开（C/C++）面经（二面）
vue学习日记Day5
;