使用 VectorDB & OpenAI CLIP 构建大规模图像搜索：从 120 小时到 1 小时，从 $$$ 到 $

从给定数据集中搜索图像是一个重要的应用，例如，找到所有我在海滩上的照片。人们一直在使用 AI 技术来实现这一目标，例如语义图像搜索。

在这篇博客文章中，我们分享了以高度分布式且成本效益高的方式构建大规模图像搜索数据库的经验——将端到端生成时间从 120 小时缩短到 1 小时，将成本从 $231 降低到 $46.2。我们将批处理数据处理功能架构在多个云平台上。

Screenshot 2025-01-30 at 2.53.30 PM.png

理解语义图像搜索

传统的图像搜索依赖于元数据，如文件名、标签或手动添加的描述。如果你曾试图在你的照片库中找到“山上的日落”，你会知道这有多么受限——除非你一丝不苟地给每张照片打了标签，否则你就束手无策了。

语义搜索通过理解图像中的实际内容和含义来解决这个问题。它的工作原理如下

向量嵌入：神经网络（如 OpenAI CLIP）将图像转换为高维向量——通常是 512 到 1024 维。这些向量捕获语义特征：颜色、形状、对象，甚至抽象概念。
相似度匹配：当你搜索时，你的查询（文本或图像）被转换到相同的向量空间。系统随后找到其向量与你的查询向量“最接近”的图像，通常使用余弦相似度或欧氏距离。
神经网络理解：这种神奇之处在于神经网络的训练过程。像 CLIP 这样的模型是在数百万个图像-文本对上训练的，学习将相关概念映射到向量空间中的彼此靠近的位置。这意味着搜索“牵狗散步的人”能够找到相关的图像，即使这些图像从未被标记过这些词语。

例如，当你搜索“山上的日落”时

查询文本“山上的日落” → 向量 [0.1, 0.8, -0.3, ...]
每个数据库图像 → 向量 [0.2, 0.7, -0.4, ...]
系统返回与查询向量最相似的图像

语义图像搜索如何工作

在此示例中，我们使用 OpenAI CLIP 的简化代码示例，CLIP 是一种嵌入大小为 512 的双编码器架构，由一个 视觉编码器（ViT）和一个 文本编码器（Transformer）组成。

为什么不使用像 DeepSeek-Janus 这样的 VLM？ CLIP 的训练方法是将文本和图像的损失函数结合起来应用于训练集，这涉及显式的文本-图像匹配目标；而一些 VLM 可能依赖于不同的训练目标（例如，掩码标记预测或图像字幕生成）。如果没有强大的对比对齐，学习到的嵌入对于直接检索任务来说可能区分度较低。

使用 CLIP 进行搜索

使用 CLIP 构建搜索系统时

生成图像嵌入：为了使图像数据集可搜索，我们需要使用 CLIP 模型为数据集中的每张图像生成嵌入向量，并将这些嵌入存储在一个数据库中，称为 向量数据库。

def process_images(image_paths: List[str], clip_model: Any) -> List[torch.Tensor]:
    embeddings = []
    for path in image_paths:
        image = preprocess_image(path) # Resize to 224x224, normalize
        with torch.no_grad():
            image_embedding = clip_model.encode_image(image)
            embeddings.append(normalize(image_embedding))
    return embeddings

查询处理：对于每个输入查询，将使用相同的 CLIP 模型为输入生成嵌入。

def process_query(query: Union[str, bytes], clip_model: Any) -> torch.Tensor:
    if is_text_query(query):
        # Text query
        text = clip.tokenize(query)
        with torch.no_grad():
            vector_embeddings = clip_model.encode_text(text)
    else:
        # Image query
        image = preprocess_image(query)
        with torch.no_grad():
            vector_embeddings = clip_model.encode_image(image)

    return normalize(vector_embeddings)

向量搜索：然后计算输入向量嵌入与数据库中图像嵌入之间的距离，并选择 topk 图像作为搜索结果。

def semantic_search(query_vector: torch.Tensor, image_vectors: torch.Tensor, k: int=10):
    similarities = torch.matmul(
        query_vector,
        torch.stack(image_vectors).T
    )
    top_k = torch.topk(similarities, k)
    return top_k.indices

如何构建包含大量图像的向量数据库？

上述步骤看起来相当简单明了，但当我们需要处理百万规模的数据集时，每个阶段的复杂性都会迅速增加，尤其是由于 AI 模型对 GPU 计算的高需求。

在本博客文章的其余部分，我们将分享如何以高度分布式且成本效益高的方式构建大规模向量数据库的经验——将端到端生成时间从 120 小时缩短到 1 小时，将成本从 $231 降低到 $46.2。我们将批处理数据处理功能架构在多个云平台上。请参阅下图，这是我们管道的总体设计，包含三个阶段

生成图像嵌入
向量数据库入库
为查询提供向量数据库服务。

Screenshot 2025-01-28 at 2.16.59 PM.png

图：构建图像搜索数据库分为三个阶段：生成图像嵌入、向量数据库入库和处理查询。

步骤 1：生成图像嵌入

