企业官网建设流程全解析

faiss 离线 IVF 搜索怎么跑：分片训练、索引与集群 search 三步流程

【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faiss

faiss 仓库的 demos/offline_ivf 目录提供了一套基于 big batch search 的离线 IVF 流程，面向亿级向量库的批量离线 KNN 检索：先train_index训练索引，再index_shard把库向量分片写入索引，最后search在 Slurm 集群上把结果落盘，最后可用consistency_check做一致性校验。本文按 README 的 Run book 顺序走一遍这三步，并说明每一步的输入、产物和验证方式。

环境准备

按 demos/offline_ivf/README.md 的要求创建 conda 环境并安装依赖：

conda create --name oivf python=3.10 conda activate oivf conda install -c pytorch/label/nightly -c nvidia faiss-gpu=1.7.4 conda install tqdm conda install pyyaml conda install -c conda-forge submitit

其中submitit是--cluster_run模式向 Slurm 集群提交任务所用的库；faiss-gpu是必需的，因为index_shard会用 GPU 量化器加速粗量化，search步骤内部直接调用faiss.knn_gpu。

后续所有run.py命令都在 demos/offline_ivf 目录下执行，因为 run.py 通过from utils import ...、from offline_ivf import OfflineIVF这类相对同目录导入加载模块。

数据与配置：shard 数据集、生成 yaml

流程的数据入口由一个 yaml 配置描述，示例见 config_ssnpp.yaml，其中的关键项：

• d：向量维度，示例为 256；
• output：结果输出目录，示例为/checkpoint/marialomeli/offline_faiss/ssnpp，需要改成你自己的存储路径；
• index和nprobe：分别给出prod与non-prod两组 factory 串（如IVF8192,PQ128）和 nprobe 值。非集群运行时使用prod组的最后一个值；
• k：每个查询取的邻居数，示例为 50；
• index_shard_size：单个索引分片包含的向量数，示例为 50000000；
• query_batch_size：search 步骤每批处理的查询向量数，示例为 50000000；
• training_sample：训练取样的向量数，示例为 1572864；
• datasets：按数据集名组织，每个数据集包含root（文件所在目录）、size（总向量数）和files列表（name、dtype、format、size）。

示例配置里datasets.ssnpp_1B.root指向/checkpoint/marialomeli/ssnpp_data，是官方示例使用的 1B ssnpp 数据集路径；换成自己的数据时，root、size、files都要与实际文件一致。dataset.py 在构造数据集时会断言root目录存在、每个文件存在，且各文件向量数之和等于size，所以这些值必须真实可核对。另外 offline_ivf.py 在初始化时断言output目录已经存在，运行前先创建好它。

可选：把原始数据库切成 shard 文件

README 第一步标注“Optionally shard your dataset”：如果库向量已经是一组可加载的分片文件（npy 或无头 raw 二进制），可以跳过这步。示例仓库里用 create_sharded_ssnpp_files.py 把单个u8bin数据库切成若干 npy 分片：

python create_sharded_ssnpp_files.py --filepath <数据库.u8bin 路径> --data_batch 50000000

--data_batch指定每个分片的向量数，必须能整除总向量数，否则脚本断言报错。该脚本还会把分片写到--filepath指定的输出目录（注意脚本里第二个参数也用了--filepath这个名，实际映射到output_dir）。

生成数据集配置

generate_config.py 里写死了示例的root、file_names、d、dt四个变量，把它改成你的分片目录和文件列表后运行：

python generate_config.py

脚本会逐个读取文件（自动区分 npy 与 raw 格式）并print出一份包含size、root、d、files的 yaml 文本，粘贴进你的配置文件的datasets段即可。文件不存在时脚本会直接断言报错，可先据此核对文件列表。

第一步：train_index 训练索引

python run.py --command train_index --config config_ssnpp.yaml --xb ssnpp_1B

--xb指定配置中的数据库数据集名。offline_ivf.py 的train_index会做三件事：读取training_sample个库向量的前缀样本并去重；用faiss.index_factory(d, index 串, metric)建索引并设置by_residual = True后训练；把训练好的空索引写成模板文件。产物路径为：

{output}/{xb}/{factory串中的逗号换成下划线}.empty.faissindex

按示例配置即{output}/ssnpp_1B/IVF8192_PQ128.empty.faissindex。验证方式就是这个文件存在。注意该方法开头有一个断言：模板文件已存在时会直接报错 “the index template file already exists”，即train_index对同一配置只能跑一次，重跑会失败而不是覆盖。

第二步：index_shard 分片建立索引

python run.py --command index_shard --config config_ssnpp.yaml --xb ssnpp_1B

这一步把全部库向量按index_shard_size切分并逐个写入索引文件（offline_ivf.py 的index_shard）：

