最近在搞一个项目,需要处理大量向量数据,还得快速找相似的东西。一开始我用的是暴力搜索,结果慢得像蜗牛爬,都快把我急死了。后来无意中发现了 Faiss 这个宝贝,简直是救命稻草啊!这玩意儿不仅快,而且还特别省内存,我马上就爱上它了。今天就给大家安利一下这个神器。
1.
Faiss 是个啥?
Faiss 是 Facebook AI 研究院开发的一个开源库,主要用来做向量的相似度搜索和聚类。它的全名是 Facebook AI Similarity Search,听起来就很高大上,对吧?但别被这名字吓到,其实用起来贼简单。
Faiss 最牛的地方在于它能处理超大规模的数据集,而且搜索速度飞快。你想啊,普通的方法处理几百万个向量就得跑好久,Faiss 能轻松应对上亿个向量,还不带喘气的。
2.
安装 Faiss
安装 Faiss 超级简单,就一行命令的事:
pip install faiss-cpu
要是你有 GPU,想用 GPU 加速,那就装这个:
pip install faiss-gpu
温馨提示:安装 GPU 版本前,确保你的 CUDA 环境配置正确哦,不然可能会踩坑。
3.
基本使用
咱们先来看看 Faiss 最基本的用法。假设你有一堆向量,想找找哪些向量离得近。
import numpy as np
import faiss
# 准备数据
d = 64 # 向量维度
nb = 100000 # 数据库中的向量数量
nq = 10000 # 查询的向量数量
# 随机生成一些向量数据
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
# 创建索引
index = faiss.IndexFlatL2(d)
# 把向量加到索引里
index.add(xb)
# 搜索最近邻
k = 4 # 找4个最近邻
D, I = index.search(xq, k)
print(f“搜索结果的距离: {D[:5]}”) # 打印前5个查询结果的距离
print(f“搜索结果的索引: {I[:5]}”) # 打印前5个查询结果的索引
这段代码看起来挺多,其实就干了这么几件事:
随便造了一堆向量数据
创建了个索引(这里用的是最简单的L2距离索引)
把数据塞进索引里
搜索最近邻
运行完这段代码,你就能看到每个查询向量的4个最近邻及其距离了。是不是感觉特别简单?
4.
Faiss 的索引类型
Faiss 提供了好多种索引类型,每种都有自己的特点。这里咱们就说说最常用的几种:
IndexFlatL2
这是最简单的索引类型,就是老老实实算欧氏距离。优点是精确,缺点是慢。
IndexIVFFlat
这个索引先把数据分成若干个桶,搜索的时候先找到可能的桶,再在里面搜索。速度比 FlatL2 快多了,但可能会漏掉一些结果。
nlist = 100 # 聚类中心的数量
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist)
IndexHNSWFlat
这是基于图的索引,搜索速度超快,而且精度也不错。缺点是建立索引时间长,占用空间大。
温馨提示:选择索引类型时要根据你的具体需求来。数据量小追求精确就用 FlatL2,数据量大追求速度就用 HNSW。
5.
实际应用:图片相似度搜索
光说不练假把式,咱们来个实际点的例子。假设你有一堆图片,想找找哪些图片比较像。可以这么搞:
from PIL import Image
import numpy as np
import faiss
import os
def image_to_vector(image_path):
# 把图片转成向量
img = Image.open(image_path).resize((224, 224))
return np.array(img).flatten().astype('float32')
# 假设图片都在 'images' 文件夹里
image_folder = 'images'
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]
# 把所有图片转成向量
vectors = []
for img_file in image_files:
vector = image_to_vector(os.path.join(image_folder, img_file))
vectors.append(vector)
vectors = np.array(vectors)
# 创建索引
d = vectors.shape[1] # 向量维度
index = faiss.IndexFlatL2(d)
index.add(vectors)
# 假设我们要找和第一张图片相似的图片
query_vector = vectors[0].reshape(1, -1)
k = 5 # 找5个最相似的
D, I = index.search(query_vector, k)
print(“最相似的图片:”)
for idx in I[0]:
print(image_files[idx])
这段代码先把图片转成向量(这里用的是最简单的像素值),然后建立索引,最后搜索相似图片。实际应用中,你可能会用更复杂的方法(比如CNN)来提取图片特征。
好啦,今天的 Faiss 入门就到这里。这玩意儿用途可广了,从推荐系统到图像检索,到处都能派上用场。赶紧去试试吧,保证你会爱上它!
推 荐 阅 读
领取专属 10元无门槛券
私享最新 技术干货