Laion-400M データセット

Laion-400M データセットには、英語の画像キャプションが付いた4億枚の画像が含まれています。現在では、Laion はさらに大規模なデータセットも提供していますが、扱い方はほぼ同じです。

このデータセットには、画像の URL、画像と画像キャプションそれぞれの埋め込み、画像と画像キャプションの類似度スコアに加え、メタデータ (画像の幅/高さ、ライセンス、NSFW フラグなど) が含まれています。このデータセットを使って、ClickHouse の近似最近傍探索を実演できます。

データの準備

生データでは、埋め込みとメタデータはそれぞれ別のファイルに保存されています。データ準備の手順では、データをダウンロードし、ファイルを結合して、 CSV に変換し、ClickHouse にインポートします。この処理には、次の download.sh スクリプトを使用できます。

number=${1}
if [[ $number == '' ]]; then
    number=1
fi;
wget --tries=100 https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/img_emb/img_emb_${number}.npy          # 画像の埋め込みをダウンロード
wget --tries=100 https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/text_emb/text_emb_${number}.npy        # テキストの埋め込みをダウンロード
wget --tries=100 https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/metadata/metadata_${number}.parquet    # メタデータをダウンロード
python3 process.py $number # ファイルをマージしてCSVに変換

スクリプト process.py は次のように定義されています。

import pandas as pd
import numpy as np
import os
import sys

str_i = str(sys.argv[1])
npy_file = "img_emb_" + str_i + '.npy'
metadata_file = "metadata_" + str_i + '.parquet'
text_npy =  "text_emb_" + str_i + '.npy'

# すべてのファイルを読み込む
im_emb = np.load(npy_file)
text_emb = np.load(text_npy) 
data = pd.read_parquet(metadata_file)

# ファイルを結合する
data = pd.concat([data, pd.DataFrame({"image_embedding" : [*im_emb]}), pd.DataFrame({"text_embedding" : [*text_emb]})], axis=1, copy=False)

# ClickHouseにインポートするカラム
data = data[['url', 'caption', 'NSFW', 'similarity', "image_embedding", "text_embedding"]]

# np.arrayをリストに変換する
data['image_embedding'] = data['image_embedding'].apply(lambda x: x.tolist())
data['text_embedding'] = data['text_embedding'].apply(lambda x: x.tolist())

# captionに様々な種類の引用符が含まれる場合があるため、このハックが必要
data['caption'] = data['caption'].apply(lambda x: x.replace("'", " ").replace('"', " "))

# データをCSVファイルとしてエクスポートする
data.to_csv(str_i + '.csv', header=False)

# 生データファイルを削除する
os.system(f"rm {npy_file} {metadata_file} {text_npy}")

データ準備パイプラインを開始するには、次を実行します。

seq 0 409 | xargs -P1 -I{} bash -c './download.sh {}'

データセットは 410 個のファイルに分割されており、各ファイルには約 100 万行が含まれています。より小さいデータの一部で作業したい場合は、seq 0 9 | ... のように範囲を調整してください。

(上記の Python スクリプトは非常に遅く (ファイルごとに約 2〜10 分) 、大量のメモリを消費し (ファイルごとに 41 GB) 、生成される CSV ファイルも大きい (各 10 GB) ため、注意してください。十分な RAM がある場合は、-P1 の値を増やして並列度を上げてください。それでも遅すぎる場合は、より良いインジェスト手順を検討してください。たとえば、.npy ファイルを Parquet に変換し、その後のすべての処理を ClickHouse で行う方法があります。)

テーブルを作成する

まずは索引なしでテーブルを作成するには、次を実行します。

CREATE TABLE laion
(
    `id` Int64,
    `url` String,
    `caption` String,
    `NSFW` String,
    `similarity` Float32,
    `image_embedding` Array(Float32),
    `text_embedding` Array(Float32)
)
ENGINE = MergeTree
ORDER BY id

CSVファイルをClickHouseにインポートするには、次の手順に従います。

INSERT INTO laion FROM INFILE '{path_to_csv_files}/*.csv'

