Files
FoodClassifier/faiss_vector_db/README.md
T

8.2 KiB

FAISS向量管理器

这是一个基于FAISS的向量管理器,专为食物分类项目设计,提供了完整的向量索引管理功能。

功能特性

  • FAISS索引管理: 支持创建、保存和加载FAISS索引
  • 向量操作: 支持向量的增删改查操作
  • 相似度搜索: 支持余弦相似度搜索和阈值过滤
  • 批量操作: 提供高效的批量向量操作接口
  • 元数据管理: 支持向量元数据的存储和检索
  • 多种索引类型: 支持不同的FAISS索引类型
  • 持久化存储: 支持索引的保存和加载

安装依赖

pip install faiss-cpu numpy
# 或者如果需要GPU支持
pip install faiss-gpu numpy

快速开始

1. 基本使用

from faiss.faiss_manager import FAISSManager
import numpy as np

# 创建管理器
manager = FAISSManager(
    dimension=512,  # 向量维度
    index_type="IndexFlatIP",  # 索引类型
    normalize_vectors=True,  # 启用向量标准化
    metric_type="cosine"  # 使用余弦相似度
)

# 添加单个向量
vector = np.random.random(512).astype(np.float32)
manager.add_vector(
    vector=vector,
    vector_id="food_001",
    metadata={"name": "回锅肉", "category": "川菜", "spicy_level": 3}
)

# 搜索相似向量
query_vector = np.random.random(512).astype(np.float32)
results = manager.search_similar(query_vector, k=5)

for vector_id, similarity, metadata in results:
    print(f"ID: {vector_id}, 相似度: {similarity:.4f}, 菜名: {metadata['name']}")

2. 批量操作

# 批量添加向量
batch_vectors = np.random.random((100, 512)).astype(np.float32)
batch_ids = [f"food_{i:03d}" for i in range(100)]
batch_metadata = [
    {"name": f"菜品{i}", "category": "川菜", "spicy_level": i % 5}
    for i in range(100)
]

results = manager.add_vectors_batch(batch_vectors, batch_ids, batch_metadata)
print(f"成功添加: {sum(results)}/{len(results)} 个向量")

# 批量删除向量
delete_ids = [f"food_{i:03d}" for i in range(10)]
delete_results = manager.delete_vectors_batch(delete_ids)
print(f"成功删除: {sum(delete_results)}/{len(delete_results)} 个向量")

3. 索引保存和加载

# 保存索引
save_path = "./my_faiss_index"
success = manager.save_index(save_path)
if success:
    print("索引保存成功")

# 加载索引
new_manager = FAISSManager(dimension=512)
success = new_manager.load_index(save_path)
if success:
    print("索引加载成功")
    
    # 验证加载结果
    info = new_manager.get_index_info()
    print(f"加载的向量数量: {info['active_vectors']}")

API 参考

FAISSManager 类

初始化参数

  • dimension (int): 向量维度
  • index_type (str): 索引类型,支持:
    • "IndexFlatIP": 内积索引(适合余弦相似度)
    • "IndexFlatL2": L2距离索引
    • "IndexIVFFlat": IVF索引(适合大规模数据)
    • "IndexHNSW": HNSW索引(高性能近似搜索)
  • normalize_vectors (bool): 是否标准化向量
  • metric_type (str): 距离度量类型 ("cosine", "l2", "ip")

主要方法

向量操作
# 添加单个向量
add_vector(vector, vector_id, metadata=None) -> bool

# 批量添加向量
add_vectors_batch(vectors, vector_ids, metadata_list=None) -> List[bool]

# 更新向量
update_vector(vector, vector_id, metadata=None) -> bool

# 删除单个向量
delete_vector(vector_id) -> bool

# 批量删除向量
delete_vectors_batch(vector_ids) -> List[bool]
搜索操作
# 相似度搜索
search_similar(query_vector, k=10, threshold=None) -> List[Tuple[str, float, Dict]]

# 根据ID获取向量信息
get_vector_by_id(vector_id) -> Optional[Tuple[np.ndarray, Dict]]
索引管理
# 保存索引
save_index(save_path) -> bool

# 加载索引
load_index(load_path) -> bool

# 清空索引
clear_index() -> bool

# 获取索引信息
get_index_info() -> Dict[str, Any]

使用场景

1. 食物图像检索

