Faiss 实战示例精讲：demos 目录下的自动调优、磁盘 IVF 与分片合并全流程

Faiss 仓库的 demos/ 目录是一组可直接运行的功能演示程序，覆盖了从自动超参调优、磁盘 IVF 索引的分布式构建与合并，到 IMI/PQ 索引、加权 KMeans、NNDescend 图构建与多样性结果处理器等各类典型场景。本文以 demos/README.md 的导读为骨架，逐一走读其中的核心演示脚本，并结合 contrib/ondisk.py、faiss/invlists/OnDiskInvertedLists.h 等源码实现，讲清楚每个 Demo 的运行方式、分步流程与底层原理。读完后，你将能够独立运行这些 Demo，并把「磁盘上构建超大 IVF 索引」的完整工作流迁移到自己的业务中。

demos 目录总览：一份可运行的功能索引

demos/README.md 将目录定位为“Demos for a few Faiss functionalities”，即若干 Faiss 功能的示例集合。目录下实际包含的内容可分为两类：

• Python 演示脚本（单文件即可运行，通常基于 sift1M 数据集）：
  • demo_auto_tune.py：自动调优功能演示
  • demo_ondisk_ivf.py：倒排文件落盘的磁盘 IVF 索引构建
  • demo_client_server_ivf.py、demo_distributed_kmeans_torch.py、demo_qinco.py、index_pq_flat_separate_codes_from_codebook.py
• C++ 演示程序（通过 demos/CMakeLists.txt 以 EXCLUDE_FROM_ALL 方式注册，需显式指定目标构建）：
  • demo_sift1M.cpp：在 sift1M 数据集上训练、构建、自动调优并评估 Recall 的完整流程
  • demo_imi_flat.cpp、demo_imi_pq.cpp：多级倒排（IMI）索引
  • demo_ivfpq_indexing.cpp：IVFPQ 索引的训练、I/O 与检索
  • demo_nndescent.cpp、demo_residual_quantizer.cpp、demo_weighted_kmeans.cpp
• 带独立子目录的主题 Demo：
  • demos/diversity_filter/：自定义 ResultHandler 实现“同一分组最多返回 M 个结果”的多样性过滤检索
  • demos/offline_ivf/：离线构建 SSNPP 文件的 IVF 流水线（含独立 README.md 与测试）
  • demos/rocksdb_ivf/：用 RocksDB 作为倒排列表存储后端的示例（含独立 README.md）

从 demos/CMakeLists.txt 可以看到，每个 C++ Demo 都单独注册为可执行目标并链接 faiss 库，例如：

add_executable(demo_sift1M EXCLUDE_FROM_ALL demo_sift1M.cpp)
target_link_libraries(demo_sift1M PRIVATE faiss)

EXCLUDE_FROM_ALL 意味着这些目标不会随默认 make 构建，需要按名称显式编译，这也提示读者：Demo 是独立可执行的入口，编译产物即为完整程序。

demo_ondisk_ivf.py：倒排文件不落内存的 IVF 索引（README 主角之一）

README 对该脚本的原始描述是：演示如何构建一个将倒排文件（inverted file）数据存放在磁盘上的 Faiss 索引——当数据放不进 RAM 时使用。脚本在小型的 sift1M 数据集上演示，按阶段推进：0 阶段训练索引；1-4 阶段构建 4 个各自包含 1/4 数据的索引（可分布在多台机器上并行）；5 阶段将 4 个索引合并为一个直接写盘的索引（不需要整体放入 RAM）；6 阶段加载并测试该索引。

脚本 demos/demo_ondisk_ivf.py 通过命令行参数 stage 选择执行阶段，完整流程如下。

阶段 0：训练并保存训练好的索引

stage = int(sys.argv[1])
tmpdir = "/tmp/"

if stage == 0:
    xt = fvecs_read("sift1M/sift_learn.fvecs")
    index = faiss.index_factory(xt.shape[1], "IVF4096,Flat")
    index.train(xt)
    faiss.write_index(index, tmpdir + "trained.index")

要点：

• 使用 index_factory 的字符串描述 "IVF4096,Flat" 一步构建出 4096 个粗量化中心的 IVF 索引；
• 训练集是 sift1M 子集中的 sift_learn.fvecs；
• 训练完成后只保存 trained.index（此时不含向量数据，体积很小），后续各分片都从它加载起步。

阶段 1–4：四台“机器”各自 add 自己的 1/4 数据

