[Faiss]. IVFPQ ANN 벡터 검색 구현하기

Faiss 설치

pip install faiss-cpu

Faiss 란?

Faiss는 Facebook AI Research에서 개발한 오픈 소스 라이브러리로, 고차원 벡터 검색을 위한 최적화된 인덱싱 및 검색 알고리즘을 제공합니다. Faiss는 대규모 데이터셋에서 빠른 검색 속도와 높은 정확도를 달성할 수 있도록 설계되었습니다.

Faiss의 장점

Faiss는 검색 엔진, 추천 시스템, 이미지 검색 등 다양한 분야에서 활용될 수 있습니다. Faiss의 장점은 대규모 데이터셋에서 빠른 검색 속도와 높은 정확도를 제공하는 것입니다.

경험을 보태서 얘기하면 Faiss를 도입하면서 검색 속도가 수십 배 이상 빨라졌고 벡터가 차지하는 메모리도 줄어서 낮은 스펙의 서버에서도 부담없이 쓸 수 있게 되었다.

하지만 무조건 좋은 것만 있는 것은 아니다.

ANN(Approximate Nearest Neighbor) 이라는 것 자체가 대략적으로 근접한 이웃을 찾아주는 것이기 때문에 "Faiss의 IVFPQ 결과가 정확하다" 보다는 "아마 제일 정확(근사)할 거야"로 해석하는 게 맞다.

속도 및 메모리 보다 정확도가 더 중요한 것이라면 Faiss 사용을 고려해볼 필요가있다.

Faiss IVFPQ 모델 학습시키기

준비물 : Numpy 형식의 Dense vector

import numpy as np
import faiss
from pathlib import Path

class FaissIVFPQ:
    """
    FaissIVFPQ 는 Faiss 라이브러리를 이용하여 IVF + PQ 알고리즘을 구현한 클래스입니다.
    IVF - 벡터를 미리 분류하여 벡터 전체를 검색하지 않고 특정 클러스터에서만 검색하는 방법입니다.
     * IVF를 사용하면 검색 속도를 높일 수 있습니다.
     
     : n_list : 클러스터의 수를 지정합니다.
     : n_probe : 검색할 최대 클러스터의 수를 지정합니다.
    
    PQ -  벡터를 양자화하는 방법입니다.
     * PQ를 사용하면 벡터의 차원을 줄일 수 있습니다.(faiss index의 메모리 사용량 감소)
     : m : 최종 벡터의 중심 수
     : bits : 각 중심에서 벡터를 양자화하는데 사용되는 비트 수를 지정합니다. 
    
    """
    def __init__(self,
                 nlist: int = 256,
                 n_probe: int = 3,
                 m: int = 384,
                 bits : int = 8):
        
        self.nlist = nlist
        self.n_probe = n_probe
        self.m = m
        self.bits = bits
        
    def train(self,
              embedding: np.ndarray,
              name: str,
              export_path: Path):
        """
        IVFPQ 모델을 학습합니다.
        IndexFlatIP : 내적을 사용하여 벡터를 정렬합니다.
        IndexIVFPQ : IVF + PQ 알고리즘을 사용합니다.
        faiss.normalize_L2 : 벡터를 L2 정규화합니다.
         -> Faiss의 결과를 dot product로 비교하기 위해 사용합니다.
        
        : embedding : 학습에 사용할 벡터를 지정합니다.
        : name : 학습된 모델의 이름을 지정합니다.
        : export_path : 학습된 모델을 저장할 경로를 지정합니다.
        
        """
        embedding_dim = embedding.shape[1]
        
        if embedding_dim % self.m != 0:
            raise ValueError(f"embedding_dim({embedding_dim}) % m({self.m}) != 0")
        
        quantizer = faiss.IndexFlatIP(embedding_dim)
        index = faiss.IndexIVFPQ(quantizer, embedding_dim,
                         self.nlist, self.m, self.bits, faiss.METRIC_INNER_PRODUCT)
        index.nprobe = self.n_probe
        
        faiss.normalize_L2(embedding)
        index.train(embedding)
        index.add(embedding)
        
        # 저장
        faiss.write_index(
            index, str(export_path) + f"/{name}_n{self.nlist}_nprobe{self.n_probe}_m{self.m}_b{self.bits}.index")

여기서 METRIC_INNER_PRODUCT 로 한 이유는 결과를 유사도로 뽑고 싶어서 입니다.
결과를 가까운 거리(L2)로 뽑고 싶은 분들은 기본 설정으로 해도됩니다.

Faiss 검색 사용

import numpy as np
import faiss

# 768차원의 랜덤 넘파이 array 생성
random_array = np.random.rand(768)
input_embedding = np.expand_dims(random_array , axis=0) #차원증가로 전체를 한번 감싸줘야함.

index = faiss.read_index("path.index")
D, I = index.search(question_embedding, topk)

"D"는 dinstance를 "I"는 index를 의미함.

Faiss IVFPQ의 평가 방법

규정된 평가 방법은 없지만 아마 이게 가장 적절하지 않을까 싶다.
먼저 Faiss를 적용하지 않은 전체검색 결과 Top N개를 뽑는다.
그 다음 Faiss를 적용한 TopN 과 비교하여 얼마나 교집합이 많은지 체크해보면 된다.

맨 위의 모델이 Faiss 미적용

메모리와 검색속도 정확도를 비교하여 중요함에 따라 모델을 결정하면 된다.