FAISS分批次添加向量后搜索仅返回单批次标签问题求助

我是FAISS新手，需基于内积实现相似向量检索。因笔记本内存有限，尝试向已训练调优的索引中分批次添加新向量，具体代码如下：

BATCH_SIZE = 100_000

for idx in tqdm(range( int( (len(df_tr) - N) / BATCH_SIZE))):
    index = faiss.read_index("indexes/trained_block.index")

    X = scipy.sparse.csr_matrix.toarray(X_full[idx * BATCH_SIZE:(idx + 1) * BATCH_SIZE]).astype('float32')
    faiss.normalize_L2(X)
    index.add(X)
    
    faiss.write_index(index, "indexes/block_{}.index".format(idx))

ivfs = []
for idx in tqdm(range( int( (len(df_tr) - N) / BATCH_SIZE))):
    index = faiss.read_index("indexes/block_{}.index".format(idx))
    ivfs.append(index.invlists)
    
    index.own_invlists = False

index = faiss.read_index("indexes/trained_block.index")

invlists = faiss.OnDiskInvertedLists(index.nlist, index.code_size, "indexes/merged_index.ivfdata")
ivf_vector = faiss.InvertedListsPtrVector() 

for ivf in ivfs: 
    ivf_vector.push_back(ivf)

ntotal = invlists.merge_from(ivf_vector.data(), ivf_vector.size())
index.ntotal = ntotal  
index.replace_invlists(invlists)  
faiss.write_index(index, "indexes/merged_index.index")  

但执行相似向量搜索时，返回的标签仅在0-100000（单批次大小）范围内，搜索代码及结果如下：

query = scipy.sparse.csr_matrix.toarray(vectorizer.transform(['sample'])).astype('float32')

index.nprobe = 10
D, I = index.search(query, 100)
print(I)

# 输出结果
[[93121 75215 99842 17907 17835 94646 93832 95062 87345 91036 87749 88507
86637 84382 82840 17261 84315 93969 78607 94330 99566 49088 95428 85836
77877 54978 91496 55231 75761 21885 64547 78052 81165  8370 81296 92231
67480 78757 16133 56417 43638 25109 77122 43178 53848 65869 49360  8440
3287 88457 21400 28398 15780 94845 35407 92137 55795 98621 13516 53323
23751 50605 62996 13813 59634 31121 86262  5930 39545 79405 91105 15471
23820 66360 46133 29015 28760 25257 15921  1079 47869 53775 26922 40162
79801 86765 82793 29220 53651 21723 11123 83319 47878 93225  2211 44512
65712 41331 83744 95585]]

问题原因 #

1. ID重复重叠：每次分批次添加时，都是从原始索引（含N个向量）重新读取并添加新批次向量，FAISS默认会从当前索引的ntotal值开始分配新ID。这导致每个批次的新向量ID范围都是[N, N+BATCH_SIZE-1]，所有批次的ID完全重叠，合并后检索返回的ID自然只会落在单批次的范围内。
2. 原始向量重复存储：每个block索引都包含原始索引的N个向量+当前批次向量，合并时原始向量的倒排表会被多次合并，导致索引中存在重复的原始向量，既浪费空间也会干扰检索结果。

解决思路与修正代码 #

步骤1：分批次添加时使用全局唯一ID #

创建仅包含当前批次向量的索引（而非基于原始索引叠加），并通过add_with_ids手动指定全局唯一ID，避免ID重叠。

import numpy as np
from tqdm import tqdm
import faiss
import scipy.sparse

BATCH_SIZE = 100_000
total_add = len(df_tr) - N

# 读取原始训练好的索引，获取关键参数
base_index = faiss.read_index("indexes/trained_block.index")
dim = base_index.d
nlist = base_index.nlist
# 根据原始索引类型生成工厂字符串，示例为IVFFlat，需根据实际调整
index_factory_str = f"IVF{nlist},Flat"

# 分批次创建仅含当前批次向量的索引
for idx in tqdm(range(int(total_add / BATCH_SIZE))):
    # 创建与原始索引同类型的空索引
    index = faiss.index_factory(dim, index_factory_str)
    # 复制原始索引的训练参数（如聚类中心、nprobe等）
    faiss.copy_parameters(base_index, index)
    # 复用原始索引的量化器（避免重复训练）
    index.quantizer = base_index.quantizer
    index.quantizer.own_fields = False  # 防止重复存储量化器
    
    # 加载当前批次向量并归一化
    X = scipy.sparse.csr_matrix.toarray(X_full[idx * BATCH_SIZE:(idx + 1) * BATCH_SIZE]).astype('float32')
    faiss.normalize_L2(X)
    
    # 生成该批次向量的全局唯一ID
    start_id = N + idx * BATCH_SIZE
    end_id = N + (idx + 1) * BATCH_SIZE
    ids = np.arange(start_id, end_id, dtype=np.int64)
    
    # 添加向量并指定全局ID
    index.add_with_ids(X, ids)
    
    # 保存批次索引
    faiss.write_index(index, f"indexes/block_{idx}.index")

步骤2：合并原始索引与所有批次索引的倒排表 #

将原始索引的倒排表与各批次的倒排表合并，确保所有向量的ID全局唯一且无重复。

ivfs = []

# 先添加原始索引的倒排表
ivfs.append(base_index.invlists)
base_index.own_invlists = False  # 释放所有权，避免后续销毁时删除

# 再添加各批次索引的倒排表
for idx in tqdm(range(int(total_add / BATCH_SIZE))):
    index = faiss.read_index(f"indexes/block_{idx}.index")
    ivfs.append(index.invlists)
    index.own_invlists = False

# 创建磁盘存储的倒排表（解决内存限制）
invlists = faiss.OnDiskInvertedLists(nlist, base_index.code_size, "indexes/merged_index.ivfdata")
ivf_vector = faiss.InvertedListsPtrVector()

for ivf in ivfs:
    ivf_vector.push_back(ivf)

# 合并所有倒排表
ntotal = invlists.merge_from(ivf_vector.data(), ivf_vector.size())

# 创建最终合并后的索引
merged_index = faiss.index_factory(dim, index_factory_str)
faiss.copy_parameters(base_index, merged_index)
merged_index.quantizer = base_index.quantizer
merged_index.quantizer.own_fields = False
merged_index.nlist = nlist
merged_index.ntotal = ntotal
merged_index.replace_invlists(invlists)

# 保存最终索引
faiss.write_index(merged_index, "indexes/merged_index.index")

步骤3：验证检索结果 #

重新加载合并后的索引进行检索，此时返回的ID应覆盖所有向量的全局范围（0到len(df_tr)-1）。

query = scipy.sparse.csr_matrix.toarray(vectorizer.transform(['sample'])).astype('float32')
faiss.normalize_L2(query)  # 内积检索需保证查询向量也归一化

index = faiss.read_index("indexes/merged_index.index")
index.nprobe = 10
D, I = index.search(query, 100)
print(I)

内容的提问来源于stack exchange，提问作者Fedor