• 先读入train_index产出的模板文件，并断言nprobe <= quantizer.ntotal，即训练出的粗聚类中心数要不少于 nprobe，否则报错提示 “check”；
• 粗量化器通过faiss.index_cpu_to_all_gpus放到 GPU 上，因此这一步也需要 GPU 环境；
• 依次处理第i个分片（从第i * index_shard_size个向量开始读index_shard_size个向量），add_with_ids写入后保存为{output}/{xb}/{factory}.shard_{i}。示例配置下 1B 向量共生成shard_0到shard_19共 20 个文件。

如果某个分片文件已存在（打开时抛FileExistsError），该分片会被跳过并记日志 “skipping shard”，所以这一步可以中断后重跑。验证方式：检查shard_0到shard_{nshards-1}全部存在，nshards = ceil(数据库总向量数 / index_shard_size)。

第三步：search 集群搜索

search 默认假设库向量就是查询向量（不传--xq时xq取xb）。结果按query_batch_size分批，每批用big_batch_search（见 contrib/big_batch_search.py）在 GPU 上计算，写入：

{output}/{xq}_in_{xb}/knn/I{批次号:010}_{factory}_np{nprobe}.npy {output}/{xq}_in_{xb}/knn/D_approx{批次号:010}_{factory}_np{nprobe}.npy

按示例配置，第一批结果文件是{output}/ssnpp_1B_in_ssnpp_1B/knn/I0000000000_IVF8192_PQ128_np512.npy和同名前缀的D_approx...npy。已存在的结果文件会被跳过（日志 “already exists”），同样支持中断重跑。

集群模式（README 主路径）

python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --cluster_run --partition <PARTITION-NAME>

--cluster_run时 run.py 使用submitit.AutoExecutor向 Slurm 提交任务，相关参数及默认值（来自 run.py 的 argparse 定义）：

两点来自代码的固定行为：job 时限在代码里写死为slurm_time=70 * 60（70 分钟），不是命令行参数；--partition被限定在一个固定的 choices 列表里（learnfair、devlab、scavenge、learnlab、nllb等），如果你的集群分区不在其中，需要先修改 run.py 里的 choices。search 启动时会读取SLURM_JOB_ID环境变量，在 knn 目录写record_*文件记录处理各批次的 job，用于识别旧 job 留下的零长度输出文件并清理。

查询向量与库向量不同的情况

README 明确说明：如果查询向量不同于库向量，需要先在第一步为查询数据准备好一个新数据集（shard 文件 + yaml 里的datasets条目），再通过--xq传入：

python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --xq <QUERIES_DATASET_NAME>

验证：consistency_check

README 建议在任何时候都可以跑一致性检查做 sanity check：

python run.py --command consistency_check --config config_ssnpp.yaml --xb ssnpp_1B

offline_ivf.py 的consistency_check依次做四组校验，任何一项不满足都会以断言失败中断：

1. 模板文件{output}/{xb}/{factory}.empty.faissindex存在；
2. 对每个索引分片，取出分片内偏移offset处的SMALL_DATA_SAMPLE（10000）个库向量，以nprobe=1搜索自己的 top-100，断言每个向量的自身 id 都在结果中；
3. search 产物检查：每个批次的I...npy/D_approx...npy存在且非空，I的列数等于k、行数等于该批查询数；
4. 用分片索引在线重搜一小批查询，与离线结果比对，断言匹配比例ratio > 0.95，且逐查询距离和与在线结果np.allclose（rtol=0.01）。

全部通过后日志打印 “done”。这是 README 给出的唯一显式校验手段，跑通它即可认为三步流程的产物是自洽的。

边界与限制

• 集群路径依赖 Slurm：--cluster_run只负责把 job 提交到集群（README 第 6 步 “Send jobs to the cluster to run search”），本地无 Slurm 时可以去掉--cluster_run在单机跑同一条 search 命令，此时会用prod组的最后一个 index 串和 nprobe 本地执行。
• search 与 index_shard 都依赖 GPU（faiss.get_num_gpus()、index_cpu_to_all_gpus、knn_gpu），纯 CPU 环境不适用该脚本。
• 除三步主路径外，代码里还有两个可选命令：merge_index把各shard_i合并为单个 ondisk 索引（依赖faiss.contrib.ondisk.merge_ondisk），evaluate在evaluation_sample个样本上对比精确与近似搜索并打印交集指标、平均距离等日志；二者不在 README 主流程内，需要时再单独运行。
• 数据集文件的dtype按配置原样加载（示例为uint8，对应 ssnpp 的整型归一化向量），换成 float 数据时同步修改files里的dtype即可，不要改动d与文件实际维度不一致。

参考资料：demos/offline_ivf/README.md、demos/offline_ivf/run.py、demos/offline_ivf/offline_ivf.py、demos/offline_ivf/config_ssnpp.yaml、contrib/big_batch_search.py。

【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faiss