轻松计算剪辑嵌入并使用它们构建剪辑检索系统
项目描述
剪辑检索
轻松计算剪辑嵌入并使用它们构建剪辑检索系统。使用 3080 可以在 20 小时内处理 100M 文本+图像嵌入。
- 剪辑客户端允许通过 python 远程查询后端。剪辑客户端笔记本
- 剪辑推理允许您快速(在 3080 上为 1500 个样本/秒)计算图像和文本嵌入
- 剪辑索引从嵌入中构建有效的索引
- 剪辑过滤器允许您使用剪辑索引过滤掉数据
- clip back 使用简单的烧瓶服务托管索引
- clip front 是一个简单的 ui 查询背面。在剪辑检索 ui中查看
- clip end2end 运行 img2dataset、推理、索引,然后前后运行,以使所有这些更容易开始
端到端这使得构建一个简单的语义搜索系统成为可能。有兴趣了解一般语义搜索吗?你可以阅读我关于这个话题的中篇文章。
另请参阅laion5B 和数十亿规模的语义搜索,以了解有关如何将这种规模扩大到数十亿样本的更多信息。
谁在使用剪辑检索?
- cah-prepro预处理 400M 图像+文本爬取的家庭数据集。剪辑检索用于计算 400M 剪辑嵌入和索引
- autofaiss使用剪辑检索来显示使用示例(请参阅那里的多模式笔记本示例)
- afiaka87 openai 演示展示了如何查看 openai 为其 DALL-E 演示发布的 1M 示例
- dzryk 的 antarctic-captions使用 autofaiss 和剪辑推理作为为图像到文本任务生成锚点的方法,取得了巨大成功
安装
点安装剪辑检索
剪辑客户端
ClipClient允许通过 python 远程查询剪辑检索后端。
有关 jupyter 笔记本示例,请参阅ClipClient- 入门笔记本。
API 初始化
在初始化期间,您可以指定一些参数:
backend_url: 后台地址。(必需的)indice_name:指定要使用的索引的名称。(必需的)aesthetic_score:由审美检测器评定的审美分数。默认为9。use_mclip:是否使用多语言版本的CLIP。默认为False。aesthetic_weight:审美分数的权重。默认为0.5modality: 搜索索引中的图像或文本,Multimodal.IMAGE或Multimodal.TEXT.之一 默认为Multimodal.IMAGE。num_images:从 API 返回的图像数量。默认为40。deduplicate: 是否通过图像嵌入对结果进行去重。默认为真。use_safety_model: 是否删除不安全的图片。默认为真。use_violence_detector:是否删除带有暴力的图片。默认为真。
例如,要使用默认参数查询 Laion5B 的托管后端:
from clip_retrieval.clip_client import ClipClient, Modality
client = ClipClient(url="https://knn5.laion.ai/knn-service", indice_name="laion5B")
文字查询
您可以找到与您提供的文本类似的字幕图像。
results = client.query(text="an image of a cat")
results[0]
> {'url': 'https://example.com/kitten.jpg', 'caption': 'an image of a kitten', 'id': 14, 'similarity': 0.2367108941078186}
按图查询
您还可以找到与您提供的图像相似的字幕图像。图像可以通过本地路径或 url 传递。
cat_results = client.query(image="cat.jpg")
dog_results = client.query(image="https://example.com/dog.jpg")
嵌入查询
您还可以找到类似于您提供的剪辑嵌入的字幕图像。
cat_results = client.query(embedding=cat_embedding)
查询图片目录
要增强具有相似文本/图像对的现有数据集,您可以查询图像目录并组合结果。
all_results = [result for result in [client.query(image=image) for image in os.listdir("my-images")]]
with open("search-results.json", "w") as f:
json.dump(all_results, f)
创建数据集
您可以使用保存的 json 结果和工具创建数据集img2dataset。
img2dataset "search-results.json" \
--input_format="json" \
--output_folder="knn_search_dataset" \
--caption_col="caption"
剪辑 end2end
首先选择图像 url 和标题(示例)的数据集,然后运行:
export CUDA_VISIBLE_DEVICES=如果没有足够的 VRAM,您可能希望避免使用 GPU。
wget https://github.com/rom1504/img2dataset/raw/main/tests/test_1000.parquet
clip-retrieval end2end test_1000.parquet /tmp/my_output
然后转到http://localhost:1234并享受在您的图片中搜索的乐趣
--run_back False如果您不想运行后端,请使用
剪辑推断
在 中获取一些图像example_folder,例如通过执行以下操作:
pip install img2dataset
echo 'https://placekitten.com/200/305' >> myimglist.txt
echo 'https://placekitten.com/200/304' >> myimglist.txt
echo 'https://placekitten.com/200/303' >> myimglist.txt
img2dataset --url_list=myimglist.txt --output_folder=image_folder --thread_count=64 --image_size=256
您还可以将与该文件夹中的图像同名的文本文件放在该文件夹中,以获取文本嵌入。
然后运行clip-retrieval inference --input_dataset image_folder --output_folder embeddings_folder
输出文件夹将包含:
- img_emb/
- img_emb_0.npy 包含图像嵌入为 numpy
- text_emb/
- text_emb_0.npy 包含文本嵌入为 numpy
- 元数据/
- metadata_0.parquet 包含图像路径、标题和元数据
这可以扩展到数百万个样本。以 3080 的 1400 个样本/秒的速度,可以在 2 小时内处理 1000 万个样本。
API
clip_inference 将一组文本+图像转换为剪辑嵌入
- input_dataset输入数据集的路径。如果 input_format 是文件,则为文件夹。Bash 大括号模式,例如“{000..150}.tar”(请参阅 https://pypi.org/project/braceexpand/),如果是 webdataset(必需)
- output_folder将保存剪辑嵌入的文件夹以及元数据(必需)
- input_format文件或 webdataset(默认文件)
- cache_path webdataset 的缓存路径(默认None)
- batch_size一次进行推理的项目数(默认256)
- num_prepro_workers进行预处理的进程数(默认8)
- enable_text启用文本处理(默认True)
- enable_image启用图像处理(默认True)
- enable_metadata启用元数据处理(默认False)
- write_batch_size写入批量大小(默认10**6)
- wds_image_key用于 webdataset 中图像的键。(默认jpg)
- wds_caption_key用于 webdataset 中标题的键。(默认txt)
- clip_model要加载的 CLIP 模型(默认ViT-B/32)。将其指定为
"open_clip:ViT-B-32-quickgelu"使用open_clip。 - mclip_model要加载的 MCLIP 模型(默认sentence-transformers/clip-ViT-B-32-multilingual-v1)
- use_mclip如果为 False,则使用 CLIP 执行推理;否则为 MCLIP(默认为False)
- use_jit将 jit 用于剪辑模型(默认True)
- distribution_strategy选择如何分配作业,详见分配部分(默认顺序)
- wds_number_file_per_input_file估计每个 tar 的样本数量,如果使用 wds 而不指定 output_partition_count(默认10000)
- output_partition_count输出分区数(默认None)
- wandb_project要使用的 wandb 项目(默认clip_retrieval)
- enable_wandb是否使用 wandb (默认False )
在 hdfs 上加载/写入文件
- 要从 hdfs 文件夹加载 Web 数据集,请在请求中设置 --input_dataset "pipe:hdfs dfs -cat path_on_hdfs" 而不使用“hdfs://”前缀。
- 要将输出写入 hdfs,请将 --output_hdfs_folder 设置为 hdfs 上以“hdfs://”为前缀的路径
使用 webdataset 格式的 hdfs 查询示例:`clip_inference --input_dataset "pipe:hdfs dfs -cat /myfolder/webdataset/{00000..00010}.tar" --output_folder "hdfs://myfolder/embeddings" --input_format webdataset
在 s3 上加载/写入文件
`clip_inference --input_dataset "pipe:aws s3 cp --quiet s3://myfolder/webdataset/{00000..00010}.tar" --output_folder "s3://myfolder/embeddings" --input_format webdataset
分布式推理
要在多个节点(和多个 gpus)上运行它,请参阅docs/distributed_clip_inference.md上的教程
剪辑索引
剪辑索引将剪辑推理的输出作为输入,并使用autofaiss 从中生成索引
clip-retrieval index --embeddings_folder embeddings_folder --index_folder index_folder
--max_index_memory_usage "16G"选项允许配置索引将消耗的内存量。更多 ram,更好的 knn 召回(默认4G)。--current_memory_available 24G允许控制在创建过程中使用多少内存(默认16G)。--image_subfolder "img_emb"允许为连接到--embeddings_folder选项(默认img_emb)的图像嵌入指定一个子文件夹。--text_subfolder "text_emb"允许为连接到--embeddings_folder选项(默认text_emb)的文本嵌入指定一个子文件夹。--copy_metadata True可以在过程结束时选择是否复制元数据(默认True)。--nb_cores 8允许控制线程数(默认None,将使用所有内核)。
输出是一个文件夹,其中包含:
- image.index 包含图像的 faiss 索引
- text.index 包含文本的 faiss 索引
- 包含镶木地板元数据的元数据文件夹
由于 autofaiss 和 faiss,这可以在几个小时内扩展到数亿个样本。
您可能需要仔细选择要为索引使用多少内存以最大化 knn 召回率。
autofaiss 索引选择 colab可以与autofaiss score_index命令一起帮助检查索引的召回。一般来说,使用更多内存的索引可以获得更好的召回率,因此更接近于幼稚(慢)knn
剪辑过滤器
计算嵌入后,您可能希望通过特定查询过滤掉数据。为此,您可以运行clip-retrieval filter --query "cat" --output_folder "cat/" --indice_folder "indice_folder"
它将在输出文件夹中复制此查询的 100 个最佳图像。使用--num_resultsor--threshold可能有助于优化过滤器
由于快速的 knn 索引,对于较大的 K 值 (100000),它可以实时运行 (<10ms),对于非常大的 K 值,可以在几分钟内运行。
此脚本适用于小型数据集。对于较大的,请查看 [notebook/simple_filter.ipynb]。
剪辑回来
Clip back 是一个简单的 knn 服务后端。如果同时使用 hdf5 和 faiss 内存映射,它只使用剪辑使用的内存,即 4GB。
运行(output_folder 是剪辑索引的输出)
echo '{"example_index": "output_folder"}' > indices_paths.json
clip-retrieval back --port 1234 --indices-paths indices_paths.json
选项:
--use_jit True对剪辑模型使用 jit--clip_model "ViT-B/32"允许选择要使用的剪辑模型--enable_mclip_option True加载 mclip 模型,从而可以使用任何语言进行搜索。--columns_to_return='["url", "image_path", "caption", "NSFW"]允许您指定应从元数据中获取哪些列并由后端返回。在 hdf5 缓存的情况下指定 less 以加快查询速度很有用。--enable_faiss_memory_mapping=True可以传递选项以使用具有内存映射的索引。这将内存使用量减少到零。--enable_hdf5 True可以传递选项以启用元数据的 hdf5 缓存。HDF5 缓存使得使用元数据几乎没有内存使用成为可能。--use_arrow True允许使用箭头而不是 hdf5。应与clip_back_prepro 一起用于非常大的数据集(十亿)--reorder_metadata_by_ivf_index True选项利用 knn ivf 索引结果的数据局部性属性:它按照 IVF 集群的顺序对元数据集合进行排序。这使得元数据检索速度更快成为可能,因为读取访问的是元数据的几个主要是顺序部分,而不是许多非顺序部分。在实践中,这意味着能够在 1 秒内检索 100 万个项目,而没有这种方法只能在 1 秒内检索到 1000 个项目。这将使用第一个图像索引对元数据进行排序。--provide_safety_model True将自动下载并加载安全模型。您需要pip install autokeras可选的依赖项才能使其正常工作。--provide_violence_detector True将加载一个暴力检测器,纸张--provide_aesthetic_embeddings True将加载美学嵌入并允许用户使查询向剪辑空间的更好点移动
这些选项也可以在配置文件中提供,以便为每个索引提供不同的选项。例子:
{
"laion5B": {
"indice_folder": "/mnt/laion5B/prepared_data",
"provide_safety_model": true,
"enable_faiss_memory_mapping": true,
"use_arrow": true,
"enable_hdf5": false,
"reorder_metadata_by_ivf_index": false,
"columns_to_return": ["url", "caption"],
"clip_model": "ViT-L/14",
"enable_mclip_option": false
},
"laion_400m": {
"indice_folder": "/mnt/laion400M/index100",
"provide_safety_model": true,
"enable_faiss_memory_mapping": true,
"enable_hdf5": true,
"use_arrow": false,
"reorder_metadata_by_ivf_index": true,
"enable_mclip_option": true,
"clip_model": "ViT-B/32"
}
}
在以下情况下使用 hdf5 或箭头缓存是一个好主意:
- 您没有足够的内存来将元数据加载到内存中
- 你的磁盘很快(即你有一个ssd)
此时,您有一个在端口 1234 上运行的简单烧瓶服务器,它可以回答以下查询:
/indices-list-> 返回索引列表/knn-service作为输入:
{
"text": "a text query",
"image": "a base64 image",
"image_url": "http://some-url.com/a.jpg",
"modality": "image", // image or text index to use
"num_images": 4, // number of output images
"indice_name": "example_index",
"num_result_ids": 4 // optional, if specified fetch this number of results in total but only num_images with metadata
}
text、image 和 image_url 互斥并返回:
[
{
"image": "base 64 of an image",
"text": "some result text",
"id": 543
},
{
"image": "base 64 of an image",
"text": "some result text",
"id": 782
}
]
如果元数据提供,每个对象也可能包含一个 url 字段。
id 是项目在索引中的位置。它可用于使用 /metadata 端点查询元数据:
{
"indice_name": "example_index",
"ids": [543, 782]
}
返回:
{
"image": "base 64 of an image",
"text": "some result text"
// any other key available in the metadata and specified in columns_to_return cli option
}
num_result_ids/knn-service和的参数/metadata可以一起用于执行大型 knn 查询,然后仅在需要时获取元数据。这样做是有道理的,因为 knn 搜索可以非常有效,这要归功于 knn IVF 索引的强引用局部性,这使得使用大 K 可以快速执行 knn,而当前的元数据 (hdf5) 磁盘实现没有那个属性,因此无法快速处理检索大量随机项目。特别是这可以用来在前端实现无限滚动。
默认情况下,后端也会公开一个前端。默认情况下,该前端将访问该后端,但是您可能需要指定这是通过 http 还是 https 发生的,在这种情况下,请使用选项--default_backend来指定后端 url。--url_column允许为前面指定列 url 的名称
回溯:基准测试和监控
如果使用内存映射索引和元数据,此后端有 50 毫秒的延迟。吞吐量约为 20 次查询/秒。对于高吞吐量,需要使用 grpc 服务器以及用于快速剪辑推断的 GPU,关闭内存映射选项也可以加快请求,但代价是高 ram 使用率。
这个后端还公开了一个普罗米修斯/metrics端点以及一个人类可读的摘要/metrics-summary。这可以(可选)用于设置grafana 仪表板进行监控:
在这个仪表板上可以看到,任何调用中最慢的部分是在图像 url 搜索的情况下通过其 url 获取图像,最多需要 300 毫秒。对于文本查询或图像查询,延迟约为 50ms。以下是指标摘要中的输出示例:
Among 20.0 calls to the knn end point with an average latency of 0.1889s per request, the step costs are (in order):
name description calls average proportion
0 download_time Time spent downloading an url 6 0.3215s 170.2%
1 metadata_get_time Time spent retrieving metadata 20 0.0415s 21.9%
2 knn_index_time Time spent doing a knn on the index 20 0.0267s 14.1%
3 image_clip_inference_time Time spent doing a image clip inference 6 0.0206s 10.9%
4 text_clip_inference_time Time spent doing a text clip inference 14 0.0186s 9.8%
5 image_prepro_time Time spent doing the image preprocessing 6 0.0097s 5.2%
6 text_prepro_time Time spent doing the text preprocessing 14 0.0020s 1.0%
前夹式
Clip front 是一个简单的 UI,可以连接到 clip back 并显示结果。您可以在剪辑检索 ui中使用它
或者你可以自己运行它:
npm install -g clip-retrieval-front
clip-retrieval-front 3005
您也可以使用clip-retrieval front或从 python 包运行它。
发展
对于开发它,去前面npm install然后运行npm start。
用于发展
无论是在本地,还是在gitpod中(在export PIP_USER=false那里做)
设置一个虚拟环境:
python3 -m venv .env
source .env/bin/activate
pip install -e .
运行测试:
pip install -r requirements-test.txt
然后
make lint
make test
您可以使用make black重新格式化代码
python -m pytest -x -s -v tests -k "test_runner"运行特定的测试
如果你想通过python后端或前端使用前端,运行
cd front
npm install
npm run build
cd ..
pip install -e .