Laion-400M 데이터셋
Laion-400M 데이터셋
영어 이미지 캡션과 함께 4억 개의 이미지를 담은 대규모 이미지-텍스트 데이터셋입니다. 이미지·캡션 임베딩과 유사도 점수, 메타데이터를 포함해 ClickHouse에서 근사 최근접 이웃 검색(ANN)을 실습하기에 좋습니다. UDF로 임베딩을 직접 생성하는 방법까지 확인할 수 있어요.
출처: 문서
본문
Laion-400M 데이터셋은 영어 이미지 캡션이 있는 4억 개의 이미지를 포함합니다. Laion은 현재 더 큰 데이터셋도 제공하지만, 다루는 방식은 비슷할 것입니다.
이 데이터셋은 이미지 URL, 이미지와 이미지 캡션 각각의 임베딩, 이미지와 이미지 캡션 사이의 유사도 점수, 그리고 메타데이터(예: 이미지 너비/높이, 라이선스, NSFW 플래그)를 포함합니다. 이 데이터셋으로 ClickHouse의 근사 최근접 이웃 검색을 시연할 수 있습니다.
데이터 준비 (Data preparation)
임베딩과 메타데이터는 원시 데이터의 별도 파일에 저장됩니다. 데이터 준비 단계는 데이터를 내려받고, 파일을 병합하고, 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 # download image embedding
wget --tries=100 https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/text_emb/text_emb_${number}.npy # download text embedding
wget --tries=100 https://deploy.laion.ai/8f83b608504d46bb81708ec86e912220/embeddings/metadata/metadata_${number}.parquet # download metadata
python3 process.py $number # merge files and convert to 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'
# load all files
im_emb = np.load(npy_file)
text_emb = np.load(text_npy)
data = pd.read_parquet(metadata_file)
# combine files
data = pd.concat([data, pd.DataFrame({"image_embedding" : [*im_emb]}), pd.DataFrame({"text_embedding" : [*text_emb]})], axis=1, copy=False)
# columns to be imported into ClickHouse
data = data[['url', 'caption', 'NSFW', 'similarity', "image_embedding", "text_embedding"]]
# transform np.arrays to lists
data['image_embedding'] = data['image_embedding'].apply(lambda x: x.tolist())
data['text_embedding'] = data['text_embedding'].apply(lambda x: x.tolist())
# this small hack is needed because caption sometimes contains all kind of quotes
data['caption'] = data['caption'].apply(lambda x: x.replace("'", " ").replace('"', " "))
# export data as CSV file
data.to_csv(str_i + '.csv', header=False)
# removed raw data files
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)
처음에 인덱스 없이 테이블을 생성하려면:
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 컬럼은 단지 설명을 위한 것이며 스크립트가 고유하지 않은 값으로 채운다는 점에 유의하세요.
브루트포스 벡터 유사도 검색 실행 (Run a brute-force vector similarity search)
브루트포스 근사 벡터 검색을 실행하려면:
SELECT url, caption FROM laion ORDER BY cosineDistance(image_embedding, {target:Array(Float32)}) LIMIT 10
target은 512개 요소의 배열이며 클라이언트 매개변수입니다. 그런 배열을 얻는 편리한 방법은 문서 마지막에 제시됩니다. 지금은 임의의 LEGO 세트 그림의 임베딩을 target으로 실행할 수 있습니다.
결과
┌─url──────...(생략)...─┬─caption──────────────────────────────────────────────────────────────────────────┐
1. │ ... │ LEGO Friends: Puppy Treats & Tricks (41304) │
2. │ ... │ Nouveau LEGO Friends 41334 Andrea s Park Performance 2018 │
3. │ ... │ 3938 LEGO Andreas Bunny House Girls Friends Heartlake Age 5-12 / 62 Pieces New! │
4. │ ... │ LEGO Friends Avonturenkamp Boomhuis - 41122 │
5. │ ... │ LEGO Friends Andrea s Theatershow - 3932 │
6. │ ... │ 41445 - LEGO Friends - Ambulanta clinicii veterinare │
7. │ ... │ LEGO FRIENDS 41336 EMMA S ART CAFÉ │
8. │ ... │ more details on LEGO Friends Stephanie s Friendship Cake Set - 41308. │
9. │ ... │ Lego Friends Gymnast 30400 Polybag 26 pcs │
10. │ ... │ 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.)
(원문의 결과 표에는 각 행의 전체 URL이 포함되어 있습니다.)
벡터 유사도 인덱스로 근사 벡터 유사도 검색 실행 (Run an approximate vector similarity search with a vector similarity index)
이제 테이블에 두 개의 벡터 유사도 인덱스를 정의해 봅시다.
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)
인덱스 생성과 검색의 매개변수 및 성능 고려사항은 문서에 설명되어 있습니다. 위 인덱스 정의는 거리 메트릭으로 "cosine distance"를 사용하는 HNSW 인덱스를 지정하며, "hnsw_max_connections_per_layer" 매개변수는 64, "hnsw_candidate_list_size_for_construction" 매개변수는 256으로 설정됩니다. 이 인덱스는 메모리 사용을 최적화하기 위해 양자화로 반정밀도 brain float(bfloat16)를 사용합니다.
인덱스를 구축하고 실체화(materialize)하려면 다음 문을 실행하세요:
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. │ ... │ LEGO Friends: Puppy Treats & Tricks (41304) │
...(원문과 동일한 10개 결과)...
└───────────...(생략)...─┴──────────────────────────────────────────────────────────────────────────────────┘
10 rows in set. Elapsed: 0.019 sec. Processed 137.27 thousand rows, 24.42 MB (7.38 million rows/s., 1.31 GB/s.)
쿼리 지연 시간이 크게 줄었는데, 그 이유는 최근접 이웃이 벡터 인덱스로 검색되었기 때문입니다. 벡터 유사도 인덱스를 사용한 검색은 브루트포스 검색 결과와 약간 다를 수 있습니다. HNSW 인덱스는 HNSW 매개변수를 신중하게 선택하고 인덱스 품질을 평가하면 (브루트포스 검색과 같은 정확도의) 1에 가까운 recall을 달성할 수 있습니다.
UDF로 임베딩 생성하기 (Creating embeddings with UDFs)
보통 새 이미지나 새 이미지 캡션에 대한 임베딩을 만들고 데이터에서 유사한 이미지/이미지 캡션 쌍을 검색하고 싶을 것입니다. UDF를 사용하면 클라이언트를 떠나지 않고 target 벡터를 만들 수 있습니다. 데이터를 생성할 때와 검색용 새 임베딩을 만들 때 같은 모델을 사용하는 것이 중요합니다. 다음 스크립트들은 데이터셋의 기반이기도 한 ViT-B/32 모델을 활용합니다.
텍스트 임베딩 (Text embeddings)
먼저 다음 Python 스크립트를 ClickHouse 데이터 경로의 user_scripts/ 디렉터리에 저장하고 실행 권한을 부여하세요(chmod +x encode_text.py).
encode_text.py:
#!/usr/bin/python3
#!Note: Change the above python3 executable location if a virtual env is being used.
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 자체는 임베딩 벡터를 계산하고 출력하는 데 몇 초가 걸릴 수 있습니다.
이미지 임베딩 (Image embeddings)
이미지 임베딩도 비슷하게 만들 수 있으며, 로컬에 파일로 저장된 이미지의 임베딩을 생성하는 Python 스크립트를 제공합니다.
encode_image.py
#!/usr/bin/python3
#!Note: Change the above python3 executable location if a virtual env is being used.
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>
검색할 예제 이미지를 가져옵니다:
# get a random image of a LEGO set
$ 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