faiss 离线 IVF 搜索怎么跑分片训练、索引与集群 search 三步流程【免费下载链接】faissA library for efficient similarity search and clustering of dense vectors.项目地址: https://gitcode.com/GitHub_Trending/fa/faissfaiss 仓库的 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 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其中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向量维度示例为 256output结果输出目录示例为/checkpoint/marialomeli/offline_faiss/ssnpp需要改成你自己的存储路径index和nprobe分别给出prod与non-prod两组 factory 串如IVF8192,PQ128和 nprobe 值。非集群运行时使用prod组的最后一个值k每个查询取的邻居数示例为 50index_shard_size单个索引分片包含的向量数示例为 50000000query_batch_sizesearch 步骤每批处理的查询向量数示例为 50000000training_sample训练取样的向量数示例为 1572864datasets按数据集名组织每个数据集包含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 定义参数默认值用途--partitionlearnlabSlurm 分区README 示例中写为PARTITION-NAME按你的集群替换--num_nodes1每个 job 的节点数--gpus_per_node8每个节点 GPU 数--cpus_per_task80每个 task 的 CPU 数--tasks_per_node1每节点任务数--logs_dir/checkpoint/marialomeli/offline_faiss/logssubmitit 日志目录建议替换--job_nameoivf集群任务名--nt96检索线程数faiss.omp_set_num_threads两点来自代码的固定行为job 时限在代码里写死为slurm_time70 * 6070 分钟不是命令行参数--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_checkREADME 建议在任何时候都可以跑一致性检查做 sanity checkpython run.py --command consistency_check --config config_ssnpp.yaml --xb ssnpp_1Boffline_ivf.py 的consistency_check依次做四组校验任何一项不满足都会以断言失败中断模板文件{output}/{xb}/{factory}.empty.faissindex存在对每个索引分片取出分片内偏移offset处的SMALL_DATA_SAMPLE10000个库向量以nprobe1搜索自己的 top-100断言每个向量的自身 id 都在结果中search 产物检查每个批次的I...npy/D_approx...npy存在且非空I的列数等于k、行数等于该批查询数用分片索引在线重搜一小批查询与离线结果比对断言匹配比例ratio 0.95且逐查询距离和与在线结果np.allclosertol0.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 都依赖 GPUfaiss.get_num_gpus()、index_cpu_to_all_gpus、knn_gpu纯 CPU 环境不适用该脚本。除三步主路径外代码里还有两个可选命令merge_index把各shard_i合并为单个 ondisk 索引依赖faiss.contrib.ondisk.merge_ondiskevaluate在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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考