if 1 <= stage <= 4:
    bno = stage - 1
    xb = fvecs_read("sift1M/sift_base.fvecs")
    i0, i1 = int(bno * xb.shape[0] / 4), int((bno + 1) * xb.shape[0] / 4)
    index = faiss.read_index(tmpdir + "trained.index")
    index.add_with_ids(xb[i0:i1], np.arange(i0, i1))
    faiss.write_index(index, tmpdir + "block_%d.index" % bno)

要点：

• 每个阶段读取同一份 trained.index，再各自 add 全量库的 1/4 切片，add_with_ids 保证合并后全局 id 不冲突；
• 四个阶段互不依赖，可以并行/分布式执行，这正是 README 所说“can be done in parallel on several machines”；
• 每个分片单独落盘为 block_0..3.index。

阶段 5：合并四个分片，倒排数据直接写入磁盘

if stage == 5:
    index = faiss.read_index(tmpdir + "trained.index")
    block_fnames = [tmpdir + "block_%d.index" % bno for bno in range(4)]
    merge_ondisk(index, block_fnames, tmpdir + "merged_index.ivfdata")
    faiss.write_index(index, tmpdir + "populated.index")

这里调用的 merge_ondisk 来自 contrib/ondisk.py，其源码揭示了“合并时不占满内存”的三处关键设计：

1. 以 MMAP 方式读入各分片：faiss.read_index(fname, faiss.IO_FLAG_MMAP)——注释明确说明“IO_FLAG_MMAP is to avoid actually loading the data, thus the total size of the inverted lists can exceed the available RAM”，即分片中的倒排数据并不真正读进内存；
2. 用 OnDiskInvertedLists 作为输出容器：faiss.OnDiskInvertedLists(nlist, code_size, ivfdata_fname) 将合并后的倒排数据直接写入 merged_index.ivfdata 文件；
3. merge_from_multiple 一次性归并：把各分片的 InvertedLists 归并进输出倒排列表，随后 replace_invlists(invlists, True) 替换索引的倒排列表，invlists.this.disown() 避免 Python GC 误释放。

该函数还有两条使用约束值得注意：assert not isinstance(trained_index, faiss.IndexIVFPQR)（IndexIVFPQR 不支持作为磁盘索引）与 assert index.ntotal == 0（只对空索引工作）。其底层由 faiss/invlists/OnDiskInvertedLists.cpp 实现，通过 mmap 打开数据文件、按偏移定位各桶的 code 区间。

阶段 6：加载磁盘索引并检索评测

if stage == 6:
    index = faiss.read_index(tmpdir + "populated.index")
    index.nprobe = 16

    xq = fvecs_read("sift1M/sift_query.fvecs")
    gt = ivecs_read("sift1M/sift_groundtruth.ivecs")
    D, I = index.search(xq, 5)

    recall_at_1 = (I[:, :1] == gt[:, :1]).sum() / float(xq.shape[0])
    print("recall@1: %.3f" % recall_at_1)

nprobe=16 控制了每次查询访问的粗量化桶数，是精度与延迟的折中参数；最后用预计算的 ground truth 计算 recall@1 验证结果。

适用前提：脚本假定当前目录下存在 sift1M/ 数据集（sift_learn.fvecs、sift_base.fvecs、sift_query.fvecs、sift_groundtruth.ivecs），临时文件写入 /tmp/。实际部署时把 tmpdir 与各阶段参数化即可，分片数不必固定为 4。

demo_auto_tune.py：自动调优功能演示（README 主角之二）

README 对 demo_auto_tune.py 的描述只有一句——“Demonstrates the auto-tuning functionality of Faiss”，脚本本身则完整展示了 Faiss 参数自动探索（ParameterSpace explore）的工作方式。

脚本的工作流程：

1. 加载数据与真值：读取 sift1M 的训练集、库、查询集与 ground truth；
2. 定义准则：faiss.OneRecallAtRCriterion(xq.shape[0], 1)，即“1 - recall@1”作为搜索质量指标，并通过 crit.set_groundtruth(None, gt)、crit.nnn = k 绑定真值；
3. 准备候选索引清单：脚本内置四组索引键，分别对应不同内存预算——

# criterion = 1-recall at 1
crit = faiss.OneRecallAtRCriterion(xq.shape[0], 1)
crit.set_groundtruth(None, gt)
crit.nnn = k

# indexes that are useful when there is no limitation on memory usage
unlimited_mem_keys = [
    "IMI2x10,Flat", "IMI2x11,Flat", "IVF4096,Flat", "IVF16384,Flat", "PCA64,IMI2x10,Flat",
]