id カラムはあくまで説明用であり、スクリプトによって重複する値が設定される点に注意してください。

ブルートフォースでベクトル類似度検索を実行する

ブルートフォースで近似ベクトル検索を実行するには、次を実行します。

SELECT url, caption FROM laion ORDER BY cosineDistance(image_embedding, {target:Array(Float32)}) LIMIT 10

target は 512 要素の配列で、クライアントパラメータです。このような配列を簡単に取得する方法は、記事の最後で紹介します。ここではひとまず、ランダムな LEGO セット画像の埋め込みを target として実行できます。

結果

    ┌─url───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┬─caption──────────────────────────────────────────────────────────────────────────┐
 1. │ https://s4.thcdn.com/productimg/600/600/11340490-9914447026352671.jpg                                                                                                                         │ LEGO Friends: Puppy Treats & Tricks (41304)                                      │
 2. │ https://www.avenuedelabrique.com/img/uploads/f20fd44bfa4bd49f2a3a5fad0f0dfed7d53c3d2f.jpg                                                                                                     │ Nouveau LEGO Friends 41334 Andrea s Park Performance 2018                        │
 3. │ http://images.esellerpro.com/2489/I/667/303/3938_box_in.jpg                                                                                                                                   │ 3938 LEGO Andreas Bunny House Girls Friends Heartlake Age 5-12 / 62 Pieces  New! │
 4. │ http://i.shopmania.org/180x180/7/7f/7f1e1a2ab33cde6af4573a9e0caea61293dfc58d.jpg?u=https%3A%2F%2Fs.s-bol.com%2Fimgbase0%2Fimagebase3%2Fextralarge%2FFC%2F4%2F0%2F9%2F9%2F9200000049789904.jpg │ LEGO Friends Avonturenkamp Boomhuis - 41122                                      │
 5. │ https://s.s-bol.com/imgbase0/imagebase/large/FC/5/5/9/4/1004004011684955.jpg                                                                                                                  │ LEGO Friends Andrea s Theatershow - 3932                                         │
 6. │ https://www.jucariicucubau.ro/30252-home_default/41445-lego-friends-ambulanta-clinicii-veterinare.jpg                                                                                         │ 41445 - LEGO Friends - Ambulanta clinicii veterinare                             │
 7. │ https://cdn.awsli.com.br/600x1000/91/91201/produto/24833262/234c032725.jpg                                                                                                                    │ LEGO FRIENDS 41336 EMMA S ART CAFÉ                                               │
 8. │ https://media.4rgos.it/s/Argos/6174930_R_SET?$Thumb150$&amp;$Web$                                                                                                                             │ more details on LEGO Friends Stephanie s Friendship Cake Set - 41308.            │
 9. │ https://thumbs4.ebaystatic.com/d/l225/m/mG4k6qAONd10voI8NUUMOjw.jpg                                                                                                                           │ Lego Friends Gymnast 30400 Polybag 26 pcs                                        │
10. │ http://www.ibrickcity.com/wp-content/gallery/41057/thumbs/thumbs_lego-41057-heartlake-horse-show-friends-3.jpg                                                                                │ lego-41057-heartlake-horse-show-friends-3                                        │
    └───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┴──────────────────────────────────────────────────────────────────────────────────┘

10 rows in set. Elapsed: 4.605 sec. Processed 100.38 million rows, 309.98 GB (21.80 million rows/s., 67.31 GB/s.)

ベクトル類似度索引を使って近似ベクトル類似度検索を実行する

それでは、テーブルに 2 つのベクトル類似度索引を定義します。

ALTER TABLE laion ADD INDEX image_index image_embedding TYPE vector_similarity('hnsw', 'cosineDistance', 512, 'bf16', 64, 256)
ALTER TABLE laion ADD INDEX text_index text_embedding TYPE vector_similarity('hnsw', 'cosineDistance', 512, 'bf16', 64, 256)

