2. 基础使用
2.1 FAISS基础语法
安装与环境配置
bash
# CPU版本
pip install faiss-cpu
# GPU版本(需要CUDA环境)
pip install faiss-gpu
核心概念与基本流程
首先解释下为什么要有训练(train)、添加(add)和搜索(search)这几个操作:
- 训练(train):对于大多数复杂索引,训练过程能让索引学习向量数据的分布特征,以便后续更高效地组织和搜索数据。例如,IVF 索引会在训练阶段使用聚类算法对向量分组,PQ 索引会在训练阶段学习如何量化向量。但像
IndexFlatL2
这种暴力搜索索引,不需要学习数据分布,所以不用训练。 - 添加(add):此操作是将待搜索的向量数据存入索引中,只有添加了数据,索引才能对这些数据进行搜索。就像我们把书放进图书馆,之后才能在图书馆里找书。
- 搜索(search):当索引中有了数据后,就可以用这个操作传入查询向量,从已添加的数据里找出与之相似的向量。
所以,当索引创建好,完成必要的训练后,我们就可以使用 add
把数据添加进去,再用 search
进行相似性搜索。
python
import faiss
import numpy as np
# 基本使用流程示例
def basic_faiss_usage():
# 向量维度
dimension = 128
# 生成示例数据
np.random.seed(1234)
# 数据库向量:1000个128维向量
database_vectors = np.random.random((1000, dimension)).astype('float32')
# 查询向量:5个128维向量
query_vectors = np.random.random((5, dimension)).astype('float32')
# 1. 创建索引
index = faiss.IndexFlatL2(dimension)
# 2. 添加数据到索引
index.add(database_vectors)
print(f"索引是否已训练:{index.is_trained}")
# 3. 执行搜索
k = 3 # 返回每个查询的最相似3个结果
distances, indices = index.search(query_vectors, k)
print("搜索结果索引:", indices)
print("对应距离:", distances)
basic_faiss_usage()
2.2 核心函数详解
索引创建函数
python
# L2距离(欧氏距离)索引
index_l2 = faiss.IndexFlatL2(dimension)
# 内积相似度索引
index_ip = faiss.IndexFlatIP(dimension)
# 带ID映射的索引(便于跟踪原始ID)
index_id = faiss.IndexIDMap(index_l2)
数据操作函数
python
# 添加向量数据
def add_vectors_example():
dimension = 64
index = faiss.IndexFlatL2(dimension)
# 生成100个64维向量
vectors = np.random.random((100, dimension)).astype('float32')
# 添加向量到索引
index.add(vectors)
print(f"索引中的向量数量:{index.ntotal}")
# 添加带自定义ID的向量
index_with_ids = faiss.IndexIDMap(faiss.IndexFlatL2(dimension))
vector_ids = np.arange(100, 200, dtype=np.int64) # 自定义ID
index_with_ids.add_with_ids(vectors, vector_ids)
add_vectors_example()
搜索与查询函数
python
import faiss
import numpy as np
def search_operations():
dimension = 64
index = faiss.IndexFlatL2(dimension)
# 生成100个64维向量
vectors = np.random.random((100, dimension)).astype('float32')
# 添加向量到索引
index.add(vectors)
print(f"索引中的向量数量:{index.ntotal}")
k = 2
# 生成10个查询向量
queries = np.random.random((4, dimension)).astype('float32')
distances, indices = index.search(queries, k)
return distances, indices
distances, indices = search_operations()
print("距离:", distances)
print("索引:", indices)
索引状态管理
python
import faiss
import numpy as np
def index_management():
dimension = 64
index = faiss.IndexFlatL2(dimension)
vectors = np.random.rand(1000, dimension).astype('float32')
index.add(vectors)
# 检查索引状态
print(f"索引是否已训练:{index.is_trained}")
print(f"索引中的向量数量:{index.ntotal}")
# 重置索引(清空所有数据)
index.reset()
print(f"重置后向量数量:{index.ntotal}")
index_management()