跳到主要内容

常用Faiss GPU加速索引

Faiss中不少常用的索引类型（Index）都可以在GPU上运行，从而获得极高的加速效果。

CPU / GPU 互操作性​

GPU版本的Index索引同时支持主机（CPU）指针和设备（GPU）指针作为add()和search()方法的输入。如果输入数据已经在与索引相同的GPU上，无需数据拷贝，执行效率最高。如果输入数据在CPU或另一块GPU上，Faiss会自动完成必要的数据拷贝（包括查询结果会自动拷回CPU）。

GPU的Index类型可以直接作为CPU索引的替代品，无需修改业务代码，Faiss会自动处理数据拷贝。例如，GpuIndexFlatL2可直接替代IndexFlatL2，且能显著提升速度。需要注意，GpuIndexIVFFlat和GpuIndexIVFPQ并未实现CPU版本（IndexIVF、IndexIVFFlat、IndexIVFPQ）的全部高阶接口（如直接操作倒排表），但支持大部分常用接口和主流功能。

索引转换​

通过index_gpu_to_cpu、index_cpu_to_gpu和index_cpu_to_gpu_multiple方法实现CPU与GPU之间的索引转换。

转为GPU的相关函数允许传入可选的GpuClonerOptions对象，用于灵活调整GPU上的资源配置。默认配置适用于内存充足的环境。如果显存紧张，可以调整相应参数。

GpuClonerOptions类定义详见github：gpu/GpuClonerOptions.h。

Python中的索引转换​

在Python中，同样可以使用index_gpu_to_cpu、index_cpu_to_gpu和index_cpu_to_gpu_multiple方法。

针对更高级别的接口，可以使用faiss.gpu_wrappers.py中提供的封装：

  • index_cpu_to_all_gpus：将一个CPU索引复制到所有可用GPU上，或通过ngpu=3指定复制到指定数量的GPU。
  • index_cpu_gpu_list：和上面类似，但可以通过gpus=[2, 4, 6]指定GPU的ID列表。
  • index_cpu_to_gpu_multiple_py：可复用GPU资源，从而节约内存，推荐在同时实例化多个GPU索引时使用。示例：
ngpu = 4
resources = [faiss.StandardGpuResources() for i in range(ngpu)]
index1_gpu = faiss.index_cpu_to_gpu_multiple_py(resources, index1)
index2_gpu = faiss.index_cpu_to_gpu_multiple_py(resources, index2)

所有上述方法均支持co=参数，传入GpuMultipleClonerOptions对象。如果只用单GPU，也可以设置ngpu=1。

直接使用Pytorch张量作为输入​

Pytorch的张量可以直接传递给search()和add()方法，具体参考torch_test_contrib_gpu.py。

GpuResources 对象​

所有GPU索引都需要传入一个StandardGpuResources对象（实现了抽象类GpuResources），其中包含每个GPU相关的资源信息，如临时显存（默认为12GB显卡分配约2GB）、cuBLAS句柄和CUDA流等。

Scratch（临时）内存​

GpuResources对象分配的临时scratch内存对高效运行和避免GPU/CPU间频繁同步至关重要。Faiss的GPU代码设计为“零分配”（allocation-free），假定临时态数据都可以放入scratch内存。你可以通过setTempMemory方法调整临时内存大小，甚至设置为0。

在GPU Faiss中，显存大体分为两类：永久分配（如索引本身的数据结构）和临时分配。临时分配利用内存堆栈（stack）机制，堆栈空间用尽时会回退到cudaMalloc。通常应预留约1GB的临时空间，以减少cudaMalloc/Free带来的性能损耗。

如果scratch内存过小，可能会出现由于频繁申请和释放显存导致的显著性能下降。可以通过资源对象查询scratch空间的使用量，并据此合理调整分配。

CUDA流（Stream）​

所有Faiss GPU操作会根据GpuResources设定的stream顺序执行，而不默认使用主stream。

提示

如果与如Pytorch等平台直接交互（数据不经过CPU内存），建议：

  • 通过setDefaultNullStreamAllDevices设置Faiss使用默认stream，或
  • 对不同stream之间做显式同步。

支持的索引类型​

以下索引类型已支持GPU加速：

  • IndexFlat → GpuIndexFlat
  • IndexIVFFlat → GpuIndexIVFFlat
  • IndexIVFScalarQuantizer → GpuIndexIVFScalarQuantizer
  • IndexIVFPQ → GpuIndexIVFPQ

这些索引构造时，需额外传入资源对象、索引存储配置和计算精度（float16/float32）等参数。

用户自定义索引（user indices）有四种存储方式：

