向量检索ANN(Approximate Nearest Neighbor Search),指的是对于一个query向量,从向量库中找到和它距离最接近的k个向量。这是一个典型的topk任务,这里要的不是“精确topk”而是“近似topk”,更多考虑的是 精度和计算时间之间的权衡,牺牲一部分精度来降低计算时间。
Faiss的全称是Facebook AI Similarity Search,是FaceBook的AI团队针对大规模相似度检索问题开发的一个工具,使用C++编写,有python接口,对10亿量级的索引可以做到毫秒级检索的性能。
Faiss提供了各种不同的算法组合,比如牺牲一些精度,来提供检索的速度;另外一种思路是,预先构建索引,通过空间换时间。所以,在看faiss的各种算法实现的时候,需要关注的几个指标,检索速度,内存占比,检索精度。
一个高效的向量检索模型网网需要满足下面三个条件才能达到工业级可用:
- 实时查询,支持海量(百亿、千亿级别)规模库量级的;
- 存储高效,要求构建的向量索引模型数据压缩比高,达到大幅缩减内存使占用的目的;
- 召回精度好,topK有比较好的召回率,跟暴力搜索的结果相比;
常用的几种向量距离:欧式距离,向量内积,夹角余弦,汉明距离,杰卡德相似距离;
源码编译
git clone https://github.com/facebookresearch/faiss.git
git reset --hard dafdff110489db7587b169a0afee8470f220d295 # 我用的是这个版本
# 如果cmake版本太低,这里会提示需要升级cmake
cmake -B build -DFAISS_ENABLE_GPU=OFF -DFAISS_ENABLE_PYTHON=ON -DCMAKE_BUILD_TYPE=Debug -DBUILD_SHARED_LIBS=ON -DBUILD_TESTING=OFF .
# 升级cmake
# 下载
wget https://github.com/Kitware/CMake/releases/download/v3.23.0/cmake-3.23.0-linux-x86_64.sh
# 安装,安装好cd到对应的目录下查看是否有cmake的bin文件,配置环境变量
sudo bash ./cmake-3.23.0-linux-x86_64.sh --skip-licence --prefix=/usr
- 编译&运行demo
#基于上一步编译出来的faiss动态库,编译demo
#写一个makefile,指定faiss动态库的路径和各种编译依赖
CC = g++
CFLAGS = -g -Wall -I/home/your_path/faiss
LDFLAGS = -L/home/your_path/faiss/faiss/build/faiss -lfaiss
RPATH = -Wl,-rpath=/home/your_path/faiss/faiss/build/faiss
all: flat ivfflat hnsw ivfpq
flat: flat.cc
$(CC) $(CFLAGS) -o flat flat.cc $(LDFLAGS) $(RPATH)
#include <faiss/IndexFlat.h>
#include <faiss/index_io.h>
#include <iostream>
int main() {
std::cout << "hello faiss" << std::endl;
// Define the dimension of the vectors
int d = 4;
// Define the number of vectors to index
int nb = 100000; // 10w
// Generate some random data to index
float *xb = new float[d * nb]; // 64 * 10w
for(int i = 0; i < nb; i++) {
for(int j = 0; j < d; j++) {
xb[d * i + j] = drand48(); // xb = [[向量1],[向量2] ... ]
}
}
// Create an index
faiss::IndexFlatL2 index(d);
std::cout << "index dimension : " << d << std::endl;
// Index the data
index.add(nb, xb); // xb = [dimension * vector num]
// Save the index to disk
// faiss::write_index(&index, "example.index");
// Load the index from disk
//faiss::IndexFlatL2 index2(d);
//faiss::read_index(&index2, "example.index");
// Search the index
const int k = 5;
const int query_vector_num = 3;
long int *I = new long int[k * query_vector_num]; // 5 * 3
float *D = new float[k * query_vector_num]; // 5 * 3
// 从xb中的前query_vector_num个向量作为输入
// 从xb中寻找最近的k个返回
index.search(query_vector_num, xb, k, D, I);
// Print the results
for(int i = 0; i < query_vector_num; i++) {
printf("Query %d:\\n", i);
for(int j = 0; j < k; j++) {
show_vector(xb, I[k * i + j], d);
printf(" %ld (distance=%g)\\n", I[k * i + j], D[k * i + j]);
}
}
// Clean up
delete[] xb;
delete[] I;
delete[] D;
return 0;
}
这样一个简单的使用flat算法的faiss demo就运行起来了。
源码阅读
- 代码中的一些专有名词:
flat | 暴力检索 |
---|---|
code_size | 一个向量长度,默认情况下向量是不压缩的,长度=sizeof(float) * dim |
codes | 存放编码后的向量叫codes,用一维的float数组表示,存放方式 [[向量1],[向量2] ... ] |
metric_type | 计算距离的不同方法 |
centroid | ivfflat算法中的聚类中心 |
probe | ivfflat算法中搜索聚类中心相邻的区域 |
- 基类设计:
faiss的基类接口设计比较简单
全量数据 :data→add→train(可?。鷖earch
增量数据 :data→add→search
Index基类提供了基本的接口,后续各种继承基类,实现具体的功能。
// 基类index类,提供了基本的添加向量,检索topk,训练(对于需要预训练的检索类型)
// 不同的派生类会继承基类,来实现自定义实现
class Index {
int d; ///< vector dimension
idx_t ntotal; ///< total nb of indexed vectors
bool is_trained;
/// type of metric this index uses for search
MetricType metric_type;
/** Add n vectors of dimension d to the index.
*
* Vectors are implicitly assigned labels ntotal .. ntotal + n - 1
* This function slices the input vectors in chunks smaller than
* blocksize_add and calls add_core.
* @param n number of vectors
* @param x input matrix, size n * d
*/
virtual void add(idx_t n, const float* x) = 0;
/** query n vectors of dimension d to the index.
*
* return at most k vectors. If there are not enough results for a
* query, the result array is padded with -1s.
*
* @param n number of vectors
* @param x input vectors to search, size n * d
* @param k number of extracted vectors
* @param distances output pairwise distances, size n*k
* @param labels output labels of the NNs, size n*k
*/
virtual void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels,
const SearchParameters* params = nullptr) const = 0;
/** Perform training on a representative set of vectors
*
* @param n nb of training vectors
* @param x training vecors, size n * d
*/
virtual void train(idx_t n, const float* x);
/** encode a set of vectors
*
* @param n number of vectors
* @param x input vectors, size n * d
* @param bytes output encoded vectors, size n * sa_code_size()
*/
virtual void sa_encode(idx_t n, const float* x, uint8_t* bytes) const;
/** decode a set of vectors
*
* @param n number of vectors
* @param bytes input encoded vectors, size n * sa_code_size()
* @param x output vectors, size n * d
*/
virtual void sa_decode(idx_t n, const uint8_t* bytes, float* x) const;
}
1. Flat暴力求解
Flat就是比较简单的暴力算法求解topk,时间复杂度是O(M*N)。思路是对于一个查询向量,遍历整个向量库依次计算和每一个向量的距离,将每一个距离push到一个堆中。
int main() {
// 定义向量的维度
int d = 4;
// 定义向量的数量
int nb = 100000; // 10w
// Generate some random data to index
float *xb = new float[d * nb]; // 64 * 10w
for(int i = 0; i < nb; i++) {
for(int j = 0; j < d; j++) {
xb[d * i + j] = drand48(); // xb = [[向量1],[向量2] ... ]
}
}
// 创建flat l2 index
faiss::IndexFlatL2 index(d);
// 在flat l2 index上添加向量库
index.add(nb, xb); // xb = [dimension * vector num]
// 检索topk
const int k = 5;
const int query_vector_num = 3;
// 保存检索结果 I返回topk向量的index,D返回topk向量的距离
long int *I = new long int[k * query_vector_num]; // 5 * 3
float *D = new float[k * query_vector_num]; // 5 * 3
// 开始检索
// 从xb中的前query_vector_num个向量作为query向量,寻找topk,结果保存到D,I数组中
index.search(query_vector_num, xb, k, D, I);
// Print the results
for(int i = 0; i < query_vector_num; i++) {
printf("Query %d:\n", i);
for(int j = 0; j < k; j++) {
show_vector(xb, I[k * i + j], d);
printf(" %ld (distance=%g)\n", I[k * i + j], D[k * i + j]);
}
}
// Clean up
delete[] xb;
delete[] I;
delete[] D;
return 0;
}
暴力的计算topk代码:
// x: query vector
// y: data base
// d: dimension
// nx: the number of query vector number
// ny: the number of data base total
// res: result
template <class BlockResultHandler, bool use_sel = false>
void exhaustive_L2sqr_seq(
const float* x,
const float* y,
size_t d,
size_t nx,
size_t ny,
BlockResultHandler& res,
const IDSelector* sel = nullptr) {
using SingleResultHandler =
typename BlockResultHandler::SingleResultHandler;
int nt = std::min(int(nx), omp_get_max_threads());
FAISS_ASSERT(use_sel == (sel != nullptr));
// #pragma omp parallel num_threads(nt) //open mp 多线程优化
{
SingleResultHandler resi(res);
// #pragma omp for
// 从查询的nx个向量里面依次遍历 (对于每个query x,找到它的topk)
for (int64_t i = 0; i < nx; i++) {
const float* x_i = x + i * d;
const float* y_j = y;
resi.begin(i);
// 和向量库里面的向量依次比较
// j 表示的向量的index, y_j 表示的是向量的起始地址
for (size_t j = 0; j < ny; j++, y_j += d) {
if (use_sel && !sel->is_member(j)) {
continue;
}
float disij = fvec_L2sqr(x_i, y_j, d); // 计算距离 这里没有开平方根
resi.add_result(disij, j); // 加入堆中
}
resi.end();
}
}
}
2. IVFFlat
IVF表示的是inverted file ,相较于暴力匹配多了一步预筛的作用。目的是减少需要计算距离的目标向量的个数,做法就是直接对所有向量做k-means。
IVFFlat的基本思路是:假设聚类产出了1024个聚类中心,那么每来一个查询向量,首先计算其与1024个聚类簇心的距离,然后选择距离最近的top N个簇,只计算查询向量与这几个簇底下的向量的距离。这样计算的复杂度降低为 O(M * N/cluster_num)
IVFFlat有什么问题,是否可以找到最优解?
如果query的向量处于两个聚类的边缘,很有可能被分到另外一个聚类中心,导致无法计算出最有结果。 但是我们可以通过各种方法提高它的精度。比如,找到与 xq 相近的几个 centroids, 然后在这几个区域内搜索. 我们每次搜索的范围(scope)覆盖的区域数量。在 faiss 里面叫 nprobe, 我们可以提高 nprobe 的数量, 以提高精确度。
int main() {
// ...
int nlist = 100; // 聚类中心
int k = 4; // top k
faiss::IndexFlatL2 quantizer(d);
faiss::IndexIVFFlat index(&quantizer, d, nlist);
assert(!index.is_trained);
// step1 :训练
// 目的是产生100个聚类中心
index.train(nb, xb);
assert(index.is_trained);
// step2:add添加向量
// 根据第一步计算出来的聚类中心,将向量添加到离他最近的聚类中心 (每个聚类中心都有对应的倒排向量)
index.add(nb, xb);
// step3:查询
// 根据聚类中心和nprob(搜索的相邻的聚类中心),找到对应聚类中心的倒排链,然后在倒排链上完成检索
idx_t* I = new idx_t[k * nq];
float* D = new float[k * nq];
index.search(nq, xq, k, D, I); // nq : the number of query , xy : query database
}
IndexIVF 类设计:
// IndexIVF 继承了 Index, 同时继承了 L1 量化器
struct IndexIVF : Index, Level1Quantizer {
InvertedLists* invlists; // 向量编码后存储在倒排表里面(包括很多个倒排文件) // 本质是两个vector嵌套的二维数组
bool own_invlists;
size_t code_size; // 一个d维向量最后被编码为多少字节
size_t nprobe; // 搜索的时候, 在几个区域里搜
size_t max_codes; // 一次搜索最多在几个区域里搜, 用于覆盖 nprobe 参数
// OpenMP 并行化模式
// 0: 默认情况, 将请求拆分
// 1: 在倒排表上进行并行化
// 2: 同时使用0和1
// 3: 将请求拆分为更细粒度
int parallel_mode;
const int PARALLEL_MODE_NO_HEAP_INIT = 1024;
// 记录向量id与倒排表中倒排文件的关系, 可选
// 用于支持 reconstruct 方法, 暂时不用关注
DirectMap direct_map;
IndexIVF(
Index* quantizer,
size_t d,
size_t nlist,
size_t code_size,
MetricType metric = METRIC_L2);
void reset() override;
void train(idx_t n, const float* x) override;
void add(idx_t n, const float* x) override;
void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
void search(
idx_t n,
const float* x,
idx_t k,
float* distances,
idx_t* labels) const override;
size_t remove_ids(const IDSelector& sel) override;
~IndexIVF() override;
...
};
step1 train:
// 训练 L1 量化器以及子量化器
// n 为训练集向量数量(向量维度d)
// x 为训练集(长度 n*d)
void IndexIVF::train(idx_t n, const float* x) {
if (verbose) // 打印详情, debug 用, 默认 false
printf("Training level-1 quantizer\n");
// 该方法继承自 Level1Quantizer, 训练 L1 量化器
train_q1(n, x, verbose, metric_type);
if (verbose)
printf("Training IVF residual\n");
// 通过残差训练子量化器, 默认的实现是啥也不干, 先不管它
train_residual(n, x);
// 标记我们的 IVF 索引已经训练好了, 可以用了
is_trained = true;
}
step2 add:
void IndexIVF::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
std::unique_ptr<idx_t[]> coarse_idx(new idx_t[n]);
// 1. 找到添加的向量对应的聚类中心
quantizer->assign(n, x, coarse_idx.get());
// 2. 开始添加
add_core(n, x, xids, coarse_idx.get());
}
// n : 添加的向量的数量
// x : 向量数据
// xids:nullptr
// coarse_idx: 对应的聚类中心数组
void IndexIVF::add_core(
idx_t n,
const float* x,
const idx_t* xids,
const idx_t* coarse_idx,
void* inverted_list_context) {
FAISS_THROW_IF_NOT(coarse_idx);
FAISS_THROW_IF_NOT(is_trained);
direct_map.check_can_add(xids);
size_t nadd = 0, nminus1 = 0;
for (size_t i = 0; i < n; i++) {
if (coarse_idx[i] < 0)
nminus1++;
}
std::unique_ptr<uint8_t[]> flat_codes(new uint8_t[n * code_size]);
encode_vectors(n, x, coarse_idx, flat_codes.get());
DirectMapAdd dm_adder(direct_map, n, xids);
#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 = coarse_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,
inverted_list_context);
dm_adder.add(i, list_no, ofs);
nadd++;
} else if (rank == 0 && list_no == -1) {
dm_adder.add(i, -1, 0);
}
}
}
ntotal += n;
}
step3 search:
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 size_t nprobe =
std::min(nlist, params ? params->nprobe : this->nprobe); // nprob 表示查找几个区域 默认是=1
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]); // 通过query的index,找到它的聚类中心的index
std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]); // 通过query的index,找到它到聚类中心的距离
double t0 = getmillisecs();
// IndexFlat::search 这里应该是找到query向量的聚类中心向量
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) {
// ... 先忽略
} else {
// 入口
sub_search_func(n, x, distances, labels, &indexIVF_stats);
}
}
堆设计
// CMin 泛型类
// 提供T对象的cmp操作,min类就是比较谁最小;
// T1类型支持传入多个参数,如果T相同,那么支持比较T1
template <typename T_, typename TI_>
struct CMin {
typedef T_ T;
typedef TI_ TI;
typedef CMax<T_, TI_> Crev; // reference to reverse comparison
// 提供cmp比较谁更小
inline static bool cmp(T a, T b) {
return a < b;
}
inline static bool cmp2(T a1, T b1, TI a2, TI b2) {
return (a1 < b1) || ((a1 == b1) && (a2 < b2));
}
// 返回T类型的最小值
inline static T neutral() {
return std::numeric_limits<T>::lowest();
}
static const bool is_max = false;
// 返回T类型的下一个值
inline static T nextafter(T x) {
return cmin_nextafter(x);
}
};
// 使用方式2:
using HeapForIP2 = CMin<uint16_t, int64_t>;
HeapForIP2 heap2;
std::cout << heap2.cmp(100, 200) << std::endl;
std::cout << heap2.neutral() << std::endl;
std::cout << heap2.nextafter(100) << std::endl;
// 使用方式1:
// IP内积实际上是算余弦相似度, 越大越相似, 使用最小堆求 topK 个最大值
// L2计算里氏距离, 越小越相似, 使用最大堆求 topK 个最小值
using HeapForIP = CMin<float, idx_t>;
using HeapForL2 = CMax<float, idx_t>;
heap_heapify<HeapForIP>(k, simi, idxi);
template <class C>
inline void heap_heapify(
size_t k,
typename C::T* bh_val,
typename C::TI* bh_ids,
const typename C::T* x = nullptr,
const typename C::TI* ids = nullptr,
size_t k0 = 0) {
if (k0 > 0)
assert(x);
if (ids) {
for (size_t i = 0; i < k0; i++)
heap_push<C>(i + 1, bh_val, bh_ids, x[i], ids[i]);
} else {
for (size_t i = 0; i < k0; i++)
heap_push<C>(i + 1, bh_val, bh_ids, x[i], i);
}
for (size_t i = k0; i < k; i++) {
bh_val[i] = C::neutral();
bh_ids[i] = -1;
}
}
template <class C>
inline void heap_push(
size_t k,
typename C::T* bh_val,
typename C::TI* bh_ids,
typename C::T val,
typename C::TI id) {
bh_val--; /* Use 1-based indexing for easier node->child translation */
bh_ids--;
size_t i = k, i_father;
while (i > 1) {
i_father = i >> 1;
if (!C::cmp2(val, bh_val[i_father], id, bh_ids[i_father])) {
/* the heap structure is ok */
break;
}
bh_val[i] = bh_val[i_father];
bh_ids[i] = bh_ids[i_father];
i = i_father;
}
bh_val[i] = val;
bh_ids[i] = id;
}
IVFFlat的主要参数:
向量维度 | d :128 |
---|---|
metrice_type | ip内积 |
nlist | nlist=3000 聚类区域个数 |
niter | niter = 25 聚类迭代次数 |
nredo | nredo = 1 聚类重复计算次数,取最小结果 |
nprobe | nprobe=3 搜索查找相邻聚类中心的数量 |
3. IVFPQ
IVFFlat 可以在牺牲一部分精度的前提下,提高检索的速度。但是flat的方式,并不会向量进行压缩,为了节约内存提出了ivfpq的算法。PQ 是 Product quantization(乘积量化) 的缩写,按照文章里面的测试,可以降低90%多的内存占用。
IVFPQ 的思路是:
- 首先对向量库的向量按照维度进行切分,比如把 n * 128维的向量切分为 n * (8 个 16维)的子向量;
- 分别对子向量进行聚类 (ivf提现在这里),比如每个子向量产出 256个聚类中心,一共 8 * 256 个聚类中心;
- 对于一个128维的向量而来,可以分为8个子向量,每个子向量都可以找到对应的聚类中心。那么,就可以用对应的聚类中心序号来代替这个向量。[1,2,xxxxxx, 127] --> [c1,c2,c3,c4,c5,c6,c7,c8]。这里c的值域空间是0~255,只需要8个bit就可以存下来;这样就可以把一个很长的向量用很小的空间保存下来了。(pq核心思想)
- 对于查询向量q来说,首先先把自己量化,变为[c1,c3,c4,c2,xxxx,cn]。然后,在量化过后的向量库上做topk,找到和自己最近的k个量化后的向量。然后,将这个k个向量转换为对应的原始向量。计算出距离。
int main() {
int nlist = 100; // 聚类中心数量
int k = 4; // topk
int m = 8; // number of subquantizers
faiss::IndexFlatL2 quantizer(d);
faiss::IndexIVFPQ index(&quantizer, d, nlist, m, 8);
index.train(nb, xb);
index.add(nb, xb);
idx_t* I = new idx_t[k * 5];
float* D = new float[k * 5];
index.search(5, xb, k, D, I);
return 0;
}
4. HNSW
HNSW是多层 negivate small world的缩写,HNSW是最早是Yury Malkov在这篇论文中提出来的 HNSW,感兴趣的可以看原文。HNSW的核心思想是利用了图论的思想(首先假设有一张图,先不考虑这张图是如何构建的),通过随机找一个入口点,从入口点开始按照各种算法去遍历图,直到直到相近的点。简单的来说,HNSW像是一张构建在图上的多层跳表,检索的时候利用跳表快速寻找到最近的点。
HNSW可以分成构建(build)和检索(search)两个阶段。构建阶段的主要目的是,根据向量去构建去一张多层的图;检索阶段,主要是根据query 向量 在上一阶段构建出来的图 进行查询,从图的最高层向下检索直到找到最近的topk个向量。
- 构建(build)阶段
- 利用指数分布的函数,随机产出每个向量点映射到的最高层;这一步选择出了哪些节点在哪些层;
- 从最高层开始向下,根据上一步得到的每一层的节点,去构建每一层的邻接关系;
- 每一层的构建逻辑是,每个节点依次插入
3.1 插入的时候,在目前正在构建的图中开始贪心检索;
3.2 从入口的entry point开始,找到entry point的所有邻居中 距离 插入向量最近的点;
3.3 把这个最近的点设置为新的entry point,重复上一步;
3.4 直到计算出来的距离 = prev_计算出来的距离 (说明达到局部最优),然后跳出循环;
3.5 这里找到了一个邻居,其他的m-1个通过bfs找到; - 剪枝(shrink)
当邻居的数量超过了m,需要进行剪枝。这里剪枝的主要出发点是,尽可能的去保留“高速高路”,即不要一堆相邻的点聚集在一起,而是保留一些离的很远的点,这么做的目的是,让跳表的的范围比较广。这个就是hnsw论文里面精华的核心思想了。
build阶段的产出是两个,一个是存储向量的文件(codes),一个是存储图的关系的文件(index)。下一步检索就是基于两个文件寻找query向量的topk邻居。
- 检索(search)阶段
- 从图的最高层entry point开始向下检索,一层一层的向下贪心查找,找到每一层最近的邻居,直到找到0层;
- 在第0层开始,从上一步找到的entry point开始,进行bfs(广度优先遍历),遍历的同时维护一个最大堆,保留这个过程找到的topk
- 直到bfs的次数达到的ef-search的上线,停止遍历。堆里面的元素就是找到的topk。
代码:
build阶段
#include <iostream>
#include <faiss/IndexHNSW.h>
#include <faiss/IndexFlat.h>
#include <faiss/index_io.h>
#include <iostream>
int main() {
int d = 64; // 向量维度
int nb = 10000; // 向量数量
// 创建一个随机的数据集
float* xb = new float[nb * d];
for (int i = 0; i < nb * d; i++) {
xb[i] = drand48();
}
// 构建索引
faiss::IndexHNSWFlat index(d, 32);
index.add(nb, xb);
// Save the index to disk
faiss::write_index(&index, "example.index");
delete[] xb;
return 0;
}
search阶段
这是另外一段代码:
#include <iostream>
#include <faiss/IndexHNSW.h>
#include <faiss/IndexFlat.h>
#include <faiss/index_io.h>
#include <iostream>
int main() {
int d = 64; // 向量维度
int k = 5; // 查询的最近邻数量
// 进行查询
faiss::Index* index = faiss::read_index("example.index");
float* xq = new float[d];
for (int i = 0; i < d; i++) {
xq[i] = drand48();
}
long int* I = new long int[k];
float* D = new float[k];
index->search(1, xq, k, D, I);
// 打印查询结果
std::cout << "Nearest neighbors:" << std::endl;
for (int i = 0; i < k; i++) {
std::cout << "Index: " << I[i] << ", Distance: " << D[i] << std::endl;
}
delete[] xq;
delete[] I;
delete[] D;
return 0;
}
4. faiss中关于计算效率的优化
faiss 大面积的使用了open mp来进行平行计算来提供计算效率,这种场景读多写少,for(i~n)循环计算的逻辑非常适合用open mp改造。举一个简单的例子:
#include <iostream>
#include <omp.h>
#include <vector>
#include <chrono>
int main() {
int n = 100000000;
std::vector<int> nums(n, 1);
int sum = 0;
auto start = std::chrono::high_resolution_clock::now();
// 无优化的串行计算
for (int i = 0; i < n; i++) {
sum += nums[i];
}
auto end = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> serial_time = end - start;
std::cout << "Serial sum: " << sum << " - Time: " << serial_time.count() << " seconds" << std::endl;
sum = 0; // 重置sum
start = std::chrono::high_resolution_clock::now();
// 使用OpenMP并行计算
#pragma omp parallel for reduction(+:sum)
for (int i = 0; i < n; i++) {
sum += nums[i];
}
end = std::chrono::high_resolution_clock::now();
std::chrono::duration<double> parallel_time = end - start;
std::cout << "Parallel sum: " << sum << " - Time: " << parallel_time.count() << " seconds" << std::endl;
std::cout << "Speedup: " << serial_time.count() / parallel_time.count() << std::endl;
return 0;
}
Serial sum: 100000000 - Time: 0.234739 seconds
Parallel sum: 100000000 - Time: 0.0501435 seconds
Speedup: 4.68134
open mp这种并行编程模型,提供了更高级的并行抽象,也更容易改造基于for循环的串行代码。
5. 小结
- Faiss不同索引的对比:
- 不在意内存:选择HNSW;
- 不在意时间:XXXFlat;
- 在意内存,在意时间,不要在意精度: PQ/IVFFALT
参考: