GPU対応インデックスの作成と転送

概要

FAISS では、CPU 上で作成したインデックスを GPU に転送して高速化する方法と、最初から GPU 上にインデックスを構築する方法の両方が用意されています。ここでは「GPU 対応インデックスの作成と転送」の手順・API・ベストプラクティスを体系的に説明します。

---

1. 前提：GPU 版 FAISS のインストールと CUDA 環境

---

2. GPU リソース (StandardGpuResources) の初期化

python
import faiss
res = faiss.StandardGpuResources()      # 一度だけ生成
res.setTempMemory(256 * 1024 * 1024)    # 必要に応じて一時メモリを調整 (既定は 512–1536 MB)

StandardGpuResources は scratch メモリ・cuBLAS ハンドル・CUDA ストリームを保持します。サイズが小さすぎると cudaMalloc が頻発し性能が低下します。github.com

---

3. インデックスの作成と転送パターン

3.1 CPU で作成 ➜ GPU へ転送

1. CPU インデックスを構築／訓練

   python
   d = 768
   nlist = 4096
   quantizer = faiss.IndexFlatL2(d)
   cpu_index = faiss.IndexIVFPQ(quantizer, d, nlist, m=16, nbits=8)
   cpu_index.train(train_vectors)            # 訓練が必要な場合
   

2. 転送

   python
   gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
   

   • 複数 GPU へ複製: index_cpu_to_all_gpus

   • 指定リストへ複製: index_cpu_gpu_list(gpus=[0,2])

   • リソース共有: index_cpu_to_gpu_multiple_py(resources, cpu_index)
     github.com

3.2 GPU で直接作成

python
config = faiss.GpuIndexIVFFlatConfig()
config.device = 0
config.useFloat16 = True        # メモリ削減 & 高速化
gpu_index = faiss.GpuIndexIVFFlat(res, d, nlist, faiss.METRIC_L2, config)
gpu_index.train(train_vectors)  # 訓練も GPU 上で実行

GPU で直接訓練すれば二重コピーを防げます。

3.3 転送後の利用

• 追加・検索: gpu_index.add(xb) / gpu_index.search(xq, k)
  入力が PyTorch tensor などで既に GPU 上にある場合はホストコピーが省けます。github.com

• パラメータ調整: gpu_index.nprobe, gpu_index.indices_options, reserveVecs() などを CPU 版と同様に変更可能。

---

4. GPU ➜ CPU への逆転送と永続化

GPU インデックスは直接ディスクに保存できません。永続化する場合は

python
cpu_back = faiss.index_gpu_to_cpu(gpu_index)
faiss.write_index(cpu_back, "index.faiss")

と一度 CPU に戻します。github.com

---

5. 複数 GPU の活用

さらに common_ivf_quantizer=True で複数 GPU で量子化器を共有しメモリを節約できます。github.com

---

6. 代表的な GPU 対応インデックスと制限

制限値 (k, nprobe ≤ 2048 など) とメモリ上限に注意。github.com

---

7. ベストプラクティスと落とし穴

1. バッチ処理: add() / search() は 8 k～16 k 件程度をまとめると GPU メモリとスループットのバランスが良い。

2. メモリ不足: indices_options=faiss.INDICES_CPU でインデックス ID を CPU 側に置く、または useFloat16 を活用。

3. I/O: 保存前に必ず index_gpu_to_cpu()。

4. ストリーム同期: 異なる CUDA ストリームを使う他ライブラリ (PyTorch など) と併用する際は setDefaultNullStreamAllDevices() または明示的同期を実施。

5. リソース共有: 複数 GPU インデックス生成時は同じ StandardGpuResources インスタンスを再利用してメモリを節約。

---

8. 最小実装例（CPU → GPU 転送パターン）

python
import faiss, numpy as np

d = 128
xb = np.random.random((100_000, d)).astype('float32')
xq = np.random.random((10, d)).astype('float32')

# 1) CPU インデックス
cpu_index = faiss.IndexFlatL2(d)

# 2) GPU へ転送
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)   # ★転送

# 3) データ追加 & 検索
gpu_index.add(xb)
D, I = gpu_index.search(xq, k=5)
print(I[:3], D[:3])

（index_cpu_to_gpu による転送例）unfoldai.com

---

まとめ

• 転送 API: index_cpu_to_gpu／index_gpu_to_cpu を中心に覚える

• リソース管理: StandardGpuResources と GpuClonerOptions を適切に設定

• パフォーマンス: バッチ処理・float16・複数 GPU レプリカで最大化

• 保存: GPU から直接ファイルへは書き出せないので必ず CPU に戻す

これらを押さえれば、大規模ベクトル検索を GPU で高速に処理しつつ、運用面のトラブルも回避できます。