【Faiss】索引选择指南（三）

最新推荐文章于 2024-08-23 19:55:45 发布

mjiansun

最新推荐文章于 2024-08-23 19:55:45 发布

阅读量2.2k

点赞数 1

分类专栏：算法与数据结构 Python

原文链接：https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index

版权

Python 同时被 2 个专栏收录

368 篇文章 33 订阅

订阅专栏

算法与数据结构

34 篇文章 7 订阅

订阅专栏

https://github.com/liqima/faiss_note/blob/master/3.Basics%20%E9%80%89%E6%8B%A9%E5%90%88%E9%80%82%E7%9A%84index%E7%B1%BB%E5%9E%8B.ipynb

选择合适的index类型

选择index类型并没有一套精准的法则可以依据，需要根据自己的实际情况选取。下面的几个问题可以作为选取index的参考。

是否需要精确的结果

如果需要，应该使用“Flat” 只有 IndexFlatL2 能确保返回精确结果。一般将其作为baseline与其他索引方式对比，以便在精度和时间开销之间做权衡。
不支持add_with_ids，如果需要，可以用“IDMap, Flat”。
支持GPU。

#导入faiss
import sys
sys.path.append('/home/maliqi/faiss/python/')
import faiss

#数据
import numpy as np 
d = 512          #维数
n_data = 2000   
np.random.seed(0) 
data = []
mu = 3
sigma = 0.1
for i in range(n_data):
    data.append(np.random.normal(mu, sigma, d))
data = np.array(data).astype('float32')

#ids, 6位随机数
ids = []
start = 100000
for i in range(data.shape[0]):
    ids.append(start)
    start += 100
ids = np.array(ids)

#不支持add_with_ids
index = faiss.index_factory(d, "Flat")
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)
[[   0  798  879  223  981 1401 1458 1174  919   26]
 [   1  981 1524 1639 1949 1472 1162  923  840  300]
 [   2 1886  375 1351  518 1735 1551 1958  390 1695]
 [   3 1459  331  389  655 1943 1483 1723 1672 1859]
 [   4   13  715 1470  608  459  888  850 1080 1654]]

index = faiss.index_factory(d, "IDMap, Flat")
index.add_with_ids(data, ids)
dis, ind = index.search(data[:5], 10)
print(ind)   # 返回的结果是我们自己定义的id

结果为：
[[100000 179800 187900 122300 198100 240100 245800 217400 191900 102600]
 [100100 198100 252400 263900 294900 247200 216200 192300 184000 130000]
 [100200 288600 137500 235100 151800 273500 255100 295800 139000 269500]
 [100300 245900 133100 138900 165500 294300 248300 272300 267200 285900]
 [100400 101300 171500 247000 160800 145900 188800 185000 208000 265400]]

关心内存开销

需要注意的是faiss在索引时必须将index读入内存。

如果不在意内存占用空间，使用“HNSWx”

如果内存空间很大，数据库很小，HNSW是最好的选择，速度快，精度高，一般4<=x<=64。不支持add_with_ids，不支持移除向量，不需要训练，不支持GPU。

index = faiss.index_factory(d, "HNSW8")
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)

结果为：
[[ 879  981   26 1132  807 1639   28 1334 1832 1821]
 [   1  981 1524 1639 1949 1472 1162  923  840  300]
 [   2 1886  375 1351  518 1958  390 1695 1707 1080]
 [   3 1459  331  389  655 1483 1723 1672 1859  650]
 [   4   13  715 1470  608  459 1080 1654  665  154]]

如果稍微有点在意，使用“..., Flat“

"..."是聚类操作，聚类之后将每个向量映射到相应的bucket。该索引类型并不会保存压缩之后的数据，而是保存原始数据，所以内存开销与原始数据一致。通过nprobe参数控制速度/精度。
支持GPU,但是要注意，选用的聚类操作必须也支持。

index = faiss.index_factory(d, "IVF100, Flat")
index.train(data)
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)