使用大型 AI 模型为数据集中的图像生成嵌入是最耗时且计算成本最高的步骤。一个典型的云端 GPU，例如 Nvidia L4，使用 openclip 每秒可以处理 7 张图像。使用一个 GPU 计算 100 万张图像需要 40 多个小时；计算包含 50 亿张图像的最大的 Laion-5B 数据集需要超过 23 年！

扩展规模

加速生成过程的自然选择是使用大量配备多个 GPU 的机器并行生成嵌入。想象一下，如果你可以使用 100 台配备类似 GPU 的机器：处理时间将缩短到只需 1 小时！虽然这是个直观的想法，但在现实中，GPU 资源

稀缺：云服务提供商在单个区域的 GPU 容量有限，限制了用户可获得的 GPU 数量。
昂贵：配备 GPU 的机器可能比仅配备 CPU 的机器昂贵 100 倍以上。

为了解决稀缺性问题，利用跨多个区域和云平台的不同类型的 GPU 会非常有帮助，因为它整合了来自多个资源池的可用 GPU。由于嵌入生成不需要节点之间的通信，我们可以将计算自由地分布在全球范围内。

对于昂贵的 GPU，我们可以利用云上的竞价实例（Spot offerings），这些实例通常比普通按需实例便宜 2-3 倍，尽管它们可能随时被抢占，即工作负载可能被中断。

Screenshot 2025-01-28 at 2.34.02 PM.png

图：使用 SkyPilot 进行图像向量计算的示意图。

启动作业

我们使用 SkyPilot 在多个区域或云平台上启动并行嵌入生成作业，并混合使用按需实例和竞价实例，以获得最佳 GPU 可用性并降低成本。具体来说，我们使用 SkyPilot 托管作业，这些作业会自动跨多云/区域获取 GPU，并从故障和竞价实例抢占中恢复作业。。生成的嵌入存储在 SkyPilot 托管的桶中，采用 Apache Parquet 格式，这使得无需重新处理即可恢复或合并部分结果。

我们将 GPU 资源指定为一组备选方案，因此 SkyPilot 可以获取其中任何一个来运行后续指定 CLIP 的嵌入生成作业。我们将数据集静态划分为子集，并将每个子集分配给一个单独的作业。SkyPilot 负责确保每个作业完成其分区，或通过重试作业从失败中恢复。

下面是一个简化配置示例（完整的 SkyPilot YAML 文件可以在这里找到）

name: clip-batch-compute-vectors

resources:
  accelerators: {T4:1, L4:1}
  memory: 32+
  use_spot: true # we use Spot VMs to save costs

envs:
  START_IDX: ${START_IDX}
  END_IDX: ${END_IDX}

file_mounts:
  /output:
    name: my-bucket-for-embedding-output
    mode: MOUNT

setup: |
  pip install numpy==1.26.4
  pip install torch==2.5.1 torchvision==0.20.1 ftfy regex tqdm  
	...

run: |
  python scripts/compute_vectors.py \
    --output-dir /output \
    --start-idx ${START_IDX} \
    --end-idx ${END_IDX} ...

我们使用一个简单的脚本启动所有作业，该脚本将分区分配给每个作业

DATA_SIZE = 1_000_000
NUM_WORKERS = 40
per_worker_size = DATA_SIZE // NUM_WORKERS
task = sky.Task.from_yaml('task.yaml')   

def start_job(start_idx, end_idx):
    sky.jobs.launch(
        task, 
        envs={
            "START_IDX": str(start_idx), 
            "END_IDX": str(end_idx)
        }, 
        detach_run=True
    )

with concurrent.futures.ThreadPoolExecutor() as executor:
    for start_idx in range(0, DATA_SIZE, per_worker_size):
        end_idx = min(DATA_SIZE, start_idx + per_worker_size)
        executor.submit(start_job, start_idx, end_idx)

使用 SkyPilot 提交所有嵌入生成作业后，我们发现这些作业成功地在不同区域并行运行，如图所示，包括 eastus2、westus3 等。

Screenshot 2025-02-04 at 10.32.12 PM.png

步骤 2：向量数据库入库

为所有图像生成嵌入向量后，我们将向量入库到向量数据库中。向量数据库构建内部结构以支持存储和快速查询。

在此步骤中，我们使用 SkyPilot 启动一个云虚拟机，从 SkyPilot 托管桶中的 Apache Parquet 文件收集（嵌入、图像）对，构建向量数据库，并将生成的向量数据库存储到另一个桶中（SkyPilot 的 yaml 文件在此）。

不同的向量数据库具有不同的底层存储机制和权衡（参阅各种向量数据库）。在我们的示例中，我们使用 ChromaDB。

Screenshot 2025-02-02 at 1.31.46 AM.png

步骤 3：提供向量数据库服务

向量数据库只有在你能够查询它时才有用。我们现在通过一个 API 端点来提供数据库服务，其他应用（或你的本地客户端）可以调用该端点执行语义搜索。我们为向量数据库设计了一个简单的搜索网站。在此查看演示。详细的 SkyPilot YAML 文件可以在这里找到。

Screenshot 2025-02-03 at 9.28.53 AM.png

理解语义图像搜索#

语义图像搜索如何工作#

使用 CLIP 进行搜索#

如何构建包含大量图像的向量数据库？#

步骤 1：生成图像嵌入#

步骤 2：向量数据库入库#

步骤 3：提供向量数据库服务#

阅读更多#