Python FAISS增量添加索引方法:解决IndexIVFPQ无法增量更新问题
Faiss IndexIVFPQ 增量添加向量解决方案
核心逻辑:IndexIVFPQ的训练步骤仅需要使用少量和全量数据分布一致的采样数据即可完成,不需要加载全量嵌入,训练完成后的索引完全支持增量添加向量,无需一次性写入全量数据。
具体实现步骤
- 先从全量句子数据中随机采样部分样本,采样量建议为你设置的聚类中心数(nlist=1000)的30~100倍,也就是至少3万条嵌入向量,足够保证聚类精度,对BERT生成的同领域嵌入来说,采样10%的全量数据即可满足要求。
- 仅用采样得到的嵌入完成索引训练,训练完成后即可释放采样数据占用的内存。
- 后续按批次处理全量句子,生成嵌入后直接增量添加到索引中,和之前使用IndexFlatL2的逻辑完全一致。
修改后的实现代码
import random import numpy as np import faiss from tqdm import tqdm class FaissIVFPQIndex: def __init__(self, sentences, model, sample_ratio=0.1, nlist=1000, m=16, nbits=8): self.sentences = sentences self.model = model self.dim = 768 # 第一步:采样训练数据 self.train_sentences = random.sample(sentences, k=int(len(sentences)*sample_ratio)) # 保证采样数不低于nlist的30倍 if len(self.train_sentences) < nlist * 30: self.train_sentences = random.sample(sentences, k=min(nlist*30, len(sentences))) # 初始化IVFPQ索引 self.quantizer = faiss.IndexFlatL2(self.dim) self.index = faiss.IndexIVFPQ(self.quantizer, self.dim, nlist, m, nbits) # 先训练索引 self._train_index() def _train_index(self): # 计算训练用的嵌入 train_embs = [] result = self.model(self.train_sentences) for toks, embs in tqdm(result, desc="训练索引采样"): train_embs.extend(embs.numpy().tolist() if hasattr(embs, 'numpy') else embs) train_embs = np.stack(train_embs) # 训练索引,训练完成后train_embs即可释放 self.index.train(train_embs) del train_embs print("索引训练完成") def process_sentences(self): result = self.model(self.sentences) self.sentence_ids = [] self.token_ids = [] self.all_tokens = [] for i, (toks, embs) in enumerate(tqdm(result, desc="增量构建索引")): # 和IndexFlatL2一样逐句增量添加 all_embeddings = [] for j, (tok, emb) in enumerate(zip(toks, embs)): self.sentence_ids.append(i) self.token_ids.append(j) self.all_tokens.append(tok) all_embeddings.append(emb) all_embeddings = np.stack(all_embeddings) # 训练完成后直接增量add即可 self.index.add(all_embeddings) faiss.write_index(self.index, "faiss_Model_mini")
注意事项
- 采样的训练数据需要和全量数据分布一致,若后续新增的业务数据和初始训练数据的语义领域差异很大,建议补充新数据重新训练索引,避免检索精度下降
- 索引训练完成后,后续新增的向量不需要重新触发训练,直接调用
add方法即可写入 - 若需要边接收新数据边构建索引,仅首次构建时需要完成一次索引训练,后续新增的所有数据都可以直接增量添加
内容的提问来源于stack exchange,提问作者DevPy
相关产品推荐
相关产品推荐