索引の作成と検索に関するパラメータおよびパフォーマンス上の考慮事項は、ドキュメントで説明されています。上記の索引定義では、距離指標として"コサイン距離"を使用する HNSW 索引を指定しており、パラメータ "hnsw_max_connections_per_layer" は 64、パラメータ "hnsw_candidate_list_size_for_construction" は 256 に設定されています。この索引では、メモリ使用量を最適化するため、量子化に半精度の brain float (bfloat16) を使用します。

索引を構築して実体化するには、次のステートメントを実行します：

ALTER TABLE laion MATERIALIZE INDEX image_index;
ALTER TABLE laion MATERIALIZE INDEX text_index;

索引の構築と保存には、行数や HNSW 索引のパラメータによって、数分から数時間かかることがあります。

ベクトル検索を実行するには、同じクエリをもう一度実行するだけです。

SELECT url, caption FROM laion ORDER BY cosineDistance(image_embedding, {target:Array(Float32)}) LIMIT 10

結果

    ┌─url───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┬─caption──────────────────────────────────────────────────────────────────────────┐
 1. │ https://s4.thcdn.com/productimg/600/600/11340490-9914447026352671.jpg                                                                                                                         │ LEGO Friends: Puppy Treats & Tricks (41304)                                      │
 2. │ https://www.avenuedelabrique.com/img/uploads/f20fd44bfa4bd49f2a3a5fad0f0dfed7d53c3d2f.jpg                                                                                                     │ Nouveau LEGO Friends 41334 Andrea s Park Performance 2018                        │
 3. │ http://images.esellerpro.com/2489/I/667/303/3938_box_in.jpg                                                                                                                                   │ 3938 LEGO Andreas Bunny House Girls Friends Heartlake Age 5-12 / 62 Pieces  New! │
 4. │ http://i.shopmania.org/180x180/7/7f/7f1e1a2ab33cde6af4573a9e0caea61293dfc58d.jpg?u=https%3A%2F%2Fs.s-bol.com%2Fimgbase0%2Fimagebase3%2Fextralarge%2FFC%2F4%2F0%2F9%2F9%2F9200000049789904.jpg │ LEGO Friends Avonturenkamp Boomhuis - 41122                                      │
 5. │ https://s.s-bol.com/imgbase0/imagebase/large/FC/5/5/9/4/1004004011684955.jpg                                                                                                                  │ LEGO Friends Andrea s Theatershow - 3932                                         │
 6. │ https://www.jucariicucubau.ro/30252-home_default/41445-lego-friends-ambulanta-clinicii-veterinare.jpg                                                                                         │ 41445 - LEGO Friends - Ambulanta clinicii veterinare                             │
 7. │ https://cdn.awsli.com.br/600x1000/91/91201/produto/24833262/234c032725.jpg                                                                                                                    │ LEGO FRIENDS 41336 EMMA S ART CAFÉ                                               │
 8. │ https://media.4rgos.it/s/Argos/6174930_R_SET?$Thumb150$&amp;$Web$                                                                                                                             │ more details on LEGO Friends Stephanie s Friendship Cake Set - 41308.            │
 9. │ https://thumbs4.ebaystatic.com/d/l225/m/mG4k6qAONd10voI8NUUMOjw.jpg                                                                                                                           │ Lego Friends Gymnast 30400 Polybag 26 pcs                                        │
10. │ http://www.ibrickcity.com/wp-content/gallery/41057/thumbs/thumbs_lego-41057-heartlake-horse-show-friends-3.jpg                                                                                │ lego-41057-heartlake-horse-show-friends-3                                        │
    └───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┴──────────────────────────────────────────────────────────────────────────────────┘

10行のセット。経過時間: 0.019秒。処理済み: 137.27千行、24.42 MB（7.38百万行/秒、1.31 GB/秒）

ベクトル索引を使って最近傍を取得したことで、クエリのレイテンシは大幅に低下しました。ベクトル類似度索引を使用するベクトル類似度検索では、総当たり検索の結果とわずかに異なる結果が返される場合があります。 HNSW パラメータを慎重に選定し、索引の品質を評価することで、HNSW 索引はリコールを 1 に近い値 (総当たり検索と同等の精度) まで高められる可能性があります。