存储选项说明
INDICES_64_BIT在GPU上用64位整数存储索引值
INDICES_32_BIT在GPU上用32位整数存储（前提是所有索引能被32位整型表达），最终API会转化为long返回
INDICES_CPU不在GPU上保存索引，而保存在CPU，仅索引本身在GPU，这样可以节省显存，但会带来CPU/GPU之间拷贝与查找的开销
INDICES_IVF仅用于倒排表ID及偏移量的组合场景

float16（半精度浮点数）和float32（单精度浮点数）可分别决定底层存储和一些中间数据的精度和资源消耗。它们均支持主流GPU（Kepler K40、Maxwell及之后产品）。float16通常表现更优秀，能显著降低显存占用，但可能牺牲部分精度。Recall@N（召回率）受影响较小，但距离计算会有差别。对于支持float16运算单元（如Pascal架构）的GPU，能获得更好加速。

例如，Pascal架构在float16模式下的GpuIndexFlat会用Hgemm，而其他架构用SgemmEx。

限制说明​

  • 所有索引的k和nprobe参数需≤2048。
  • GpuIndexIVFPQ每个编码向量允许的码字节数为：1、2、3、4、8、12、16、20、24、28、32、48、56、64、96 字节。
备注

GPU一般更容易遇到内存瓶颈，所以请注意：

  • 当GpuIndexIVFPQ每个编码≥56字节时，需强制使用float16模式（受限于GPU共享内存），例如 48×28×sizeof(float)=4915248 \times 2^8 \times \text{sizeof(float)}=49152 字节；
  • 对于GpuIndexIVFPQ，预计算表可能占用巨大内存，如遇cudaMalloc失败，可关闭预计算表；
  • 倒排表对应的索引可通过indices_options = INDICES_CPU，转存在CPU；
  • 极端显存压力下，倒排表几何扩容可能溢出，此时应先估算索引大小并调用reserveVecs预留空间，从而提升添加速度；
  • StandardGpuResources默认分别为4GB及以下分配512MiB、8GB及以下分配1024MiB、其他GPU分配最多1536MiB的临时空间。如果分配过多可酌情调小，但会影响速度；
  • 添加或查询大批量向量建议分批处理（常用分批大小为8192），具体看bench_gpu_1bn.py示例代码。
提示

从CPU索引用index_cpu_to_gpu转换为GPU索引时，默认的GpuClonerOptions设置是优先速度但会消耗更多内存。

important

磁盘I/O操作尚不支持GPU索引存取。请先用index_gpu_to_cpu将索引转换为CPU版本后再进行持久化存储。

Python中的GPU类​

如果Faiss是用GPU支持编译的，Python接口也能直接使用所有GPU相关类。

详细使用实例见：benchs/bench_gpu_sift1m.py。

多GPU使用说明​

多GPU支持可通过以下方式实现：

  • 使用IndexReplicas，将数据集分别复制到多个GPU，批量并行查询（适合大批量数据，查询吞吐量高）；
  • 使用IndexShards，将数据集分片存储在多块GPU（突破单卡数据量限制，适合大规模数据）。

调用index_cpu_to_gpu_multiple自动实现上述功能，默认创建IndexReplicas。要切换为IndexShards，需传入GpuMultipleClonerOptions对象，并将shard设置为true。

若为倒排索引（IVF），还可通过设置common_ivf_quantizer=true合并量化器以复用计算资源。

性能表现​

important

请参阅 GPU与CPU对比 页面以获取最新性能数据和参考信息。

使用GPU版Faiss通常能比同等CPU实现快5-10倍（详见官方benchmark和性能报告）。在拥有多块GPU的系统中，通过“复制”方式可接近线性加速（如8卡可获6~7倍以上增益）；而数据分片（sharding，N卡各持1/N数据）加速效果略逊于复制。

绝大多数索引（除可能GpuIndexFlat*外）主要受限于全局带宽（global bandwidth）或共享带宽（shared bandwidth），而非计算单元本身。因此和CPU一样，推荐使用批量查询方式（如批量size为8192左右），不要每次只查单条。

需要注意，随着k（最近邻数量）增大，查询效率会下降，特别是大于512时更明显。

支持的GPU设备​

GPU加速依赖CUDA，机器上应有至少一块支持CUDA计算能力3.5及以上（Kepler及之后架构，如K40等）。某些特殊加速功能（如warp shuffles，CC3.0+；读缓存ld.nc/__ldg，CC3.5+）以及float16支持（需CUDA 7.5+编译）涉及较新硬件特性。

官方大量测试了K40、Titan X、M40和P100，部分操作在GTX 970、GeForce GTX 1080上可能出现崩溃（详情见issue #8、issue #34）。欢迎反馈其他平台的运行经验，无论正面或负面。

faiss GPU Index 工作示意图