当前位置：首页 > article >正文

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

article 2024/12/29 6:16:32

流程

在这里插入图片描述

代码

void IndexIVF::search(
        idx_t n,
        const float* x,
        idx_t k,
        float* distances,
        idx_t* labels,
        const SearchParameters* params_in) const {
    FAISS_THROW_IF_NOT(k > 0);
    const IVFSearchParameters* params = nullptr;
    if (params_in) {
        params = dynamic_cast<const IVFSearchParameters*>(params_in);
        FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
    }
    const size_t nprobe =
            std::min(nlist, params ? params->nprobe : this->nprobe);
    FAISS_THROW_IF_NOT(nprobe > 0);

    // search function for a subset of queries
    auto sub_search_func = [this, k, nprobe, params](
                                   idx_t n,
                                   const float* x,
                                   float* distances,
                                   idx_t* labels,
                                   IndexIVFStats* ivf_stats) {
        std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
        std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);

        double t0 = getmillisecs();
        quantizer->search(
                n,
                x,
                nprobe,
                coarse_dis.get(),
                idx.get(),
                params ? params->quantizer_params : nullptr);

        double t1 = getmillisecs();
        invlists->prefetch_lists(idx.get(), n * nprobe);

        search_preassigned(
                n,
                x,
                k,
                idx.get(),
                coarse_dis.get(),
                distances,
                labels,
                false,
                params,
                ivf_stats);
        double t2 = getmillisecs();
        ivf_stats->quantization_time += t1 - t0;
        ivf_stats->search_time += t2 - t0;
    };

    if ((parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT) == 0) {
        int nt = std::min(omp_get_max_threads(), int(n));
        std::vector<IndexIVFStats> stats(nt);
        std::mutex exception_mutex;
        std::string exception_string;

#pragma omp parallel for if (nt > 1)
        for (idx_t slice = 0; slice < nt; slice++) {
            IndexIVFStats local_stats;
            idx_t i0 = n * slice / nt;
            idx_t i1 = n * (slice + 1) / nt;
            if (i1 > i0) {
                try {
                    sub_search_func(
                            i1 - i0,
                            x + i0 * d,
                            distances + i0 * k,
                            labels + i0 * k,
                            &stats[slice]);
                } catch (const std::exception& e) {
                    std::lock_guard<std::mutex> lock(exception_mutex);
                    exception_string = e.what();
                }
            }
        }

        if (!exception_string.empty()) {
            FAISS_THROW_MSG(exception_string.c_str());
        }

        // collect stats
        for (idx_t slice = 0; slice < nt; slice++) {
            indexIVF_stats.add(stats[slice]);
        }
    } else {
        // handle paralellization at level below (or don't run in parallel at
        // all)
        sub_search_func(n, x, distances, labels, &indexIVF_stats);
    }
}

代码解析

IndexIVF::search 函数是 FAISS 的 IndexIVF 类中实现的一个核心函数，用于在倒排文件（Inverted File List, IVF）索引中执行搜索操作。以下是对函数的详细解析：

函数功能

在倒排文件索引中搜索最近的 k 个向量，返回它们的距离和对应的索引。
支持多线程并行化以提高查询性能。

参数说明

void IndexIVF::search(
    idx_t n,                    // 查询向量的数量
    const float* x,             // 查询向量（每个向量有 d 个维度）
    idx_t k,                    // 每个查询向量要找到的最近邻个数
    float* distances,           // 输出的距离数组，大小为 n*k
    idx_t* labels,              // 输出的索引数组，大小为 n*k
    const SearchParameters* params_in // 搜索参数，可选
) const;

n：查询向量的数量。
x：指向查询向量的指针，形状为 (n, d)。
k：每个查询向量需要返回的最近邻数量。
distances：保存结果的距离数组。
labels：保存结果的索引数组。
params_in：可选的搜索参数对象，通常包括 nprobe（控制搜索的倒排列表数量）等。

函数实现解析

参数验证

FAISS_THROW_IF_NOT(k > 0);

确保 k > 0，即需要找到至少一个最近邻。
2. 处理搜索参数

const IVFSearchParameters* params = nullptr;
if (params_in) {
    params = dynamic_cast<const IVFSearchParameters*>(params_in);
    FAISS_THROW_IF_NOT_MSG(params, "IndexIVF params have incorrect type");
}
const size_t nprobe = std::min(nlist, params ? params->nprobe : this->nprobe);
FAISS_THROW_IF_NOT(nprobe > 0);

检查输入的搜索参数 params_in 是否是 IVFSearchParameters 类型。
从参数中提取 nprobe，即查询时访问的倒排列表数量：
- 如果参数提供了 nprobe，则使用参数中的值。
- 如果未提供，则使用索引默认的 nprobe。
确保 nprobe > 0。

定义子搜索函数

auto sub_search_func = [this, k, nprobe, params](
    idx_t n,
    const float* x,
    float* distances,
    idx_t* labels,
    IndexIVFStats* ivf_stats) {
    ...
};

定义一个局部 lambda 函数 sub_search_func，处理子查询任务。参数包括当前的查询向量、结果存储位置和统计信息。
内部实现的步骤：
量化查询向量：

quantizer->search(n, x, nprobe, coarse_dis.get(), idx.get(), params ? params->quantizer_params : nullptr);

使用量化器将查询向量分配到 nprobe 个倒排列表中。idx 保存分配的倒排列表索引。coarse_dis 保存量化后的距离。
倒排列表的预取：

invlists->prefetch_lists(idx.get(), n * nprobe);

预取倒排列表数据以提高内存访问性能。
实际搜索：

search_preassigned(n, x, k, idx.get(), coarse_dis.get(), distances, labels, false, params, ivf_stats);

在分配好的倒排列表中执行搜索，返回最近邻结果的距离和索引。
更新统计信息：

ivf_stats->quantization_time += t1 - t0;
ivf_stats->search_time += t2 - t0;

记录量化时间和搜索时间。
4. 选择并行模式

if ((parallel_mode & ~PARALLEL_MODE_NO_HEAP_INIT) == 0) {
    ...
} else {
    sub_search_func(n, x, distances, labels, &indexIVF_stats);
}

根据 parallel_mode 决定并行模式：如果启用了并行模式，则使用 OpenMP 进行多线程查询。否则直接调用 sub_search_func 处理整个查询。
5. 并行查询

int nt = std::min(omp_get_max_threads(), int(n));
std::vector<IndexIVFStats> stats(nt);
std::mutex exception_mutex;
std::string exception_string;

#pragma omp parallel for if (nt > 1)
for (idx_t slice = 0; slice < nt; slice++) {
    ...
}