소개

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 예제

입력

Output

2. Bag

입력

Output

3. Pullover

입력

Output

4. T-shirt/top

입력

Output


... 모두 좋은 느낌으로 검색할 수 있습니다!

요약

처음에는 차원 삭감(PCA라든지 t-SNE라든지 선형 판별 분석이라든지) 하려고 했는데, 그대로 돌진해 예상 이상으로 퍼포먼스 나왔기 때문에 「뭐 어쩐지」가 되어 버렸습니다. 그 정도 ES의 벡터 검색은 사용하기 편리하다고 생각합니다.
본래라면 100차원이라든지 좁혀 사용한다고 생각하고, 비교적 좋은 것이 아닐까 생각합니다.

코사인 유사도의 결과도 예상한 대로의 내용으로 비교적 좋은 느낌으로 사용할 수 있을 것이라고 생각했습니다.
(어쩌면 장래에 사라질지도 모르는 것 같고, 사라지면 슬프다고 생각하거나 하기도 했습니다. 나중에 조사하지 않았기 때문에 모릅니다만, X-Pack 와의 관련이 신경이 쓰입니다.)

참고 자료

ELASTICSEARCH에서 분산 표현을 사용한 유사한 문서 검색

pytorch 입문【MNIST, FASHION-MNIST】