# memory limited to 16 bytes / vector
keys_mem_16 = [
    "IMI2x10,PQ16", "IVF4096,PQ16", "IMI2x10,PQ8+8", "OPQ16_64,IMI2x10,PQ16",
]

# limited to 32 bytes / vector
keys_mem_32 = [
    "IMI2x10,PQ32", "IVF4096,PQ32", "IVF16384,PQ32",
    "IMI2x10,PQ16+16", "OPQ32,IVF4096,PQ32", "IVF4096,PQ16+16", "OPQ16,IMI2x10,PQ16+16",
]

# indexes that can run on the GPU
keys_gpu = [
    "PCA64,IVF4096,Flat", "PCA64,Flat", "Flat", "IVF4096,Flat", "IVF16384,Flat", "IVF4096,PQ32",
]

这四组列表本身就是一份“索引字符串语法 + 内存预算”的速查示例：PQ16/PQ32 表示乘积量化编码字节数，PQ8+8 表示两段不同子码字长的残差式 PQ，OPQ/PCA 前缀表示前置正交/降维变换，IMI 表示多级倒排。

4. 逐索引探索操作点：核心调用是 params.initialize(index) 后 opi = params.explore(index, xq, crit)——ParameterSpace 自动枚举该索引上所有可调参数（如 nprobe、efSearch 等）的取值组合，对每种组合执行检索并按准则评估，得到“性能—耗时”操作点序列 OperatingPoints；各索引的结果通过 op.merge_with(opi, index_key + " ") 汇入全局 Pareto 前沿；
5. 图形输出（可选）：若安装了 matplotlib，脚本会把各索引的操作点绘制为“1-recall@1 vs 每查询耗时(ms)”曲线，保存到 tmp/demo_auto_tune.png。

默认配置 keys_to_test = unlimited_mem_keys、use_gpu = False；若要跑 GPU 分支，需将 use_gpu 置为 True，并且要求 Faiss 编译了 GPU 支持（脚本会断言 faiss.StandardGpuResources 可用），并通过 faiss.index_cpu_to_gpu(res, dev_no, index) 把索引转移到设备 0。

C++ 侧对应实现是 faiss/AutoTune.h 与 faiss/AutoTune.cpp，Python 接口由 faiss/python/extra_wrappers.py 暴露 ParameterSpace、OperatingPoints 等封装。C++ Demo demo_sift1M.cpp 演示了同样的流程并多走了一步：在 ops.optimal_pts 中挑选第一个 perf > 0.5（即 recall@1 > 0.5）的操作点，用 params.set_index_parameters(index, selected_params) 把参数串写回索引，然后检索并手工统计 R@1/R@10/R@100：

faiss::OneRecallAtRCriterion crit(nq, 1);
crit.set_groundtruth(k, nullptr, gt);
crit.nnn = k;

faiss::ParameterSpace params;
params.initialize(index);
faiss::OperatingPoints ops;
params.explore(index, nq, xq, crit, &ops);

其余 C++ Demo 与主题子目录速览

demo_ivfpq_indexing.cpp（demos/demo_ivfpq_indexing.cpp）：用 128 维、20 万条向量演示 IVFPQ 的完整生命周期。粗量化器为独立对象 IndexFlatL2 coarse_quantizer(d)，中心数按经验公式 ncentroids = int(4 * sqrt(nb)) 选取（约 566 个）；IndexIVFPQ index(&coarse_quantizer, d, ncentroids, 4, 8) 中 4 是每码字节数（要求 d 为其倍数）、8 是每个子码的比特数；训练后通过 write_index 落盘，再走 add、读回、search 的 I/O 演示。源码中的注释特别提醒“the coarse quantizer should not be deallocated before the index”，这是 IVF 系索引的一个常见生命周期陷阱。

demo_weighted_kmeans.cpp（demos/demo_weighted_kmeans.cpp）：演示带权 KMeans 聚类。核心是 Clustering 对象：clus.train(n, input, *index, weights) 传入逐向量权重；Demo 用四种后端对照——IndexFlatL2、IndexFlatIP、IndexFlatIP + spherical（球面约束）与 IndexHNSWFlat(d, 32)（efSearch=128，用于展示聚类时近似最近邻加速）。训练结束后质心存于 clus.centroids，iteration_stats.back().obj 给出最后一次迭代的聚类目标值。

