FAISS 시작하기 — 첫 인덱스와 검색

FAISS 시작하기 — 첫 인덱스와 검색

FAISS는 고정 차원 d의 벡터 모음(보통 수십~수백 차원)을 다뤄요. 이 벡터들은 행렬로 저장되는데, FAISS는 32비트 부동소수점 행렬만 사용한다는 점이 포인트입니다.

출처: https://github.com/facebookresearch/faiss/wiki/Getting-started

데이터 준비

두 개의 행렬이 필요해요. xb 는 인덱스에 넣고 검색할 데이터베이스 벡터, xq 는 최근접 이웃을 찾고자 하는 질의 벡터입니다. 아래는 d=64차원, 데이터베이스 100,000개, 질의 10,000개의 예시예요.

import numpy as np
d = 64                           # dimension
nb = 100000                      # database size
nq = 10000                       # nb of queries
np.random.seed(1234)             # make reproducible
xb = np.random.random((nb, d)).astype('float32')
xb[:, 0] += np.arange(nb) / 1000.
xq = np.random.random((nq, d)).astype('float32')
xq[:, 0] += np.arange(nq) / 1000.

Python에서는 행렬을 항상 numpy 배열로 표현합니다. dtype은 반드시 float32여야 해요.

인덱스 만들고 벡터 추가

FAISS는 Index 객체를 중심으로 동작해요. 가장 단순한 버전인 IndexFlatL2는 데이터베이스 벡터에 대해 무차별(brute-force) L2 거리 검색을 수행합니다. 인덱스를 만들 때는 벡터의 차원 d를 알려줘야 해요.

import faiss                   # make faiss available
index = faiss.IndexFlatL2(d)   # build the index
print(index.is_trained)
index.add(xb)                  # add vectors to the index
print(index.ntotal)

IndexFlatL2는 훈련 단계를 건너뛸 수 있어요. 인덱스에는 add와 search 두 연산을 수행할 수 있고, 상태 변수 is_trained(훈련 필요 여부)와 ntotal(인덱스된 벡터 수)도 확인할 수 있습니다.

k-최근접 이웃 검색

기본 검색은 각 질의 벡터에 대해 데이터베이스에서 k개의 최근접 이웃을 찾는 거예요. 결과로 거리순으로 정렬된 이웃 ID 행렬 I와 그에 해당하는 제곱 거리 행렬 D를 돌려줍니다. 먼저 데이터베이스 벡터 일부를 검색해, 최근접 이웃이 자기 자신인지 sanity check을 해볼게요.

k = 4                          # we want to see 4 nearest neighbors
D, I = index.search(xb[:5], k) # sanity check
print(I)
print(D)
D, I = index.search(xq, k)     # actual search
print(I[:5])                   # neighbors of the 5 first queries
print(I[-5:])                  # neighbors of the 5 last queries

출력을 보면 각 질의의 최근접 이웃이 실제로 벡터 자신의 인덱스이고, 그때 거리는 0이에요. 한 행 안에서도 거리가 점점 커지는 걸 볼 수 있어요. 벡터 첫 번째 성분에 값을 더해 데이터를 첫 축 방향으로 늘렸기 때문에, 앞쪽 질의의 이웃도 데이터셋 앞부분에 몰려 있는 것을 확인할 수 있습니다.

더 알아보기