faiss Offline IVF 实战基于 big batch search 的十亿级向量离线 KNN 全流程【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faissOffline IVF 是 faiss 仓库中位于 demos/offline_ivf 的一个端到端演示工程它把 faiss 的 big batch search 机制、分片索引、GPU 加速的 kNN 计算与集群作业调度组合在一起用于对超大规模示例为 SSNPP 10 亿条 256 维向量数据集做离线全库最近邻检索。读完本文你将掌握从数据集分片、索引训练、分片建库到集群检索与一致性校验的完整操作流程并理解每个步骤在 offline_ivf.py 中的具体实现逻辑。一、环境准备按照 demos/offline_ivf/README.md 的说明工程依赖 faiss 的 GPU 版本官方给出的环境创建方式如下conda create --name oivf python3.10 conda activate oivf conda install -c pytorch/label/nightly -c nvidia faiss-gpu1.7.4 conda install tqdm conda install pyyaml conda install -c conda-forge submitit注意其中的适用前提安装的是faiss-gpu1.7.4的 nightly 版本检索阶段search命令会实际调用faiss.knn_gpu与StandardGpuResources因此必须有可用的 GPUsubmitit是集群作业调度库--cluster_run模式依赖它向 Slurm 集群提交作业见 run.pytqdm、pyyaml分别用于进度显示和读取 yaml 配置。二、数据集分片与配置文件生成2.1 分片脚本原始数据库例如 big-ann-benchmarks 格式的.u8bin文件需要先切分成多个可独立读取的分片文件。仓库中对应的脚本是 create_sharded_ssnpp_files.pyREADME 中提及的create_sharded_dataset.py即此脚本的现名python create_sharded_ssnpp_files.py \ --filepath /path/to/FB_ssnpp_database.u8bin \ --data_batch 50000000 \ --output_dir /path/to/ssnpp_data脚本逻辑很直接先用xbin_mmap以np.memmap方式零拷贝映射整个.u8bin文件跳过 8 字节的 n/d 头部然后按--data_batch默认 5000 万条/文件要求能整除总数逐批np.save成ssnpp_0000000000.npy、ssnpp_0000000001.npy这样的 npy 分片。2.2 配置生成generate_config.py 遍历指定目录下的分片文件自动探测每个文件是raw无头部的原始二进制还是npy格式并统计每个文件的向量条数最终yaml.dump出一份数据集描述配置。它接受脚本顶部硬编码的root目录与文件名列表示例中是 20 个ssnpp_{i:010}.npy维度 d256、dtypeuint8。2.3 配置文件结构仓库自带的示例配置 config_ssnpp.yaml 展示了完整的配置格式各字段含义如下字段示例值含义d256向量维度output/checkpoint/.../ssnpp索引与结果的输出根目录index.prod/index.non-prodIVF8192,PQ128等index factory 字符串prod 用于生产运行non-prod 用于多参数扫描nprobe.prod/nprobe.non-prod512 / 256~8192各实验组使用的 nprobe 值k50每个查询返回的近邻数index_shard_size50000000每个索引分片包含的向量条数query_batch_size50000000检索阶段每批处理的查询条数evaluation_sample10000评估evaluate时的查询采样条数training_sample1572864训练索引时采样的向量条数datasets.名称ssnpp_1B数据集描述root目录、size总数、files列表每个文件含dtype/format/name/size示例中ssnpp_1B数据集共 20 个 5000 万条的 uint8 npy 分片总计 10 亿条向量这正是 README Run book 中--xb ssnpp_1B引用的名称。三、Run book五步操作命令3.1 train_index训练索引模板python run.py --command train_index --config config_ssnpp.yaml --xb ssnpp_1B对应 offline_ivf.py 的 train_index()用self.xb_ds.get_first_n(training_sample, np.float32)读取数据前 1572864 条并经np.unique(axis0)去重对去重场景的数据集尤为重要faiss.index_factory(d, IVF8192,PQ128, metric)构建索引并train通过faiss.extract_index_ivf取出内部 IVF 节点强制设置by_residual True以残差编码提升 PQ 精度用faiss.write_index保存为*.empty.faissindex模板文件——只含训练好的量化工具不含任何向量后续所有分片都基于它构建。命令已存在模板文件时会直接断言失败避免重复训练。3.2 index_shard生成索引分片python run.py --command index_shard --config config_ssnpp.yaml --xb ssnpp_1B对应 index_shard()读取索引模板把 CPU quantizer 用faiss.index_cpu_to_all_gpus迁移到 GPU 加速粗量化按index_shard_size切分数据10 亿 / 5000 万 20 个分片逐分片index.reset()后流式add_with_ids其中 ID 是全局连续的start jj保证检索结果可以直接定位回原始数据若配置了 OPQ 等前置变换IndexPreTransform输入会先经过_iterate_transformed做变换每写完一个分片立即faiss.write_index落盘为*.shard_{i}已存在的分片文件会跳过FileExistsError分支支持断点续建。分片数量在构造函数中计算nshards ceil(xb_ds.size / index_shard_size)见 offline_ivf.py#L95-L98。3.3 search集群批量检索python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B \ --cluster_run --partition PARTITION-NAME这是整个工程的核心对应 search()。其流程打开分片索引_open_sharded_index读取模板索引后把所有分片的InvertedLists收集进一个指针向量用faiss.HStackInvertedLists水平拼接后replace_invlists挂回索引offline_ivf.py#L768-L793。这样无需合并物理文件内存映射即可逻辑上得到一个完整索引按 query_batch_size 批量处理每批查询先经 GPU quantizer 做粗量化得到q_assign每个查询的 nprobe 个候选倒排列表调用 big_batch_search这是离线检索效率的关键。它以methodknn_function方式传入自定义_knn_function——内部用faiss.knn_gpu在指定 GPU 上以 float16 计算一个倒排列表内全部查询子集 × 库向量的距离并设置查询/向量内存上限4GB / 8GB防止显存溢出容错与断点每个批次输出到I{offset}_...npy/D_approx{offset}_...npy文件已存在则整批跳过Slurm 任务还会写record_{offset}.txt记录 job id若同一 job 重跑且结果文件为空会清理后重算。big_batch_search 的实现在 contrib/big_batch_search.pyreorder_assign把所有查询按其粗量化结果做桶排序bucket sort随后每个倒排列表只处理被分配到该列表的查询子集桶级计算结果再统一汇入全局ResultHeap。相比逐查询检索它把每个桶的计算变成一次大型矩阵运算天然契合 GPU GEMMthreadedfaiss.get_num_gpus() * 8、prefetch_threads等参数让预取下一桶 计算当前桶 写回结果流水线并行还支持checkpoint/checkpoint_freq定期落盘search 中每 7200 秒一次进程被杀后可从检查点恢复。关于查询集README 特别说明search 默认假设数据库向量即查询向量args.xq缺省时自动等于args.xb见 offline_ivf.py#L52-L53。如果查询集不同需要先在配置中准备新数据集第 1 步再传入--xqpython run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --xq QUERIES_DATASET_NAME3.4 evaluate 与 consistency_checkREADME 最后一条命令用于任意步骤后的健全性检查python run.py --command consistency_check --config config_ssnpp.yaml --xb ssnpp_1Bconsistency_check()的校验链相当完整offline_ivf.py#L817-L899索引模板存在模板文件可读索引分片自洽从每个分片偏移i*shard_size offset处抽 10000 条向量以nprobe1搜 100 近邻断言每条向量都能在自己的 100 个近邻中找到自己拼接后索引一致同样的抽样在HStackInvertedLists拼接索引上重做一遍检索结果抽查从 knn 输出目录加载每个批次的I/D结果文件与在线index.search的前 10000 条对比要求索引重合率 0.95、距离之和allclose(rtol0.01)。此外工程还支持 README 未展开的evaluate命令对evaluation_sample条查询分别做精确检索faiss.contrib.exhaustive_search的knn_ground_truth与近似检索输出逐 k 的knn_intersection_measure近邻交集率、精确/近似/精化距离均值以及交集基数分布utils.py 的 get_intersection_cardinality_frequencies。配置中若开启evaluate_by_margin还会基于2 * D(a,b) / (mean_a mean_b)这一双向检索 margin 公式utils.py 的 margin统计双向互为近邻的可信近邻对这是去重类应用的典型评估口径。四、run.py 命令行参数全解run.py 是所有命令的统一入口--command取值即OfflineIVF的方法名train_index/index_shard/search/merge_index/evaluate/consistency_check等通过eval(foivf.{args.command}())动态分发run.py#L32-L39。主要参数参数默认值说明--command必填—要执行的步骤--config必填—数据集/索引配置 yaml--xb必填—数据库向量数据集名对应 config 的datasets键--xq等于--xb查询向量数据集名--nt96搜索 OpenMP 线程数启动时调用faiss.omp_set_num_threads--no_residuals关闭残差训练时不设by_residual--cluster_runFalse通过 submitit 提交 Slurm 作业数组--job_name/--num_nodes/--tasks_per_nodeoivf / 1 / 1作业命名与节点配置--gpus_per_node/--cpus_per_task8 / 80单节点 GPU/CPU 配额--logs_dir/checkpoint/.../logs集群日志目录--partitionlearnlabSlurm 分区名choices 中列出了若干集群分区集群模式下process_options_and_run_jobs的调度策略是run.py#L42-L75evaluate命令会为indexfactory 字符串×nprobe的每个组合各提交一个作业non-prod 列表用于网格扫描其他命令则按 factory 字符串各提交一个作业本地运行则只取prod列表的最后一个值。五、输出文件布局所有产物都组织在output/xb/、output/xq/、output/xq_in_xb/三棵目录下xb/{factory}.empty.faissindex训练模板xb/{factory}.shard_{i}索引分片xq_in_xb/knn/检索结果I{offset}_{factory}_np{nprobe}.npy、D_approx{offset}_...npy及检查点文件CP...xq_in_xb/eval/评估中间量I_a_gt.npy、D_a_ann_...npy、margin_*.npy等。文件命名把 factory、nprobe 编码进文件名使得多组参数的扫描结果互不覆盖可独立对比。六、测试用例demos/offline_ivf/tests 提供了用小规模数据8 维、数千条向量、IVF256,PQ4或OPQ4,IVF256,PQ4跑通全流程的单元测试test_train_index验证 train_index 产出*.empty.faissindex模板test_index_shard_equal_file_sizes/test_index_shard_unequal_file_sizes覆盖分片大小整除/不整除库大小两种切分情形逐一确认.shard_i文件生成test_search串联 train → index_shard → search确认knn/I...结果文件落盘test_evaluate_without_margin含 OPQ 变体额外走merge_index后跑 evaluate检查 eval 目录下全部中间文件test_consistency_check验证在未完成前置步骤时一致性检查会正确报错。测试数据由 tests/testing_utils.py 的TestDataCreator生成可直接作为小数据集跑通本工程的参考配置。七、适用边界小结该工程面向离线批处理场景全库/全查询集的 KNN 计算产出 npy 结果文件不是在线服务索引search依赖 GPUfaiss.knn_gpu而train_index/index_shard/consistency_check主要在 CPU 侧完成数据集要求能被index_shard_size/query_batch_size整除search 中有整除性断言配置中的每个文件 size 之和必须等于数据集总 sizedataset.py 的断言从源码结构看big_batch_search 目前支持IndexIVFFlat、IndexIVFPQ、IndexIVFScalarQuantizer三类 IVF 索引BlockComputer示例配置的IVF8192,PQ128即属于 IVFPQ 路径。【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faiss创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 faiss Offline IVF 实战：基于 big batch search 的十亿级向量离线 KNN 全流程
faiss Offline IVF 实战基于 big batch search 的十亿级向量离线 KNN 全流程【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faissOffline IVF 是 faiss 仓库中位于 demos/offline_ivf 的一个端到端演示工程它把 faiss 的 big batch search 机制、分片索引、GPU 加速的 kNN 计算与集群作业调度组合在一起用于对超大规模示例为 SSNPP 10 亿条 256 维向量数据集做离线全库最近邻检索。读完本文你将掌握从数据集分片、索引训练、分片建库到集群检索与一致性校验的完整操作流程并理解每个步骤在 offline_ivf.py 中的具体实现逻辑。一、环境准备按照 demos/offline_ivf/README.md 的说明工程依赖 faiss 的 GPU 版本官方给出的环境创建方式如下conda create --name oivf python3.10 conda activate oivf conda install -c pytorch/label/nightly -c nvidia faiss-gpu1.7.4 conda install tqdm conda install pyyaml conda install -c conda-forge submitit注意其中的适用前提安装的是faiss-gpu1.7.4的 nightly 版本检索阶段search命令会实际调用faiss.knn_gpu与StandardGpuResources因此必须有可用的 GPUsubmitit是集群作业调度库--cluster_run模式依赖它向 Slurm 集群提交作业见 run.pytqdm、pyyaml分别用于进度显示和读取 yaml 配置。二、数据集分片与配置文件生成2.1 分片脚本原始数据库例如 big-ann-benchmarks 格式的.u8bin文件需要先切分成多个可独立读取的分片文件。仓库中对应的脚本是 create_sharded_ssnpp_files.pyREADME 中提及的create_sharded_dataset.py即此脚本的现名python create_sharded_ssnpp_files.py \ --filepath /path/to/FB_ssnpp_database.u8bin \ --data_batch 50000000 \ --output_dir /path/to/ssnpp_data脚本逻辑很直接先用xbin_mmap以np.memmap方式零拷贝映射整个.u8bin文件跳过 8 字节的 n/d 头部然后按--data_batch默认 5000 万条/文件要求能整除总数逐批np.save成ssnpp_0000000000.npy、ssnpp_0000000001.npy这样的 npy 分片。2.2 配置生成generate_config.py 遍历指定目录下的分片文件自动探测每个文件是raw无头部的原始二进制还是npy格式并统计每个文件的向量条数最终yaml.dump出一份数据集描述配置。它接受脚本顶部硬编码的root目录与文件名列表示例中是 20 个ssnpp_{i:010}.npy维度 d256、dtypeuint8。2.3 配置文件结构仓库自带的示例配置 config_ssnpp.yaml 展示了完整的配置格式各字段含义如下字段示例值含义d256向量维度output/checkpoint/.../ssnpp索引与结果的输出根目录index.prod/index.non-prodIVF8192,PQ128等index factory 字符串prod 用于生产运行non-prod 用于多参数扫描nprobe.prod/nprobe.non-prod512 / 256~8192各实验组使用的 nprobe 值k50每个查询返回的近邻数index_shard_size50000000每个索引分片包含的向量条数query_batch_size50000000检索阶段每批处理的查询条数evaluation_sample10000评估evaluate时的查询采样条数training_sample1572864训练索引时采样的向量条数datasets.名称ssnpp_1B数据集描述root目录、size总数、files列表每个文件含dtype/format/name/size示例中ssnpp_1B数据集共 20 个 5000 万条的 uint8 npy 分片总计 10 亿条向量这正是 README Run book 中--xb ssnpp_1B引用的名称。三、Run book五步操作命令3.1 train_index训练索引模板python run.py --command train_index --config config_ssnpp.yaml --xb ssnpp_1B对应 offline_ivf.py 的 train_index()用self.xb_ds.get_first_n(training_sample, np.float32)读取数据前 1572864 条并经np.unique(axis0)去重对去重场景的数据集尤为重要faiss.index_factory(d, IVF8192,PQ128, metric)构建索引并train通过faiss.extract_index_ivf取出内部 IVF 节点强制设置by_residual True以残差编码提升 PQ 精度用faiss.write_index保存为*.empty.faissindex模板文件——只含训练好的量化工具不含任何向量后续所有分片都基于它构建。命令已存在模板文件时会直接断言失败避免重复训练。3.2 index_shard生成索引分片python run.py --command index_shard --config config_ssnpp.yaml --xb ssnpp_1B对应 index_shard()读取索引模板把 CPU quantizer 用faiss.index_cpu_to_all_gpus迁移到 GPU 加速粗量化按index_shard_size切分数据10 亿 / 5000 万 20 个分片逐分片index.reset()后流式add_with_ids其中 ID 是全局连续的start jj保证检索结果可以直接定位回原始数据若配置了 OPQ 等前置变换IndexPreTransform输入会先经过_iterate_transformed做变换每写完一个分片立即faiss.write_index落盘为*.shard_{i}已存在的分片文件会跳过FileExistsError分支支持断点续建。分片数量在构造函数中计算nshards ceil(xb_ds.size / index_shard_size)见 offline_ivf.py#L95-L98。3.3 search集群批量检索python run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B \ --cluster_run --partition PARTITION-NAME这是整个工程的核心对应 search()。其流程打开分片索引_open_sharded_index读取模板索引后把所有分片的InvertedLists收集进一个指针向量用faiss.HStackInvertedLists水平拼接后replace_invlists挂回索引offline_ivf.py#L768-L793。这样无需合并物理文件内存映射即可逻辑上得到一个完整索引按 query_batch_size 批量处理每批查询先经 GPU quantizer 做粗量化得到q_assign每个查询的 nprobe 个候选倒排列表调用 big_batch_search这是离线检索效率的关键。它以methodknn_function方式传入自定义_knn_function——内部用faiss.knn_gpu在指定 GPU 上以 float16 计算一个倒排列表内全部查询子集 × 库向量的距离并设置查询/向量内存上限4GB / 8GB防止显存溢出容错与断点每个批次输出到I{offset}_...npy/D_approx{offset}_...npy文件已存在则整批跳过Slurm 任务还会写record_{offset}.txt记录 job id若同一 job 重跑且结果文件为空会清理后重算。big_batch_search 的实现在 contrib/big_batch_search.pyreorder_assign把所有查询按其粗量化结果做桶排序bucket sort随后每个倒排列表只处理被分配到该列表的查询子集桶级计算结果再统一汇入全局ResultHeap。相比逐查询检索它把每个桶的计算变成一次大型矩阵运算天然契合 GPU GEMMthreadedfaiss.get_num_gpus() * 8、prefetch_threads等参数让预取下一桶 计算当前桶 写回结果流水线并行还支持checkpoint/checkpoint_freq定期落盘search 中每 7200 秒一次进程被杀后可从检查点恢复。关于查询集README 特别说明search 默认假设数据库向量即查询向量args.xq缺省时自动等于args.xb见 offline_ivf.py#L52-L53。如果查询集不同需要先在配置中准备新数据集第 1 步再传入--xqpython run.py --command search --config config_ssnpp.yaml --xb ssnpp_1B --xq QUERIES_DATASET_NAME3.4 evaluate 与 consistency_checkREADME 最后一条命令用于任意步骤后的健全性检查python run.py --command consistency_check --config config_ssnpp.yaml --xb ssnpp_1Bconsistency_check()的校验链相当完整offline_ivf.py#L817-L899索引模板存在模板文件可读索引分片自洽从每个分片偏移i*shard_size offset处抽 10000 条向量以nprobe1搜 100 近邻断言每条向量都能在自己的 100 个近邻中找到自己拼接后索引一致同样的抽样在HStackInvertedLists拼接索引上重做一遍检索结果抽查从 knn 输出目录加载每个批次的I/D结果文件与在线index.search的前 10000 条对比要求索引重合率 0.95、距离之和allclose(rtol0.01)。此外工程还支持 README 未展开的evaluate命令对evaluation_sample条查询分别做精确检索faiss.contrib.exhaustive_search的knn_ground_truth与近似检索输出逐 k 的knn_intersection_measure近邻交集率、精确/近似/精化距离均值以及交集基数分布utils.py 的 get_intersection_cardinality_frequencies。配置中若开启evaluate_by_margin还会基于2 * D(a,b) / (mean_a mean_b)这一双向检索 margin 公式utils.py 的 margin统计双向互为近邻的可信近邻对这是去重类应用的典型评估口径。四、run.py 命令行参数全解run.py 是所有命令的统一入口--command取值即OfflineIVF的方法名train_index/index_shard/search/merge_index/evaluate/consistency_check等通过eval(foivf.{args.command}())动态分发run.py#L32-L39。主要参数参数默认值说明--command必填—要执行的步骤--config必填—数据集/索引配置 yaml--xb必填—数据库向量数据集名对应 config 的datasets键--xq等于--xb查询向量数据集名--nt96搜索 OpenMP 线程数启动时调用faiss.omp_set_num_threads--no_residuals关闭残差训练时不设by_residual--cluster_runFalse通过 submitit 提交 Slurm 作业数组--job_name/--num_nodes/--tasks_per_nodeoivf / 1 / 1作业命名与节点配置--gpus_per_node/--cpus_per_task8 / 80单节点 GPU/CPU 配额--logs_dir/checkpoint/.../logs集群日志目录--partitionlearnlabSlurm 分区名choices 中列出了若干集群分区集群模式下process_options_and_run_jobs的调度策略是run.py#L42-L75evaluate命令会为indexfactory 字符串×nprobe的每个组合各提交一个作业non-prod 列表用于网格扫描其他命令则按 factory 字符串各提交一个作业本地运行则只取prod列表的最后一个值。五、输出文件布局所有产物都组织在output/xb/、output/xq/、output/xq_in_xb/三棵目录下xb/{factory}.empty.faissindex训练模板xb/{factory}.shard_{i}索引分片xq_in_xb/knn/检索结果I{offset}_{factory}_np{nprobe}.npy、D_approx{offset}_...npy及检查点文件CP...xq_in_xb/eval/评估中间量I_a_gt.npy、D_a_ann_...npy、margin_*.npy等。文件命名把 factory、nprobe 编码进文件名使得多组参数的扫描结果互不覆盖可独立对比。六、测试用例demos/offline_ivf/tests 提供了用小规模数据8 维、数千条向量、IVF256,PQ4或OPQ4,IVF256,PQ4跑通全流程的单元测试test_train_index验证 train_index 产出*.empty.faissindex模板test_index_shard_equal_file_sizes/test_index_shard_unequal_file_sizes覆盖分片大小整除/不整除库大小两种切分情形逐一确认.shard_i文件生成test_search串联 train → index_shard → search确认knn/I...结果文件落盘test_evaluate_without_margin含 OPQ 变体额外走merge_index后跑 evaluate检查 eval 目录下全部中间文件test_consistency_check验证在未完成前置步骤时一致性检查会正确报错。测试数据由 tests/testing_utils.py 的TestDataCreator生成可直接作为小数据集跑通本工程的参考配置。七、适用边界小结该工程面向离线批处理场景全库/全查询集的 KNN 计算产出 npy 结果文件不是在线服务索引search依赖 GPUfaiss.knn_gpu而train_index/index_shard/consistency_check主要在 CPU 侧完成数据集要求能被index_shard_size/query_batch_size整除search 中有整除性断言配置中的每个文件 size 之和必须等于数据集总 sizedataset.py 的断言从源码结构看big_batch_search 目前支持IndexIVFFlat、IndexIVFPQ、IndexIVFScalarQuantizer三类 IVF 索引BlockComputer示例配置的IVF8192,PQ128即属于 IVFPQ 路径。【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faiss创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考