image.png
<meta charset="utf-8">
Faiss 源碼解析
faiss 是 facebook 開源的一個專門用于做高維向量的相似性搜索的庫鸭廷，有 c++ 和 python 的接口胳蛮；目前項目地址在 https://github.com/facebookresearch/faiss。本文主要結(jié)合 faiss 的官方示例，介紹如何使用 faiss 以及 暴力/IVF/IVFPQ 檢索算法在 faiss 的具體實現(xiàn)。
檢索算法介紹
檢索算法的介紹可以參考 科普个盆，本文主要關(guān)注3種檢索算法：
- 暴力搜索：顧名思義，
query和base一一比對朵栖，選擇最近的 -
IVF：首先在具有代表性的數(shù)據(jù)上訓(xùn)練聚類中心颊亮，然后將base加入到最近的聚類中心的桶里，在search的時候陨溅，query先和聚類中心比對终惑，再在一定數(shù)目的桶里做暴力搜索 -
IVFPQ：在IVF的基礎(chǔ)上，將base做PQ量化门扇，加速比對
faiss 的編譯與安裝
可以參考官方給出的編譯方法狠鸳，這里我沒有安裝 cuda，所以采用的命令是
./configure --without-cuda && make
在編譯完 faiss 之后悯嗓，我們對官方提供的示例也進行編譯，路徑在 ./tutorial/cpp 下卸察，cd到目錄下直接 make 就可以了
如何使用 faiss
官方總共提供了五個示例脯厨，其中有兩個是 gpu 版本的，三個是 cpu 版本的坑质，我們這里主要關(guān)注 cpu 的合武，分別是 1-Flat.cpp，2-IVFFLAT.cpp涡扼，3-IVFPQ.cpp稼跳，分別對應(yīng)著暴力算法檢索，IVF 算法檢索吃沪，IVFPQ 算法檢索汤善。不同的算法在用戶側(cè)代碼基本一致，我們選取 IVFPQ 做簡單介紹。
#include <cstdio>
#include <cstdlib>
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFPQ.h>
int main() {
int d = 64; // 特征維度
int nb = 100000; // base 樣本數(shù)量
int nq = 10000; // query 樣本數(shù)量
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] = drand48();
xb[d * i] += i / 1000.;
} // 隨機初始化 base 數(shù)據(jù)
for(int i = 0; i < nq; i++) {
for(int j = 0; j < d; j++)
xq[d * i + j] = drand48();
xq[d * i] += i / 1000.;
} // 隨機初始化 query 數(shù)據(jù)
int nlist = 100; // 聚類中心個數(shù)
int k = 4;
int m = 8; // bytes per vector
faiss::IndexFlatL2 quantizer(d); // 初始化用 L2 暴力 search 的 index
faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8); // 初始化 ivfpq 的 index红淡，用 L2 暴力 search 的 index 初始化
index.train(nb, xb); // 訓(xùn)練 index
index.add(nb, xb); // 將 base 數(shù)據(jù)加入到 index 中不狮，用于之后的搜索
{ // search xq
long *I = new long[k * nq];
float *D = new float[k * nq];
index.nprobe = 10; // 搜索 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("%5ld ", I[i * k + j]);
printf("\n");
}
delete [] I;
delete [] D;
}
delete [] xb;
delete [] xq;
return 0;
}
這段代碼主要包括了四個部分，分別是
- 初始化
base/query數(shù)據(jù)和index - 訓(xùn)練
index - 加入
base到index - 對
query做search
其中在旱，使用 faiss 主要包含了三步摇零。初始化數(shù)據(jù)準(zhǔn)備不用多說，faiss 中要求的數(shù)據(jù)格式都是 n * d 的矩陣格式桶蝎，然后被展平到一維 float 數(shù)組中驻仅。剩下的兩步，都是對 index 進行操作登渣。
源碼解析
檢索流程
參考官方給的例子噪服，檢索分為三步：train，add绍豁，search芯咧，不同的檢索算法，體現(xiàn)在使用不同的 index 進行這三步上
-
train：選取有代表性的數(shù)據(jù)竹揍，訓(xùn)練index -
add：將base數(shù)據(jù)加入到index中 -
search：對于給定的query敬飒，返回其對應(yīng)的在底庫中的topk
重要類
Index
index 的基類，后續(xù)各種各樣的檢索算法芬位，都會繼承這個基類或者這個類的派生類无拗，然后實現(xiàn)具體的方法，在這個類中昧碉，有如下的數(shù)據(jù)成員：
-
d：維度英染，每個向量的維度 -
ntotal：索引的向量的數(shù)目，可以理解成檢索時的base數(shù)目 -
metric_type：檢索時使用的metric類型被饿，比如L2四康，內(nèi)積等
IndexFlat
用于做暴力搜索的 index 類，直接繼承 index狭握。暴力搜索思路很簡單孽江，無需 train滑频，add 的所有 base 都被存儲起來，然后在 search 的時候把 query 和所有 base 進行比對，選取最近的似袁。我們看下具體實現(xiàn)凳宙。
add
add 就是把所有的 base 都存儲起來
void IndexFlat::add (idx_t n, const float *x) {
xb.insert(xb.end(), x, x + n * d);
ntotal += n;
}
Search
Search 的時候殴胧，根據(jù) metric type 的不同荚坞，返回 query 的 topk。具體計算時采用了 openmp 和sse/avx 優(yōu)化
void IndexFlat::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels) const
{
// we see the distances and labels as heaps
if (metric_type == METRIC_INNER_PRODUCT) {
float_minheap_array_t res = {
size_t(n), size_t(k), labels, distances};
knn_inner_product (x, xb.data(), d, n, ntotal, &res); //函數(shù)內(nèi)部有并行優(yōu)化
} else if (metric_type == METRIC_L2) {
float_maxheap_array_t res = {
size_t(n), size_t(k), labels, distances};
knn_L2sqr (x, xb.data(), d, n, ntotal, &res);
} else {
float_maxheap_array_t res = {
size_t(n), size_t(k), labels, distances};
knn_extra_metrics (x, xb.data(), d, n, ntotal,
metric_type, metric_arg,
&res);
}
}
Clustering
實現(xiàn) K-means 聚類的類今妄，提供train 郑口，需要訓(xùn)練數(shù)據(jù)和 index（用于 search 最近的向量）鸳碧，結(jié)果得到訓(xùn)練數(shù)據(jù)的類中心向量，如果是量化的向量潘酗，那么還需要提供量化使用的 index codec杆兵，我們?nèi)コ炕牟糠郑豢?float 數(shù)據(jù)
核心代碼如下仔夺，包括如下部分：
-
search過程琐脏，將聚類中心作為底庫加入到index中，并對訓(xùn)練數(shù)據(jù)做search缸兔，得到assign -
計算新的聚類中心日裙，計算新的聚類中心的代碼在
compute_centroids中，具體就是對于相同的類別的向量惰蜜，將向量的均值作為新的中心昂拂，在實現(xiàn)上，利用 openmp 進行了并行優(yōu)化
重復(fù)以上兩步抛猖，就可以得到最優(yōu)的聚類中心
void Clustering::train_encoded (idx_t nx, const uint8_t *x_in,
const Index * codec, Index & index,
const float *weights) {
// 前處理省略
for (int redo = 0; redo < nredo; redo++) {
if (verbose && nredo > 1) {
printf("Outer iteration %d / %d\n", redo, nredo);
}
// initialize (remaining) centroids with random points from the dataset
centroids.resize (d * k);
std::vector<int> perm (nx);
rand_perm (perm.data(), nx, seed + 1 + redo * 15486557L);
for (int i = n_input_centroids; i < k ; i++) {
memcpy (¢roids[i * d], x + perm[i] * line_size, line_size);
}
post_process_centroids ();
// prepare the index
if (index.ntotal != 0) {
index.reset();
}
index.add (k, centroids.data());
// k-means iterations
float err = 0;
for (int i = 0; i < niter; i++) {
double t0s = getmillisecs();
index.search (nx, reinterpret_cast<const float *>(x), 1,
dis.get(), assign.get());
InterruptCallback::check();
t_search_tot += getmillisecs() - t0s;
// accumulate error
err = 0;
for (int j = 0; j < nx; j++) {
err += dis[j];
}
// update the centroids
std::vector<float> hassign (k);
size_t k_frozen = frozen_centroids ? n_input_centroids : 0;
compute_centroids (
d, k, nx, k_frozen,
x, codec, assign.get(), weights,
hassign.data(), centroids.data()
);
index.reset ();
if (update_index) {
index.train (k, centroids.data());
}
index.add (k, centroids.data());
InterruptCallback::check ();
}
}
//保存最優(yōu)聚類中心
if (nredo > 1) {
centroids = best_centroids;
iteration_stats = best_obj;
index.reset();
index.add(k, best_centroids.data());
}
}
void compute_centroids (size_t d, size_t k, size_t n,
size_t k_frozen,
const uint8_t * x, const Index *codec,
const int64_t * assign,
const float * weights,
float * hassign,
float * centroids)
{
k -= k_frozen;
centroids += k_frozen * d;
memset (centroids, 0, sizeof(*centroids) * d * k);
size_t line_size = codec ? codec->sa_code_size() : d * sizeof (float);
#pragma omp parallel
{
int nt = omp_get_num_threads();
int rank = omp_get_thread_num();
// this thread is taking care of centroids c0:c1
size_t c0 = (k * rank) / nt;
size_t c1 = (k * (rank + 1)) / nt;
std::vector<float> decode_buffer (d);
for (size_t i = 0; i < n; i++) {
int64_t ci = assign[i];
assert (ci >= 0 && ci < k + k_frozen);
ci -= k_frozen;
if (ci >= c0 && ci < c1) {
float * c = centroids + ci * d;
const float * xi;
if (!codec) {
xi = reinterpret_cast<const float*>(x + i * line_size);
} else {
float *xif = decode_buffer.data();
codec->sa_decode (1, x + i * line_size, xif);
xi = xif;
}
if (weights) {
float w = weights[i];
hassign[ci] += w;
for (size_t j = 0; j < d; j++) {
c[j] += xi[j] * w;
}
} else {
hassign[ci] += 1.0;
for (size_t j = 0; j < d; j++) {
c[j] += xi[j];
}
}
}
}
}
#pragma omp parallel for
for (size_t ci = 0; ci < k; ci++) {
if (hassign[ci] == 0) {
continue;
}
float norm = 1 / hassign[ci];
float * c = centroids + ci * d;
for (size_t j = 0; j < d; j++) {
c[j] *= norm;
}
}
}
IndexIVF
用于做 IVF 搜索的 index 類格侯。
train
ivf 算法會把給定的數(shù)據(jù)進行聚類，得到固定數(shù)目的聚類中心财著。具體的联四，就是 train_q1? 的過程，train_residual 在 ivf 中是一個空函數(shù)
void IndexIVF::train (idx_t n, const float *x)
{
train_q1 (n, x, verbose, metric_type);
train_residual (n, x);
is_trained = true;
}
void IndexIVF::train_residual(idx_t /*n*/, const float* /*x*/) {
if (verbose)
printf("IndexIVF: no residual training\n");
// does nothing by default
}
train_q1用的是 Level1Quantizer 的具體實現(xiàn)撑教，如下朝墩，對訓(xùn)練數(shù)據(jù)進行聚類，得到聚類中心并保存下來
void Level1Quantizer::train_q1 (size_t n, const float *x, bool verbose, MetricType metric_type)
{
// 省略無關(guān)代碼
Clustering clus (d, nlist, cp);
quantizer->reset();
if (clustering_index) {
clus.train (n, x, *clustering_index);
quantizer->add (nlist, clus.centroids.data());
} else {
clus.train (n, x, *quantizer);
}
quantizer->is_trained = true;
}
add分片伟姐。根據(jù)輸入的大小收苏，按照固定的大小依次進行
add建立
invlists。根據(jù)train得到的聚類中心（保存在quantizer中）愤兵，每一個類中心對應(yīng)invlists中的一個桶鹿霸。向
invlists的桶里加入base。利用了openmp進行了并行加速
void IndexIVF::add (idx_t n, const float * x)
{
add_with_ids (n, x, nullptr);
}
void IndexIVF::add_with_ids (idx_t n, const float * x, const idx_t *xids)
{
// do some blocking to avoid excessive allocs
idx_t bs = 65536;
if (n > bs) {
for (idx_t i0 = 0; i0 < n; i0 += bs) {
idx_t i1 = std::min (n, i0 + bs);
if (verbose) {
printf(" IndexIVF::add_with_ids %ld:%ld\n", i0, i1);
}
add_with_ids (i1 - i0, x + i0 * d,
xids ? xids + i0 : nullptr);
}
return;
}
std::unique_ptr<idx_t []> idx(new idx_t[n]);
quantizer->assign (n, x, idx.get());
size_t nadd = 0, nminus1 = 0;
#pragma omp parallel reduction(+: nadd)
{
int nt = omp_get_num_threads();
int rank = omp_get_thread_num();
// each thread takes care of a subset of lists
for (size_t i = 0; i < n; i++) {
idx_t list_no = idx [i];
if (list_no >= 0 && list_no % nt == rank) {
idx_t id = xids ? xids[i] : ntotal + i;
size_t ofs = invlists->add_entry (
list_no, id,
flat_codes.get() + i * code_size
);
dm_adder.add (i, list_no, ofs);
nadd++;
} else if (rank == 0 && list_no == -1) {
dm_adder.add (i, -1, 0);
}
}
}
ntotal += n;
}
searchSearch corse_dis秆乳。搜索離query最近的聚類中心Search invlists懦鼠。在最近的nprobe個聚類中心對應(yīng)的invlists中進行暴力heap搜索，得到topk
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]);
double t0 = getmillisecs();
quantizer->search (n, x, nprobe, coarse_dis.get(), idx.get());
indexIVF_stats.quantization_time += getmillisecs() - t0;
t0 = getmillisecs();
invlists->prefetch_lists (idx.get(), n * nprobe);
search_preassigned (n, x, k, idx.get(), coarse_dis.get(),
distances, labels, false);
indexIVF_stats.search_time += getmillisecs() - t0;
}
ProductQuantizer
用來做 PQ 量化算法的類矫夷，關(guān)于 PQ 量化算法，可以參考 pq算法憋槐。簡單來說双藕，我們需要得到用來量化的碼本，然后我們可以對輸入的向量進行解碼和編碼阳仔。得到碼本的過程在 ProductQuantizer::train 中忧陪，包含
- 將輸入向量按照維度切分成
PQ段扣泊，每段的維度是dsub - 得到每段的聚類中心，這就是碼本
編碼和解碼的過程就是將輸入向量轉(zhuǎn)化為碼本里的 idx嘶摊，可以看出延蟹，量化是存在一定的誤差，其中叶堆，PQ 越大阱飘，誤差越小
void ProductQuantizer::train (int n, const float * x)
{
if (train_type != Train_shared) {
train_type_t final_train_type;
final_train_type = train_type;
if (train_type == Train_hypercube ||
train_type == Train_hypercube_pca) {
if (dsub < nbits) {
final_train_type = Train_default;
printf ("cannot train hypercube: nbits=%ld > log2(d=%ld)\n",
nbits, dsub);
}
}
float * xslice = new float[n * dsub];
ScopeDeleter<float> del (xslice);
for (int m = 0; m < M; m++) {
for (int j = 0; j < n; j++)
memcpy (xslice + j * dsub,
x + j * d + m * dsub,
dsub * sizeof(float));
Clustering clus (dsub, ksub, cp);
// we have some initialization for the centroids
if (final_train_type != Train_default) {
clus.centroids.resize (dsub * ksub);
}
switch (final_train_type) {
case Train_hypercube:
init_hypercube (dsub, nbits, n, xslice,
clus.centroids.data ());
break;
case Train_hypercube_pca:
init_hypercube_pca (dsub, nbits, n, xslice,
clus.centroids.data ());
break;
case Train_hot_start:
memcpy (clus.centroids.data(),
get_centroids (m, 0),
dsub * ksub * sizeof (float));
break;
default: ;
}
if(verbose) {
clus.verbose = true;
printf ("Training PQ slice %d/%zd\n", m, M);
}
IndexFlatL2 index (dsub);
clus.train (n, xslice, assign_index ? *assign_index : index);
set_params (clus.centroids.data(), m);
}
} else {
Clustering clus (dsub, ksub, cp);
if(verbose) {
clus.verbose = true;
printf ("Training all PQ slices at once\n");
}
IndexFlatL2 index (dsub);
clus.train (n * M, x, assign_index ? *assign_index : index);
for (int m = 0; m < M; m++) {
set_params (clus.centroids.data(), m);
}
}
}
IndexIVFPQ
ivfpq 算法在 ivf 的基礎(chǔ)上，對 base 做 pq虱颗。大家可以自行參考代碼