UDFを使用した埋め込みの作成

通常は、新しい画像や新しい画像キャプションの埋め込みを作成し、データ内で類似する画像／画像キャプションのペアを検索したいものです。クライアントを離れることなく target ベクトルを作成するには、UDF を使用できます。データの作成時と検索用の新しい埋め込みの作成時には、同じモデルを使用することが重要です。以下のスクリプトでは、データセットでも使用されている ViT-B/32 モデルを利用します。

テキスト埋め込み

まず、次の Python スクリプトを ClickHouse のデータパス内の user_scripts/ ディレクトリに保存し、実行可能にします (chmod +x encode_text.py) 。

encode_text.py:

#!/usr/bin/python3
#!注意: 仮想環境を使用している場合は、上記のpython3実行ファイルのパスを変更してください。
import clip
import torch
import numpy as np
import sys

if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load("ViT-B/32", device=device)
    for text in sys.stdin:
        inputs = clip.tokenize(text)
        with torch.no_grad():
            text_features = model.encode_text(inputs)[0].tolist()
            print(text_features)
        sys.stdout.flush()

次に、ClickHouseのサーバー設定ファイルで <user_defined_executable_functions_config>/path/to/*_function.xml</user_defined_executable_functions_config> として参照されている場所に、encode_text_function.xml を作成します。

<functions>
    <function>
        <type>executable</type>
        <name>encode_text</name>
        <return_type>Array(Float32)</return_type>
        <argument>
            <type>String</type>
            <name>text</name>
        </argument>
        <format>TabSeparated</format>
        <command>encode_text.py</command>
        <command_read_timeout>1000000</command_read_timeout>
    </function>
</functions>

これで、次のように簡単に使えます。

SELECT encode_text('cat');

初回の実行はmodelの読み込みがあるため遅くなりますが、2回目以降は高速に実行されます。その後、出力を SET param_target=... にコピーすれば、クエリを簡単に記述できます。あるいは、encode_text() 関数を cosineDistance 関数のargumentとして直接使用することもできます：

SELECT url
FROM laion
ORDER BY cosineDistance(text_embedding, encode_text('a dog and a cat')) ASC
LIMIT 10

encode_text() UDF 自体が、埋め込みベクトルを計算して出力するまでに数秒かかる場合があることに注意してください。

画像の埋め込み

画像の埋め込みも同様に作成できます。ローカルのファイルとして保存された画像の埋め込みを生成できる Python スクリプトも用意しています。

encode_image.py

#!/usr/bin/python3
#!注意: 仮想環境を使用している場合は、上記のpython3実行可能ファイルのパスを変更してください。
import clip
import torch
import numpy as np

from PIL import Image
import sys

if __name__ == '__main__':
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load("ViT-B/32", device=device)
    for text in sys.stdin:
        image = preprocess(Image.open(text.strip())).unsqueeze(0).to(device)
        with torch.no_grad():
            image_features = model.encode_image(image)[0].tolist()
            print(image_features)
        sys.stdout.flush()

encode_image_function.xml

<functions>
    <function>
        <type>executable_pool</type>
        <name>encode_image</name>
        <return_type>Array(Float32)</return_type>
        <argument>
            <type>String</type>
            <name>path</name>
        </argument>
        <format>TabSeparated</format>
        <command>encode_image.py</command>
        <command_read_timeout>1000000</command_read_timeout>
    </function>
</functions>

検索に使用するサンプル画像を取得します：

# LEGO セットのランダムな画像を取得する
$ wget http://cdn.firstcry.com/brainbees/images/products/thumb/191325a.jpg

次に、次のクエリを実行して上の画像の埋め込みを生成します：

SELECT encode_image('/path/to/your/image');

検索クエリ全体は次のとおりです：

SELECT
    url,
    caption
FROM laion
ORDER BY cosineDistance(image_embedding, encode_image('/path/to/your/image')) ASC
LIMIT 10