실습 5 — 인덱스의 속도와 정확도 재기

import time
import numpy as np
import pandas as pd
import faiss
import matplotlib.pyplot as plt
plt.rcParams["font.family"] = "Malgun Gothic"            # 그래프의 한글 글꼴
plt.rcParams["axes.unicode_minus"] = False

rng = np.random.default_rng(42)
N, d = 100_000, 384
# 32개의 숨은 방향을 섞어 384차원으로 펼치고 잡음을 조금 더한다(실제 임베딩처럼 구조가 있는 벡터)
X = rng.normal(size=(N, 32)).astype("float32") @ rng.normal(size=(32, d)).astype("float32")
X += 0.05 * np.sqrt(32) * rng.normal(size=(N, d)).astype("float32")
X /= np.linalg.norm(X, axis=1, keepdims=True)           # 길이를 1 로 맞춘다
# 질문 200개: 들어 있는 벡터를 조금 흔든 것
Q = X[rng.choice(N, 200, replace=False)] + 0.02 * rng.normal(size=(200, d)).astype("float32")
Q = (Q / np.linalg.norm(Q, axis=1, keepdims=True)).astype("float32")
print(X.shape, Q.shape, f"{X.nbytes / 1e6:.0f}MB")

flat = faiss.IndexFlatIP(d)
flat.add(X)
t = time.time()
_, truth = flat.search(Q, 10)                            # 모두 견준 결과가 "정확한 상위 10개"
flat_ms = (time.time() - t) / len(Q) * 1000


def recall(found):
    """정확한 상위 10개 가운데 몇 개를 찾았나(질문마다 잰 뒤 평균)"""
    return np.mean([len(set(a) & set(b)) / 10 for a, b in zip(truth, found)])


rows = [{"인덱스": "Flat", "손잡이": "-", "만들기(초)": 0.0, "질문당(ms)": flat_ms, "재현율": 1.0}]
print(f"Flat: 질문 하나에 {flat_ms:.2f}ms")

# IVF: 칸 256개로 나누고, 볼 칸 수(nprobe)를 바꿔 가며 잰다
t = time.time()
ivf = faiss.IndexIVFFlat(faiss.IndexFlatIP(d), d, 256, faiss.METRIC_INNER_PRODUCT)
ivf.train(X)                                             # 칸을 나누는 기준을 배운다
ivf.add(X)
build = time.time() - t
for nprobe in (1, 8, 32, 64):
    ivf.nprobe = nprobe
    t = time.time()
    _, found = ivf.search(Q, 10)
    rows.append({"인덱스": "IVF", "손잡이": f"nprobe {nprobe}", "만들기(초)": build,
                 "질문당(ms)": (time.time() - t) / len(Q) * 1000, "재현율": recall(found)})

# HNSW: 이웃 32개씩 잇는 그물을 만들고, 살펴볼 후보 수(efSearch)를 바꿔 가며 잰다
t = time.time()
hnsw = faiss.IndexHNSWFlat(d, 32)                        # 길이 1 인 벡터라 거리(L2) 순서가 내적 순서와 같다
hnsw.add(X)
build = time.time() - t
for ef in (16, 64, 128):
    hnsw.hnsw.efSearch = ef
    t = time.time()
    _, found = hnsw.search(Q, 10)
    rows.append({"인덱스": "HNSW", "손잡이": f"efSearch {ef}", "만들기(초)": build,
                 "질문당(ms)": (time.time() - t) / len(Q) * 1000, "재현율": recall(found)})

result = pd.DataFrame(rows).round({"만들기(초)": 1, "질문당(ms)": 3, "재현율": 3})
print(result.to_string(index=False))

for name, g in result.groupby("인덱스", sort=False):
    plt.plot(g["질문당(ms)"], g["재현율"], "o-", label=name)       # 손잡이를 올릴수록 오른쪽 위로 간다
plt.xscale("log")
plt.gca().xaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:g}"))
plt.xlabel("질문 하나의 검색 시간(ms, 로그 눈금)")
plt.ylabel("재현율(상위 10개)")
plt.title("벡터 10만 개에서 인덱스별 속도와 정확도")
plt.legend()
plt.grid(alpha=0.3)
plt.show()