faiss入門: ANN(近似最近傍探索)を動かしてみた

ANN(Approximate Nearest Neighbor)のPythonパッケージである **[faiss](https://github.com/facebookresearch/faiss)を動かしてみました**。いくつかあるANNのPythonパッケージの中でfaissを選んだのには、特に深い理由はありません(たまたま仕事で関係あったから)。あまり時間かけて調べてないので、書いている内容に間違いあるかもしれません。

# 環境
普通にpipでPythonパッケージはインストールし、ローカルのMac PCでJupyterで動かしています。

|種類|バージョン|備考|
|:-:|:-:|:--|
|Python|3.9.13|今どき少し古め|
|faiss-cpu|1.7.4|CPUで動かしいます|
|numpy|1.26.3||
|pandas|2.1.4|Jupyter表示に少しだけ使用|

# 動かした内容
フラットインデックスでの総当りの検索と、"IndexIVFFlat"というインデックスでの近傍検索のインデックス作成時間、検索時間およびリコールを調べています。

## 処理速度結果
まずは結果から。512次元×100万レコードに対する10クエリのTop 10を検索。

### インデックス構築時間
インデックス`IndexIVFFlat`で3000のボロノイ領域(nlist)にすると28秒と少しだけ遅いです。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/2cfc6104-3edc-9be3-7282-89fba1a4aa7c.png)

### 検索時間と再現率
分割しすぎると再現率が非常に悪いですね。少し期待しすぎました。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/134273/85f3e032-9d56-d860-a4de-b847dd4a8bde.png)

## 処理
### 0. Package Import
`time`パッケージは時間計測に使っています。

```python
import time

import numpy as np
import faiss
import pandas as pd
```

### 1. 定数定義

```python
DIMENSION = 512
NUM_NEIGHBOR = 100 # Top N Nearest Neighbors
NUM_QUERY = 10 # Number of Query for search
```

### 2. 乱数の配列生成
以下の配列を0から1までの乱数で生成します。
- vectors: 検索先(インデックス)の配列
- query: 検索元の配列

```python
def get_random_array(dimension, size, seed=1234):
    np.random.seed(seed)
    vectors = np.random.random((size, dimension)).astype('float32') #from 0 to 1
    print(vectors.shape)
    return vectors

vectors = get_random_array(DIMENSION, 1_000_000)
query = get_random_array(DIMENSION, NUM_QUERY)
```

```: Output(各配列の次元)
(1000000, 512)
(10, 512)
```

