本文解析了Faiss中IndexIVFPQ在CPU上的搜索过程，包括Python接口调用、量化器搜索、预分配搜索等关键步骤，并深入探讨了search_preassigned函数的工作原理。
1. 说明
之前分析过了faiss 在GPU中的search过程，这里分析一下IndexIVFPQ在CPU中的search过程，即不将index拷贝到GPU中。
2. 过程分析
2.1 python接口
CPU search的python接口与GPU的完全一致，没有差别。
D, I = gpu_index.search(xq_t[x],top_k)
2.2 faiss core
IndexIVF::search
因为IndexIVFPQ没有override search，所以在实际运行过程中会调用父类IndexIVF中的实现。
void IndexIVF::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels) const
{
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
// 1. quantizer中的粗量搜索
quantizer->search (n, x, nprobe, coarse_dis.get(), idx.get());
// 2. prefetch_lists函数为空
invlists->prefetch_lists (idx.get(), n * nprobe);
// 3. 预分配搜索
search_preassigned (n, x, k, idx.get(), coarse_dis.get(),
distances, labels, false);
}
search过程主要包含三个部分：
-
quantizer->search
量化器搜索聚类中心，该过程在quantizer实例中进行，计算所有(nlist个)聚类中心与每一条搜索向量的距离，并从中找出最近的nprobe个聚类中心，输出到idx和coarse_dis中。
所以这里的idx和coarse_dis是大小为 nnprobe 的二维数组，分别存放nnprobe个向量label和与原向量的距离。 -
invlists->prefetch_lists
这里没有用到这类功能，所以实际调用的该函数为空。 -
search_preassigned
经过第一步计算出粗聚类中心向量后，在这里进行二次计算，即在选定的聚类里再次计算出top_k个近邻向量。由于在add时已经对向量进行了预计算形成残差，所以这里只要进行向量和运算就可以了。
在实际测试时发现，search的主要时间消耗是在这里，占比达到96%以上，所以针对这一过程进行进一步分析。
IndexIVF::search_preassigned
/** search a set of vectors, that are pre-quantized by the IVF
* quantizer. Fill in the corresponding heaps with the query
* results. The default implementation uses InvertedListScanners
* to do the search.
*
* @param n nb of vectors to query
* @param x query vectors, size nx * d
* @param assign coarse quantization indices, size nx * nprobe
* @param centroid_dis
* distances to coarse centroids, size nx * nprobe
* @param distance
* output distances, size n * k
* @param labels output labels, size n * k
* @param store_pairs store inv list index + inv list offset
* instead in upper/lower 32 bit of result,
* instead of ids (used for reranking).
* @param params used to override the object's search parameters
*/
void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
const idx_t *keys,
const float *coarse_dis ,
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params) const
{
// max_codes是默认值0
long nprobe = params ? params->nprobe : this->nprobe;
long max_codes = params ? params->max_codes : this->max_codes;
size_t nlistv = 0, ndis = 0, nheap = 0;
// 根据计算类型定义堆
using HeapForIP = CMin<float, idx_t>;
using HeapForL2 = CMax<float, idx_t>;
bool interrupt = false;
// don't start parallel section if single query
bool do_parallel =
parallel_mode == 0 ? n > 1 :
parallel_mode == 1 ? nprobe > 1 :
nprobe * n > 1;
#pragma omp parallel if(do_parallel) reduction(+: nlistv, ndis, nheap)
{
InvertedListScanner *scanner = get_InvertedListScanner(store_pairs);
ScopeDeleter1<InvertedListScanner> del(scanner);
/****************************************************
* Actual loops, depending on parallel_mode
****************************************************/
if (parallel_mode == 0) {
#pragma omp for
for (size_t i = 0; i < n; i++) {
if (interrupt) {
continue;
}
// 在inverted_list中设置搜索的起始点
scanner->set_query (x + i * d);
// 根据i设置distances和labels的地址
float * simi = distances + i * k;
idx_t * idxi = labels + i * k;
init_result (simi, idxi);
long nscan = 0;
// 依次在nprobe个聚类中进行搜索
for (size_t ik = 0; ik < nprobe; ik++) {
nscan += scan_one_list (
keys [i * nprobe + ik],
coarse_dis[i * nprobe + ik],
simi, idxi
);
if (max_codes && nscan >= max_codes) {
break;
}
}
ndis += nscan;
//对搜索结果进行排序
reorder_result (simi, idxi);
if (InterruptCallback::is_interrupted ()) {
interrupt = true;
}
} // parallel for
} else if (parallel_mode == 1) {
std::vector <idx_t> local_idx (k);
std::vector <float> local_dis (k);
for (size_t i = 0; i < n; i++) {
scanner->set_query (x + i * d);
init_result (local_dis.data(), local_idx.data());
#pragma omp for schedule(dynamic)
for (size_t ik = 0; ik < nprobe; ik++) {
ndis += scan_one_list
(keys [i * nprobe + ik],
coarse_dis[i * nprobe + ik],
local_dis.data(), local_idx.data());
// can't do the test on max_codes
}
// merge thread-local results
float * simi = distances + i * k;
idx_t * idxi = labels + i * k;
#pragma omp single
init_result (simi, idxi);
#pragma omp barrier
// 将各个线程产生的堆合并到结果堆中，临界访问
#pragma omp critical
{
if (metric_type == METRIC_INNER_PRODUCT) {
heap_addn<HeapForIP>
(k, simi, idxi,
local_dis.data(), local_idx.data(), k);
} else {
heap_addn<HeapForL2>
(k, simi, idxi,
local_dis.data(), local_idx.data(), k);
}
}
#pragma omp barrier
#pragma omp single
reorder_result (simi, idxi);
}
} else {
FAISS_THROW_FMT ("parallel_mode %d not supported\n",
parallel_mode);
}
} // parallel section
if (interrupt) {
FAISS_THROW_MSG ("computation interrupted");
}
indexIVF_stats.nq += n;
indexIVF_stats.nlist += nlistv;
indexIVF_stats.ndis += ndis;
indexIVF_stats.nheap_updates += nheap;
}
这里的params使用默认值nullptr.
流程
从代码中可以看出，虽然根据parallel_mode值的不同，程序处理上会有部分差异，但差异主要体现在并行运算的时间点和内容，主要流程是一致的：
- 首先根据原向量下标确定要输出的distances和labels的地址；
- 在nprobe个倒序列表中进行搜索，并对搜索结果进行排序
堆的使用
Faiss使用堆对搜索结果进行排序，不同的搜索类型可能使用不同的堆，在L2的范式搜索中使用大根堆进行排序。
do_parallel
Faiss可以直接使用OpemMP的并行运算指令，do_parallel是打开并行运算的标志值，在以下三种情况下为1：
- parallel_mode == 0 && n > 1
- parallel_mode == 1 && nprobe > 1
- parallel_mode == 2 && nprobe * n > 1
parallel_mode
该值确定采用何种并行模式进行查询。0表示在查询时开启并行，1表示在inverted_list计算残差时开启并行，2表示在上述两个阶段都使用并行模式。
并行运算
Faiss默认添加了对OpenMP的支持，具体命令待添加
intialize + reorder resule heap
这部分原本是定义在search_preassigned 函数体内的函数，这里把它们单拎出来。顾名思义，这里是分别对堆进行初始化和排序，因为搜索结果通常是两个列表，distances和labels，所以这里也是两个堆。
auto init_result = [&](float *simi, idx_t *idxi) {
if (metric_type == METRIC_INNER_PRODUCT) {
heap_heapify<HeapForIP> (k, simi, idxi);
} else {
heap_heapify<HeapForL2> (k, simi, idxi);
}
};
auto reorder_result = [&] (float *simi, idx_t *idxi) {
if (metric_type == METRIC_INNER_PRODUCT) {
heap_reorder<HeapForIP> (k, simi, idxi);
} else {
heap_reorder<HeapForL2> (k, simi, idxi);
}
};
loop probes分析
每个线程在这一过程中会循环遍历nprobe个聚类，计算残差，以最后得出top个最近邻向量。代码内容如下：
// loop over probes
for (size_t ik = 0; ik < nprobe; ik++) {
nscan += scan_one_list (
keys [i * nprobe + ik],
coarse_dis[i * nprobe + ik],
simi, idxi
);
if (max_codes && nscan >= max_codes) {
break;
}
}
其中scan_one_list是函数内定义的函数，主要完成下列工作:
- invlist->list_size：计算倒序列表的大小，为空则直接跳过；
- scanner->set_list：根据key值，从coarse_dis_i列表中找到入口地址；
- ScopedCodes：根据key值和invlists的内容生成ccode；
- sids->get：重置并获取id号；
- scanner->scan_codes：扫描一组代码，计算到当前查询的距离，并更新结果堆。
scan_one_list
这个函数的功能是在单个的聚类中进行搜索。
/*
* key: nprobe中的invetred list编号
* coarse_dis_i: key对应的聚类中心的distance
* simi：存放distances结果的堆
* idxi： 存放labels结果的堆
*/
auto scan_one_list = [&] (idx_t key, float coarse_dis_i,
float *simi, idx_t *idxi) {
if (key < 0) {
// not enough centroids for multiprobe
return (size_t)0;
}
FAISS_THROW_IF_NOT_FMT (key < (idx_t) nlist,
"Invalid key=%ld nlist=%ld\n",
key, nlist);
size_t list_size = invlists->list_size(key);
// don't waste time on empty lists
if (list_size == 0) {
return (size_t)0;
}
scanner->set_list (key, coarse_dis_i);
nlistv++;
InvertedLists::ScopedCodes scodes (invlists, key);
std::unique_ptr<InvertedLists::ScopedIds> sids;
const Index::idx_t * ids = nullptr;
if (!store_pairs) {
sids.reset (new InvertedLists::ScopedIds (invlists, key));
ids = sids->get();
}
nheap += scanner->scan_codes (list_size, scodes.get(),
ids, simi, idxi, k);
return list_size;
}
这个函数的绝大部分时间消耗在scanner->scan_codes内。
scanner
从代码看，scanner是一个搜索引擎的实例，用于在InvertedList中进行搜索的具体实现。值得单独分析。
所有的线程都会单独生成一个自己的scanner。
InvertedListScanner *scanner = get_InvertedListScanner(store_pairs);
scanner->scan_codes
scanner是struct IVFPQScanner结构体的实例，定义在IndexIVFPQ.h中，scan_codes函数内容如下：
size_t scan_codes (size_t ncode,
const uint8_t *codes,
const idx_t *ids,
float *heap_sim, idx_t *heap_ids,
size_t k) const override
{
KnnSearchResults<C> res = {
/* key */ this->key,
/* ids */ this->store_pairs ? nullptr : ids,
/* k */ k,
/* heap_sim */ heap_sim,
/* heap_ids */ heap_ids,
/* nup */ 0
};
if (this->polysemous_ht > 0) {
assert(precompute_mode == 2);
this->scan_list_polysemous (ncode, codes, res);
} else if (precompute_mode == 2) {
this->scan_list_with_table (ncode, codes, res);
} else if (precompute_mode == 1) {
this->scan_list_with_pointer (ncode, codes, res);
} else if (precompute_mode == 0) {
this->scan_on_the_fly_dist (ncode, codes, res);
} else {
FAISS_THROW_MSG("bad precomp mode");
}
return res.nup;
}
实际运行中polysemous_ht为0，precompute_mode为2，故之后调用scan_list_with_table函数。
scan_list_with_table
函数内容如下：
template<class SearchResultType>
void scan_list_with_table (size_t ncode, const uint8_t *codes,
SearchResultType & res) const
{
for (size_t j = 0; j < ncode; j++) {
float dis = dis0;
const float *tab = sim_table;
for (size_t m = 0; m < pq.M; m++) {
dis += tab[*codes++];
tab += pq.ksub;
}
res.add(j, dis);
}
}
程序运行到这里已经进入了IndexIVFPQ的聚类里面进行残差计算（浮点加）。
3. 总结
虽然不够细致，但经过本文的梳理，可以大致看出IndexIVF的搜索的过程，这其中最重要的两个步骤分别与上一篇文档中提到的Product Quantizer和Inverted File System对应，所以说只要搞清楚这两个实例的过程，便能完全了解整个search的流程了。
扫一扫
3453
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