结果为：
[[   0  879  981 1401  919  143    2  807 1515 1393]
 [   1  511 1504  987  747  422 1911  638  851 1198]
 [   2  879  807  981 1401 1143  733  441 1324 1280]
 [   3  740  155 1337 1578 1181 1743  290  588 1340]
 [   4 1176  256 1186  574 1459  218  480 1828  942]]

如果很在意，使用”PCARx,...,SQ8“

如果保存全部原始数据的开销太大，可以用这个索引方式。包含三个部分，
1.降维
2.聚类
3.scalar 量化，每个向量编码为8bit 不支持GPU

index = faiss.index_factory(d, "PCAR16,IVF50,SQ8")  #每个向量降为16维
index.train(data)
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)

结果为：
[[   0  671  196 1025  624 1521  724  879 1281  533]
 [   1 1008  698  206  101  657  294  383  700  574]
 [   2 1594  754 1850  266  559  154 1723 1949 1910]
 [   3 1758  820  869 1067   14  211 1214   78 1445]
 [   4 1457  466  557 1604 1951  912  736 1974  836]]

如果非常非常在意，使用"OPQx_y,...,PQx"

y需要是x的倍数，一般保持y<=d，y<=4*x。支持GPU。

index = faiss.index_factory(d, "OPQ32_512,IVF50,PQ32")  
index.train(data)
index.add(data)
dis, ind = index.search(data[:5], 10)
print(ind)
[[   0 1686 1186 1552   47  517 1563 1738 1748  125]
 [   1  747 1816   41 1599  380 1179  803 1964  422]
 [   2 1610 1886  928  397  874  676  535 1401  929]
 [   3  548   89  509 1337  865 1472 1210 1181 1578]
 [   4  260 1781 1001 1179   41   20  747 1803 1055]]

当数据集很大该怎么办

这个就需要使用上面"..."的聚类方法。数据集被聚类成多个桶，在搜索时，只访问了一小部分桶(nprobe个桶)。

聚类是在数据集向量的一个有代表性的样本上执行的，通常是数据集的一个样本。我们指出了这个样本的最佳尺寸。（这里也不懂什么意思）

如果低于1M条矢量："...,IVFx,..."

N是数据集中向量个数，x一般取值[4sqrt(N),16sqrt(N)],需要30x ～ 256x个向量的数据集去训练。

GPU支持：是。

老版本推荐：

如果在1M-10M，使用"...,IMI2x10,..."

使用k-means将训练集聚类为2^10个类，但是执行过程是在数据集的两半部分独立执行，即聚类中心有2^(2*10)个。

如果在10M-100M，使用"...,IMI2x12,..."

如果1M-10M条矢量："...,IVF65536_HNSW32,..."

IVF与HNSW联合使用HNSW进行聚类分配（不懂）。你将需要在30*65536和256*65536之间数量的数据进行训练。

GPU支持：否(在GPU上，使用上述IVF)。

如果10M-100M："...,IVF262144_HNSW32,..."

与上述相同，将65536替换为262144(2^18)。请注意，训练将是很慢的。仅在gpu上进行训练是可能的，所有在cpu上运行的东西，请参阅列车IVF_与__gpu.ipynb.

import numpy as np
import faiss

faiss.get_num_gpus()
#Out[2]:
#8

index = faiss.index_factory(128, "PCA64,IVF16384_HNSW32,Flat")
# get some training data
xt = faiss.rand((1000000, 128))


# baseline training without GPU
index.train(xt)
#CPU times: user 48min 58s, sys: 3min 44s, total: 52min 42s
#Wall time: 6min 28s

index2 = faiss.index_factory(128, "PCA64,IVF16384_HNSW32,Flat")
index_ivf = faiss.extract_index_ivf(index2)
clustering_index = faiss.index_cpu_to_all_gpus(faiss.IndexFlatL2(64))
index_ivf.clustering_index = clustering_index
# training with GPU
index2.train(xt)
#CPU times: user 8.82 s, sys: 2.51 s, total: 11.3 s
#Wall time: 3.62 s

如果100M-1B条矢量："...,IVF1048576_HNSW32,..."

与上述相同，将65536替换为1048576(2^20)。训练会更慢！