### 3. インデックス作成
インデックスを作成する関数`create_index`です。
パラメータ`nlist`を渡すと`IndexIVFFlat`のインデックスを作成します。パラメータ`nlist`を渡さないと`IndexFlatL2`というフラットなインデックスを作成します。
細かく調べていませんが、ベクトル空間をボロノイ領域に分割してANNを実現しているようです。パラメータ`nlist`はその分割数。
※「ボロノイ領域」は[「はじパタ全力解説: 第5章 k最近傍法(kNN法)」](https://qiita.com/FukuharaYohei/items/0314f73cebb471deb515)で勉強しました。

実用としては、各レコードにIDが必要なので、`add_with_ids`関数を使って配列とIDを登録しています(コード上では`np.arange`でID生成)。`IndexFlatL2`と`IndexIVFFlat`で少しお作法が異なるので注意が必要です。[「Faiss ID mapping」](https://github.com/facebookresearch/faiss/wiki/Pre--and-post-processing#faiss-id-mapping)参照。

```python
def create_index(vectors, nlist=0):
    index_flat = faiss.IndexFlatL2(vectors.shape[1])   # build the index
    if nlist == 0:
        index = faiss.IndexIDMap2(index_flat)
    else:
        # IndexIVF doesn't need IndexIDMap
        # https://github.com/facebookresearch/faiss/wiki/Pre--and-post-processing#ids-in-the-indexivf
        index = faiss.IndexIVFFlat(index_flat, vectors.shape[1], nlist, faiss.METRIC_L2)
        index.train(vectors)

    index.add_with_ids(vectors, np.arange(vectors.shape[0]))
    # print(f'Trained: {index.is_trained}, Size: {index.ntotal}')
    return index
```

faissは他にも色々なインデックスの種類がありますが、公式Tutorialの[「Faster search
」](https://github.com/facebookresearch/faiss/wiki/Faster-search)に記載のあった`IndexFlatL2`を使っています。


### 4. 検索 
検索の関数です。`IndexFlatL2`でも`IndexIVFFlat`でも検索できるようにしていますが、`IndexIVFFlat`だとパラメータ`nprobe`を渡して、検索対象のボロノイ領域数を指定します。

```python
def search(index, query, k, nprobe=0):
    if nprobe > 0:
        index.nprobe = nprobe
    # 本当は1番目に来る自身のベクトルを除外した方がいいかも
    distances, neighbors = index.search(query, k)
    #print(f'Distance shape: {distances.shape}, Neighbor shape: {neighbors.shape}')
    #print(f'{distances=}')
    #print(f'{neighbors=}')
    return distances, neighbors
```

### 5. フラットインデックス登録と正解データ取得
フラットインデックスを登録して近傍検索の正解データを取得。この正解データから後で再現率を求めます。

```python
def get_truth(vectors, query, k):
    start = time.time()
    index = create_index(vectors) # create FLAT index
    end = time.time()
    print(f'build FLAT: {end-start}, {start}, {end}')
    start = time.time()
    _, neighbors = search(index, query, NUM_NEIGHBOR)
    end = time.time()
    print(f'search: {end-start}, {start}, {end}')
    return neighbors
truth = get_truth(vectors, query, NUM_NEIGHBOR)
print(f'{truth}')
```

インデックス構築に0.4秒、総当り検索に0.15秒かかっています。512次元の100万レコードなので、こんなものでしょう。

```:出力結果
build FLAT: 0.4003911018371582, 1706175708.877467, 1706175709.277858
search: 0.14706897735595703, 1706175709.277965, 1706175709.425034
[[    0 26726 43139 47473 67963 60394 28560  7757 23261 14442  4391 80261
...中略
  34653 49011 58559 24668]]
```

ちなみに検索結果で出力される距離はデフォルトでL2ノルムです。平方根になっていないので注意しましょう。

https://github.com/facebookresearch/faiss/wiki/MetricType-and-distances#metric_l2

https://stackoverflow.com/questions/72548651/faiss-how-to-get-euclidean-distance-between-the-vectors

### 6. `IndexIVFFlat`インデックス登録と検索
`IndexIVFFlat`インデックス登録と検索をして、各処理時間とリコールを一覧出力します。
インデックス登録時には、`nlist`(ボロノイ領域数)を変更。検索時には`nprobe`(検索対象ボロノイ領域数)を変更しています。

```python
results = []
for nlist in [10, 100, 500, 1_000, 2_000, 3_000]:
    start = time.time()
    index = create_index(vectors, nlist) 
    end = time.time()
    build_time = end-start
    print(f'build {nlist=}: {build_time}, {start}, {end}')
    for nprobe in [1, 3, 10, 30, 100]:
        result = {}
        result['build_time'] = build_time
        result['nlist'] = nlist
        start = time.time()
        _, neighbors = search(index, query, NUM_NEIGHBOR, nprobe)
        end = time.time()
        result['nprobe'] = nprobe
        result['search_time'] = end-start
        recalls = [np.count_nonzero(np.isin(neighbors[i], truth[i]) )
                                      for i in range(NUM_QUERY)]
        result['recall'] = sum(recalls) / NUM_NEIGHBOR / NUM_QUERY
        print(f'search {nlist=}, {nprobe=}: {end-start}, {result["recall"]}')
        results.append(result)

pd.DataFrame(results)
```

# 参考リンク
- [ANN Benchmarks](https://ann-benchmarks.com/)
- [aiss解説シリーズ（第一回）基本編](https://crumbjp.hateblo.jp/entry/2021/05/05/Faiss%E8%A7%A3%E8%AA%AC%E3%82%B7%E3%83%AA%E3%83%BC%E3%82%BA%EF%BC%88%E7%AC%AC%E4%B8%80%E5%9B%9E%EF%BC%89): 細かくやる場合にここを参考にしたい