はじめに

Elasticsearchのベクトル検索機能について、以前から試したいと思っていたので、実際にどんなものか試してみました。

やったこと

fashion-mnistをサンプルデータとしてElasticsearchのベクトル検索を試してみた
インデクシングしてみた
- 784次元6万件をインデクシングできるか
- どのくらいの速さでインデクシングできるか
検索してみた
- どのくらいの速度出るのか
- 精度どんなもんか（コサイン類似度なのでここで試す必要はないかもですが）

以下、やったこと

実験用Elasticsearchをローカルに立てる

Install Elasticsearch with Dockerをみてやると良いですが、以下のコマンドでローカルにElasticsearch（Docker ver.)が起動すると思うので見なくても大丈夫かもしれません。

# Dokcerイメージをpullする
docker pull docker.elastic.co/elasticsearch/elasticsearch:7.3.2

# ポート開けたり必要な環境変数を設定してDocke起動する
docker run -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" docker.elastic.co/elasticsearch/elasticsearch:7.3.2

http://localhost:9200/ にアクセスして起動していることを確認できたらOK。

画像をベクトル化してElasticsearchにインデクシング

何も考えずにそのままベクトルとしてインデクシングする

Fashion-MNISTで得られるベクトルをそのままの形で突っ込むといった内容。

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import FashionMNIST

# データセットをダウンロード
fmnist_data = FashionMNIST('./data/fashion-mnist', train=True, download=True, transform=transforms.ToTensor())
data_loader = DataLoader(fmnist_data, batch_size=4, shuffle=False)
data = data_loader.dataset

# Elasticsearchにインデクシング
es = Elasticsearch("http://localhost:9200/")

# インデックスの作成と、マッピングの設定

mapping = {
    "images": {
        "properties": {
            # 画像のインデックス
            "image_index":{
                "type": "integer"
            },
            # targetのクラスを入れておく（入れなくても良い）
            "category": {
                "type": "integer"
            },
            # Fashion-MNISTは28*28の画像なので784次元
            # ちなみに、公式でサポートしている上限は1024
            "image_vector": {
                "type": "dense_vector",
                "dims": 784
            }
        }
    }
}

# settingsを指定してインデックスを作成
es.indices.create(index='raw-images')

# 作成したインデックスのマッピングを指定
es.indices.put_mapping(index='raw-images', doc_type='images', body=mapping, include_type_name=True)

# インデクシング用の関数
def _load_data(index_, type_, data, batch_num):
    for i, row in enumerate(data):
        image, label = row
        body = {
            "image_index": batch_num + i,
            "category": int(label),
            # ここややこしいけれど普通のfloatに変換しているだけ
            # numpy系はそのまま突っ込めないっぽい
            "image_vector": list(map(lambda x: x.item(), image.numpy().flatten()))
        }
        yield {"_index": index_, "_type": type_, "_source": body}

# インデクシング（60000件は1分かからないくらいでインデクシング終了）
batch_size = 1000
dataset_with_class = list(zip(data.data, data.targets))
for i in range(math.ceil(len(dataset_with_class) / batch_size)):
    slice_data = dataset_with_class[i*batch_size:i*batch_size + batch_size]
    helpers.bulk(es, _load_data("raw-images", "images", slice_data, i*batch_size))

検索を投げる

input_image_index = 0
image, label = dataset_with_class[input_image_index]

# 検索
query_vector = list(map(lambda x: x.item(), image.numpy().flatten()))

# 画像インデックス取得
body = {
    "query": {
        "script_score": {
            "query": {"match_all": {}},
            "script": {
                "source": "cosineSimilarity(params.query_vector, doc['image_vector']) + 1.0",
                "params": {"query_vector": query_vector}
            }
        }
    },
    "_source": {"includes": ["image_index", "category"]},
    "from": 0,
    "size": 10
}

res = es.search(index="raw-images", body=body)

以下のような結果が得られる。

{'took': 292,
 'timed_out': False,
 '_shards': {'total': 1, 'successful': 1, 'skipped': 0, 'failed': 0},
 'hits': {'total': {'value': 10000, 'relation': 'gte'},
  'max_score': 2.0,
  'hits': [{'_index': 'raw-images',
    '_type': 'images',
    '_id': '0rgjOW0B8Qw73ee7lR9P',
    '_score': 2.0,
    '_source': {'image_index': 0, 'category': 9}},
   {'_index': 'raw-images',
    '_type': 'images',
    '_id': 'SbgjOW0B8Qw73ee72oR1',
    '_score': 1.9564186,
    '_source': {'image_index': 25719, 'category': 9}},
   ....

実際どんなもんか実験

インデックスや検索にかかる時間

実験に使ったDockerの根本設定

Dockerは8GBのメモリ制限かけてましたが、実行中8GBでべたっと張り付いていたので、ふるふる使ってる感じだと思います。

計測

インデクシング60,000件で30秒〜1分くらいかと思います。（雑）

クエリ実行ですが、上記のハイパー雑クエリ（784次元6万件対象）で平均0.2663秒って感じでした。

ちょい遅いかなという感じがしますが、クエリを事前に何らかの条件で絞り込めば、実際のサービスでも問題なく使えるかなと思います。
速度がそんなにシビアに求められないサービスなら普通に使えそうな感じですね。

検索結果

実際に検索してみた結果を以下に示します。
基本的にinput画像と、上位5件とかそんな感じでやっていきます。

1. Ankle boot の例

Input

Output

2. Bag

Input

Output

3. Pullover

Input

Output

4. T-shirt/top

Input

Output

... どれも良い感じに検索できていますね！

まとめ

最初は次元削減（PCAとかt-SNEとか線形判別分析とか）やろうと思ってたんですが、そのまま突っ込んで予想以上にパフォーマンス出たので「まあいっか」となってしまいました。それくらいESのベクトル検索は使い勝手が良いと思います。
本来なら100次元とかに絞り込んで使うと思いますし、割と良いのではないかと思ってます。

コサイン類似度の結果も予期した通りの内容で割と良い感じに使えるんだろうなと思いました。
（もしかすると将来的に消えちゃうかもしれないらしく、消えちゃうと悲しいなと思ったりするのでした。あと調べてないので分からないですが、X-Pack との関連が気になります。）

参考資料