대용량 데이터의 유사도 검색 라이브러리 faiss 사용하기

faiss

faiss는 대용량의 데이터 간의 유사도를 빠르게 계산해주는 유사도 검색 라이브러리입니다. 

 

예를 들어 유사한 단어를 찾고 싶다! 하시면 미리 임베딩된 단어들로 "인덱스"를 생성해 주시고, "검색"을 통해 관련된 단어를 얻으시면 됩니다. 

 

지금은 텍스트의 예로 말씀드렸는데, 이미지 검색에서도 활용 가능합니다. 

 

사용 방법

1. 임베딩된 벡터 준비

2. 인덱스 생성

3. 검색

faiss는 위의 3가지 과정을 거치면 사용하실 수 있습니다. 

 

본격적으로 사용해 보기 전! 

설치를 먼저 해주세요~

• cpu 버전

pip install faiss

• gpu 버전

pip install faiss-gpu

 

1. 임베딩된 벡터 준비

이미지면 이미지를 임베딩한 벡터, 텍스트면 텍스트를 임베딩한 벡터가 필요합니다. 

저는 예시로 SBERT를 이용하도록 하겠습니다. 

from sentence_transformers import SentenceTransformer

# 모델 불러오기
model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1')

# 데이터 임베딩 하기
query_embedding = model.encode(['cat','dog','puppy'])

2. 인덱스 생성

인덱스는 마치 하나의 데이터 베이스를 만들어 준다고 생각하시면 됩니다.

import faiss

# 내적을 이용한 인덱스 생성(Inner Product)
index = faiss.IndexFlatIP(768)

# 유사도 함수로 내적을 사용하는 인덱스 생성
# 각 벡터에 대한 고유한 ID가 있는 경우 이를 이용하여 인덱스를 구성할 수 있게 IndexIDMap을 써줌
# index = faiss.IndexIDMap(faiss.IndexFlatIP(768))

# L2 정규화 진행(cosine similarity)
faiss.normalize_L2(embs)

# 인덱스에 데이터 추가
index.add(embs)

# IndexIDMap 사용시 아래와 같이 추가
# index.add_with_ids(embs, np.array(range(0, len(embs))))

• 인덱스 생성 : 말 그대로 초기 인덱스를 정의해 주는 부분입니다. 모두 입력으로는 차원을 넣어줍니다.
  • IndexFlatIP(dimension) : 내적으로 검색하는 인덱스 생성
  • IndexFlatL2(dimension) : L2거리를 이용하여 검색하는 인덱스 생성
• 정규화 진행 : 이건 필수는 아니고, 코사인 similarity를 이용하기 위해서 추가한 부분입니다. 
  • normalize_L2(embedding) : 임베딩 벡터를 넣어주는 부분으로, L2정규화를 시켜줌
• 데이터 추가 : 생성한 인덱스에 데이터를 넣어주는 부분입니다. 
  • add_with_ids(embedding,ids) : 벡터와 해당 벡터의 고유 ID를 함께 추가하는 방법. IndexIDMap을 사용할 경우 이 함수로 데이터를 넣어줘야함.
  • add(embedding) : 인덱스에 벡터만 추가하는 방법
  • train(embedding) : 인덱스를 학습하는 방법

3. 검색

# 검색할 키워드를 벡터로 바꿔줌
keyword = model.encode(['kitty'])

# L2 정규화를 시켜줌
norm_keyword = faiss.normalize_L2(keyword)

# 검색해줌
d, i = index.search(norm_keyword,1)

• 검색 키워드 임베딩&정규화 : 인덱스에 들어있는 데이터와 유사도 측정을 하기 위해 우리가 검색해줄 키워드를 벡터로 바꿔주는 과정입니다. 저는 위에서 L2정규화를 시켜주었기 때문에, 키워드 또한 L2정규화를 시켜주었습니다. 
• 검색 : 준비해둔 인덱스에 검색하는 부분입니다. 
  • search(keyword_vector, topK) : 임베딩된 벡터와, 몇 개 까지 검색할지 설정해주는 topK를 입력. 결과로는 유사도와 인덱스를 얻음

 

전체 코드 및 실행 결과

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# 모델 불러오기
model = SentenceTransformer('multi-qa-MiniLM-L6-cos-v1')

# 데이터 임베딩 하기
query_embedding = model.encode(['cat','dog','puppy'])

# 내적을 이용한 인덱스 생성(Inner Product)
index = faiss.IndexIDMap(faiss.IndexFlatIP(384))

# L2 정규화 진행(cosine similarity)
faiss.normalize_L2(query_embedding)

# 인덱스에 데이터 추가
index.add_with_ids(query_embedding, np.array(range(0, len(query_embedding))))

# 검색할 키워드를 벡터로 바꿔줌
keyword = model.encode(['kitty'])

# L2 정규화를 시켜줌
faiss.normalize_L2(keyword)

# 검색해줌
d, i = index.search(keyword,1)

실행 결과

[[0.807919]] [[0]]

해석해 보자면 0번째 인덱스를 가지는 단어와 0.8의 distance를 가진다는 것입니다. 

kitty와 제일 가까운건 cat이 맞죠~!

 

인덱스 저장&로드

추가적으로 기본 과정은 아니지만 인덱스를 저장하고, 불러오는 부분도 한번 살펴보겠습니다.

# 인덱스 저장하는 방법
faiss.write_index(index,"YOUR_PATH")

# 인덱스 불러오는 방법
index = faiss.read_index("YOUR_PATH")

• 인덱스 저장 : 말 그대로 위에서 만들었던 인덱스를 저장하는 과정입니다. 
  • write_index(index, file_name) : 인덱스와 파일 이름을 입력으로 줌. 파일은 형식 따로 지정 안해도 됨.
  • read_index(file_name) : 저장한 파일과 동일하게 넣어주면 됨.

 

GPU 사용 설정

faiss에서 gpu설정을 따로 안해주시면 검색시 시간이 매우 느립니다. 

gpu 버전을 설치하시더라도 따로 설정을 안해주시면 cpu기반으로 검색을 합니다. 

아래 코드를 사용하기 위해서는 faiss-gpu 버전이 먼저 설치되어 있어야합니다. 오류가 난다면 이부분 체크 꼭 해주세요!

 

1. Single GPU 사용

import faiss

# GPU 장치를 초기화
res = faiss.StandardGpuResources()

# 인덱스를 생성
index = faiss.IndexFlatL2(d)

# GPU 인덱스로 변환
gpu_index = faiss.index_cpu_to_gpu(res, 0, index)

# 임베딩을 GPU 인덱스에 추가
gpu_index.add(xb)

# 인덱스를 사용하여 검색을 수행
d, i = gpu_index.search(xq, k)

2. Multi GPU 사용

import faiss

# GPU 장치를 초기화
res = faiss.StandardGpuResources()

# 인덱스를 생성
index = faiss.IndexFlatL2(d)

# 여러 개의 GPU에서 인덱스를 생성
ngpus = faiss.get_num_gpus()
co = faiss.GpuMultipleClonerOptions()
co.shard = True
index = faiss.index_cpu_to_all_gpus(index, co=co, ngpu=ngpus)

# 임베딩을 인덱스에 추가
index.add(xb)

# 인덱스를 사용하여 검색을 수행
D, I = index.search(xq, k)