demo_qinco.py（demos/demo_qinco.py）：复现 QINCo 论文结果的 Python Demo。脚本头部 docstring 明确说明“training is not implemented in Faiss”，因此流程是加载外部训练好的 PyTorch 模型（/tmp/bigann_8x8_L2.pt），先在 PyTorch 侧 encode/decode，再用 faiss.QINCo(qinco) 包装同一模型在 Faiss 侧 encode/decode，比较两者的编码一致性与解码 MSE。运行前提按 docstring 说明：克隆 Qinco 参考代码、下载 bigann1M 数据集与模型文件，且需要 PyTorch 环境。

demo_client_server_ivf.py / demo_distributed_kmeans_torch.py：分别演示 IVF 索引的客户端-服务端检索架构（可配合 contrib/client_server.py 理解）与基于 PyTorch 的分布式 KMeans 训练，是大规模场景下“数据并行”的两条典型路径。

diversity_filter 子目录（demos/diversity_filter/demo_diversity_result_handler.cpp）：演示自定义 ResultHandler 实现“同一分组（group）最多返回 M 个 id”的多样性过滤。实现思路是从源码可见的：在候选距离相等（tie）的情况下允许同组 id 以任意顺序出现，比较函数 compare_results 会识别“同组不同序”属于合法情况（源码注释：“this is a valid case - same group, just different tie-breaking”），从而验证结果的正确性；处理器定义在 demos/diversity_filter/diversity_result_handler.h，构建规则见 demos/diversity_filter/CMakeLists.txt。这是“Faiss 检索漏斗可插拔”特性的典型示例——ResultHandler 接口见 faiss/impl/ResultHandler.h。

offline_ivf / rocksdb_ivf 子目录：前者是一条“离线产出 SSNPP 格式倒排文件”的流水线（配置见 demos/offline_ivf/config_ssnpp.yaml，入口 demos/offline_ivf/run.py，测试在 demos/offline_ivf/tests/）；后者展示把 InvertedLists 后端换成 RocksDB 的写法（demos/rocksdb_ivf/RocksDBInvertedLists.h），与本文的 on-disk IVF 形成互补：一个把倒排数据落到普通 mmap 文件，一个落到 KV 存储。

运行方式与适用前提

结合 INSTALL.md 与目录结构，各 Demo 的运行前提归纳如下：

1. 数据集：demo_auto_tune.py、demo_ondisk_ivf.py、demo_sift1M.cpp 均假定当前工作目录下存在 sift1M/ 数据集（sift1M 的 fvecs/ivecs 文件）。demo_sift1M.cpp 头部注释明确给出了数据集的获取来源（Texmex 的 ANN_SIFT1M），解压到 sift1M/ 子目录即可。
2. Python Demo：python demos/demo_ondisk_ivf.py <stage>，stage 依次取 0、1、2、3、4、5、6；python demos/demo_auto_tune.py 直接运行（matplotlib 为可选依赖，用于画操作点曲线，输出到 tmp/demo_auto_tune.png）。
3. C++ Demo：先编译 Faiss 库，再按目标名构建，例如 cmake --build build --target demo_sift1M demo_ivfpq_indexing demo_weighted_kmeans demo_nndescent demo_imi_flat demo_imi_pq（这些目标因 EXCLUDE_FROM_ALL 不会自动构建）。
4. GPU 分支：demo_auto_tune.py 的 GPU 路径要求编译时启用 GPU 支持，脚本会主动断言并给出提示“Faiss was not compiled with GPU support”。

从源码结构看，这套 Demo 与 Faiss 主库的对应关系很清晰：自动调优对应 faiss/AutoTune.*，磁盘倒排对应 faiss/invlists/OnDiskInvertedLists.* 与 contrib/ondisk.py，索引字符串构建对应 faiss/index_factory.cpp，而 sift1M 上的 Recall 评测方式（读取 ground truth、统计 R@k）与 benchs/ 下基准测试脚本一脉相承。

小结

• demos/README.md 点名的两个核心 Demo 分别对应 Faiss 的两大进阶能力：ParameterSpace 自动调优（demo_auto_tune.py）与倒排文件落盘的超大 IVF 索引构建（demo_ondisk_ivf.py + contrib/ondisk.py 的 MMAP 分片合并）。
• demo_ondisk_ivf.py 的“训练 → 分片 add → 合并落盘 → 检索评测”七阶段工作流，配合 IO_FLAG_MMAP 与 OnDiskInvertedLists，是构建超出 RAM 容量的 IVF 索引的参考实现。
• 其余 C++ Demo（IMI、IVFPQ、加权 KMeans、NNDescend、QINCo、多样性过滤）与 offline_ivf / rocksdb_ivf 子目录，共同覆盖了索引构建、存储后端替换与结果策略定制等典型场景，均可在 demos/CMakeLists.txt 中找到对应的构建目标。