FAISS GPU 사용 — 인덱스 가속과 변환

FAISS GPU 사용 — 인덱스 가속과 변환

FAISS의 가장 유용한 인덱스 중 상당수가 GPU에서 구현돼 있어요. IndexFlatL2 대신 GpuIndexFlatL2를 쓰면 별다른 코드 변경 없이 큰 가속을 얻을 수 있습니다.

출처: https://github.com/facebookresearch/faiss/wiki/Faiss-on-the-GPU

CPU/GPU 상호 운용

GPU Index는 add()와 search()의 입력으로 호스트 포인터와 디바이스 포인터 모두를 받아줘요. 입력이 이미 인덱스와 같은 GPU에 있다면 복사 없이 가장 빠르게 실행되고, 그렇지 않으면 CPU→GPU 복사가 일어나고 결과는 다시 복사돼 돌아옵니다.

GPU 인덱스는 CPU 인덱스를 기대하는 코드의 드롭인 대체재로 쓰기 좋아요. 다만 GpuIndexIVFFlat과 GpuIndexIVFPQ는 CPU 버전(IndexIVFFlat, IndexIVFPQ)의 API 일부(예: 역파일 직접 조작)를 의도적으로 지원하지 않지만, 기본 Index API 대부분과 중요한 인덱싱 기능은 모두 지원합니다.

변환 함수

GPU로 가거나 GPU에서 내려올 때 쓰는 함수는 index_gpu_to_cpu, index_cpu_to_gpu, index_cpu_to_gpu_multiple이에요. GPU로 보내는 두 함수는 선택적으로 GpuClonerOptions 객체를 받아 GPU 저장 방식을 조절할 수 있습니다. 메모리가 넉넉하지 않을 때 주로 조정해요.

Python에서는 더 높은 수준의 편의 래퍼도 있습니다.

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)
  • index_cpu_to_all_gpus: CPU 인덱스를 가용 GPU 전부로 복제 (ngpu=3으로 개수 지정 가능)
  • index_cpu_gpu_list: 같은 동작에 GPU ID 목록을 추가로 받음 (gpus=[2, 4, 6])
  • index_cpu_to_gpu_multiple_py: GPU 인덱스 여러 개를 만들 때 리소스를 공유해 메모리 낭비를 줄임

이런 함수들은 공통으로 co= 인자(GpuMultipleClonerOptions)를 받고, GPU 수를 1로 설정하면 단일 GPU 관리에도 쓸 수 있어요. 또 search()와 add()에 PyTorch 텐서를 직접 입력할 수도 있습니다.

더 알아보기