Laion-400M 数据集 | ClickHouse Docs

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.arrays 转换为列表
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.)

使用向量相似度索引执行近似向量相似度搜索

现在，我们来在该表上定义两个向量相似度索引。

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。该索引使用半精度脑浮点数 (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 创建嵌入向量

通常需要为新图像或新的图像说明创建嵌入向量，并在数据中搜索相似的图像/图像说明对。我们可以使用 UDF 直接在客户端内创建 target 向量，而无需离开客户端。务必使用同一个模型来生成数据以及用于搜索的新嵌入向量。以下脚本使用 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');

第一次运行会比较慢，因为需要加载模型，但之后再次运行就会很快。然后，我们可以将输出复制到 SET param_target=... 中，这样就能轻松编写查询。或者，也可以直接将 encode_text() 函数作为 cosineDistance 函数的参数使用：

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>

下载一张用于搜索的示例图片：

# 获取一张随机的乐高积木图片
$ wget http://cdn.firstcry.com/brainbees/images/products/thumb/191325a.jpg

然后执行以下查询，为上述图片生成 embedding：

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