# 为食物分类项目设计的示例
manager = FAISSManager(dimension=2048, index_type="IndexFlatIP")

# 添加食物向量(通过CNN模型提取的特征)
food_features = extract_features_from_images(food_images)  # 假设的特征提取函数
food_metadata = [
    {"name": "回锅肉", "category": "川菜", "ingredients": ["猪肉", "青椒", "豆瓣酱"]},
    {"name": "西红柿鸡蛋", "category": "家常菜", "ingredients": ["西红柿", "鸡蛋"]},
    # ... 更多食物数据
]

manager.add_vectors_batch(food_features, food_ids, food_metadata)

# 查询相似食物
query_image_feature = extract_features_from_image(query_image)
similar_foods = manager.search_similar(query_image_feature, k=5, threshold=0.7)

for food_id, similarity, metadata in similar_foods:
    print(f"相似食物: {metadata['name']} (相似度: {similarity:.3f})")

2. 文本向量检索

# 用于文本相似度检索
manager = FAISSManager(dimension=768, index_type="IndexFlatIP")

# 添加文本向量(通过BERT等模型编码)
text_embeddings = encode_texts(recipe_texts)  # 假设的文本编码函数
recipe_metadata = [
    {"title": "回锅肉制作方法", "difficulty": "中等", "time": "30分钟"},
    # ... 更多菜谱数据
]

manager.add_vectors_batch(text_embeddings, recipe_ids, recipe_metadata)

# 搜索相关菜谱
query_embedding = encode_text("如何做川菜")
related_recipes = manager.search_similar(query_embedding, k=3)

性能优化建议

1. 索引类型选择

  • 小规模数据 (<10K向量): 使用 IndexFlatIPIndexFlatL2
  • 中等规模数据 (10K-1M向量): 使用 IndexIVFFlat
  • 大规模数据 (>1M向量): 使用 IndexHNSW 或更复杂的索引

2. 内存优化

# 对于大规模数据,考虑使用IVF索引
manager = FAISSManager(
    dimension=512,
    index_type="IndexIVFFlat",
    normalize_vectors=True
)

# 训练IVF索引(需要足够的训练数据)
if hasattr(manager.index, 'train') and not manager.index.is_trained:
    training_vectors = np.random.random((10000, 512)).astype(np.float32)
    manager.index.train(training_vectors)

3. 批量操作

# 优先使用批量操作而不是循环调用单个操作
# 好的做法
manager.add_vectors_batch(vectors, ids, metadata_list)

# 避免的做法
for vector, id, metadata in zip(vectors, ids, metadata_list):
    manager.add_vector(vector, id, metadata)

测试

运行测试以验证功能:

# 运行基本测试
python faiss/test_faiss_manager.py

# 运行使用示例
python faiss/example_usage.py

注意事项

  1. 向量维度: 所有向量必须具有相同的维度
  2. ID唯一性: 向量ID必须唯一,重复ID会导致更新操作
  3. 内存使用: FAISS索引会占用内存,大规模数据需要考虑内存限制
  4. 删除操作: FAISS不支持真正的删除,删除操作只是标记,需要重建索引来释放空间
  5. 向量标准化: 使用余弦相似度时建议启用向量标准化

故障排除

常见问题

  1. 维度不匹配错误

    # 确保所有向量维度一致
    assert vector.shape[-1] == manager.dimension
    
  2. 索引未训练错误

    # 对于IVF索引,需要先训练
    if hasattr(manager.index, 'train') and not manager.index.is_trained:
        manager.index.train(training_data)
    
  3. 内存不足

    # 使用更节省内存的索引类型
    manager = FAISSManager(dimension=512, index_type="IndexIVFFlat")
    

扩展功能

自定义距离度量

class CustomFAISSManager(FAISSManager):
    def custom_similarity_function(self, query_vector, k=10):
        """自定义相似度计算"""
        # 实现自定义逻辑
        pass

多模态检索

# 结合图像和文本特征
image_manager = FAISSManager(dimension=2048, index_type="IndexFlatIP")
text_manager = FAISSManager(dimension=768, index_type="IndexFlatIP")

# 融合检索结果
def multimodal_search(image_query, text_query, alpha=0.7):
    image_results = image_manager.search_similar(image_query, k=20)
    text_results = text_manager.search_similar(text_query, k=20)
    
    # 融合结果逻辑
    # ...

许可证

本项目采用 MIT 许可证。