语义搜图检索方案

图像检索（包括文搜图和图搜图）是各个行业中常见的一个应用场景。比如在电商场景中，基于以图搜图做相似商品查找；在云相册场景中，基于文搜图来找寻所需的图像素材。

传统基于标签的图像检索方式，即先使用目标检测/图像分类产品（如 Amazon Rekognation）产生图片标签，再存储标签并在检索时匹配标签，存在一些局限性。第一是该方法受限于产品本身支持的标签数量，比如 Amazon Rekognation 支持 3000 个左右的标签，对于标签之外的目标无法检测并检索。第二是标签检索对动词，数量，方位等信息不敏感。诸如“两个孩子在一栋房子前面奔跑”，会出现包含“孩子，房子”标签的图像，但容易忽略“两个，前面，奔跑”等信息。

在本文中，我们将介绍两种基于语义的搜图检索方式以及部署方案，您可以实际测试两种方式在特定应用场景下的搜索率以及搜索的准确性，并选择更加适合的方案使用。

方案一：

利用 CLIP 模型的文本编码器和图像编码器，将图像和文本放入同一维度的潜在向量空间中，找到与文本向量最相近的图像向量以实现基于语义的图像搜索。如下图所示，我们可以看到图像对应的 Image embedding 在向量空间的位置，不同文本会同样被编码至向量空间中，当图像和文字语义最相近时，其 Image embedding 和 Text embedding 的余弦距离最近。

方案二：

利用图生文模型（在本文中使用 vit-gpt2-image-captioning 模型，可替换）先对图像产生对应的文本描述，再将生成的图片描述文本与查询文本放入同一维度的潜在向量空间中，找到与查询文本向量最相近的描述文本向量所对应的图像，以实现基于语义的图像搜索。如下图所示，我们可以看到图像经过 vit-gpt2-image-captioning 模型后，产生 “a black and white dog and a brown and white dog”的描述文本。对于语义最相近的文本，其 Text embedding 之间的余弦距离最近。

背景知识

CLIP 模型介绍

CLIP（Contrastive Language-Image Pre-Training）是由 OpenAI 于 2021 年推出的一种深度学习模型。它主要包含两个组成部分：一个用于图像编码的卷积神经网络（Convolutional Neural Network，CNN）和一个用于文本编码的 Transformer 模型。这两个编码器将图像和文本编码输入至相同的潜在向量空间中，并使得相似的图像和文本在潜在向量空间中的距离更近。

vit-gpt2-image-captioning 模型介绍

vit-gpt2-image-captioning 模型利用 Vision Transformer（例如 ViT、BEiT、DeiT、Swin）对图像进行编码，然后基于预训练的语言模型（例如 RoBERTa、GPT2、BERT、DistilBERT）产生文本，从而实现“看图说话”的能力。

OpenSearch k-NN 介绍

Amazon OpenSearch 支持 k-NN 搜索，可让您在向量空间中通过欧几里得距离或余弦距离查找点之间的“最近邻”。在语义搜图检索方案中，主要包括两个阶段：

构建 k-NN index——在此阶段，将产生的图像向量，或描述文本向量，保存到 OpenSearch 服务中的 k-NN index 中。k-NN 的基本概念是相似数据点存在于嵌入空间中非常接近的位置。例如，文本“一朵红花”、文本“红玫瑰”和“红玫瑰”图像语义相似，因此这些文本和图像在潜在向量空间的位置是彼此接近的。

查询 k-NN index——在此阶段，通过实时机器学习模型对查询文本或者图像进行编码，形成向量。随后使用该文本向量来查询存储在 OpenSearch 服务中的 k-NN index，k-NN index 从潜在向量空间返回相似的向量。例如，如果您的查询文本是“一朵红花”，它将返回红玫瑰图像的向量作为类似的项目。

方案架构

整体方案架构及介绍如下：

通过预训练的 CLIP 模型和 vit-gpt2-image-captioning 模型创建一个用来进行批量和实时推理的 SageMaker 模型。
使用 SageMaker 批量转换作业生成图像 Embedding 或图像描述文本 Embedding 。
使用 SageMaker 实时推理（或无服务器推理）来将图像和文本进行实时编码，产生图像 Embedding 或图像描述文本 Embedding。
使用 Amazon Simple Storage Service（Amazon S3）作为主要存储服务，存储原始图像以及 SageMaker 批量转换作业产生的 Embedding 文件。
使用 OpenSearch 服务来存储向量并查找相关的向量。
使用一个查询函数来编排查询编码并执行 k-NN 搜索。

该方案具有以下优势：

灵活构建。您可以非常方便的替换 CLIP 模型或 vit-gpt2-image-captioning 模型至其他模型，只需要在 SageMaker 上创建相关模型并部署终端节点即可；您也可以替换 OpenSearch 使用其他的向量数据库。
可快速构建大规模生产应用。该方案基于 SageMaker，OpenSearch，S3 等服务构建，生产应用时无需过多担心资源限制及拓展，帮助您在快速进行场景测试后进入大规模生产应用。

方案搭建

您可以在该 github 上找到完整代码。

准备工作

您需要提前创建 SageMaker Studio 与 OpenSearch 集群。由于篇幅原因本文不赘述基础资源创建方法，请参考博客自行搭建。该博客中提供创建相关基础资源以及权限配置的 CloudFormation 模板。

数据集准备

我们使用 Unsplash Dataset Lite 数据集用作测试，其中包含大约两万五千张图像，由超过 250,000 多名摄影师贡献。实际应用请参阅有关 Unsplash 数据集的 License。

进入 SageMaker Studio ，利用以下代码从相关路径下载数据集至本地 “./photos”目录下，并解压缩：

!cd ./photos
!wget http://sbert.net/datasets/unsplash-25k-photos.zip
!unzip unsplash-25k-photos.zip

模型准备

构建 CLIP 模型：

CLIP 模型包括文本编码器和图像编码器，分别用于对图像和文本进行编码。我们也分别构建这两个模型。

首先我们先构建模型压缩包，并上传至 S3 存储桶中，如下图所示：

接下来我们准备创建模型所需要的代码，以下代码均在本地“code”路径下：

准备自定义推理脚本 clip_inference.py。我们在 model_fn 中进行模型加载，在 input_fn 中进行数据读取和预处理，在 predict_fn 定义推理逻辑，并在 output_fn 中进行数据后处理。具体如下：

定义相关的 requirements.txt 文件以允许 SageMaker 在托管模型时安装额外的包：

最后我们基于 PyTorchModel 类创建一个对象来包含模型的 S3 位置信息和推理入口点详细信息。

我们分别创建了 Clip_image_model 用于对图像编码以及 Clip_text_model 用于对文本编码。详见下面的代码：

构建 vit-gpt2-image-captioning 模型：

我们用上述方法构建 vit-gpt2-image-captioning 模型。

首先先打一个假模型压缩包并上传至 S3 存储桶。该压缩包不包含真实的模型文件，我们将在 inference.py 里直接对模型进行加载。

接下来编写同样定义模型的自定义推理脚本 inference.py。

以及定义 requirements.txt 文件以允许 SageMaker 在托管模型时安装额外的包。

最后我们依旧基于 PyTorchModel 类创建一个对象来包含模型的 S3 位置信息和推理入口点详细信息，并创建 image2text_model 模型用于输入图像产生对应文字描述。详见下面的代码：

部署 SageMaker 实时终端节点

现在我们已经创建了模型，接下来我们将基于模型创建 SageMaker 实时终端节点来实现对文本和图像的实时向量编码。我们可以直接使用 model.deploy()的方式来直接部署模型。

部署 CLIP 模型：

部署 vit-gpt2-image-captioning 模型：

模型部署成功后，我们可以在 SageMaker 控制台看到终端节点呈现 “Inservice” 状态。

产生图像和文本的 Embedding

通过 SageMaker Transform Job 实现批处理

基于 SageMaker Transform Job ，我们可以将图像批量处理成 Image embedding，或者将 vit-gpt2-image-captioning 产生的图像描述文本批量处理成 Text Embedding 进行存储。您可以每天定时定量进行处理，该方式仅在批量处理期间对使用机器进行收费，对于成本十分友好。

下文用 Clip_image_model 模型进行示例。代码中针对数据集所在的 S3 存储桶位置，以及期望输出的 Image Embedding 存储位置进行了定义；同时也定义了 Clip_image_transformer，并通过 Clip_image_transformer.transform() 开启批量处理任务。

您可以在 SageMaker 控制台上查看到正在运行或运行完成的批量处理任务。

批量转换作业完成后，我们可以在对应的 S3 存储桶中看到 Embedding 结果，结果以.out 的文件格式保存。

通过 SageMaker 实时推理产生 Embedding

除了批量作业，当然，您也可以实时产生 Embedding，并注入至 OpenSearch 中。下面我们利用 vit-gpt2-image-captioning 模型展示如何实时对描述文本进行编码。

对于 lyStEjlKNSw.jpg 图像，vit-gpt2-image-captioning 模型返回的文字描述是“a black and white dog and a brown and white dog”。

接下来我们利用 CLIP 模型的文本编码器 Clip_text_model 模型来对“a black and white dog and a brown and white dog” 描述文本产生 Embedding 向量进行编码，具体如下：

批量导入 Embedding 至 OpenSearch 中

在准备工作中，我们已经创建了 OpenSearch 集群，下文仅用 CLIP 方案来示例如何将图像 Embedding 批量导入 OpenSearch。

首先我们创建一个 k-NN index，用于存储图片 Embedding 或者图像描述文本 Embedding。

随后，我们将存储于 S3 的图像 Embedding，先导入至本地机器中，再与其对应的图片名称共同导入至 OpenSearch 中。注意，此步仅为 Demo 应用，在实际生产应用中您需要考虑更多效率以及拓展性问题，建议您修改这部分内容再应用于生产。

进行实时语义图像检索

利用 clip-text-model 模型先对查询文本进行实时编码，得到编码向量。

利用查询函数，从 OpenSearch 中进行查询。我们可以指定查询的 Index，当 Index 里面存储的是图像 Embedding 时，查询结果会返回相似度最高的 3 个图像向量；当 Index 里面存储的是图像描述文本 Embedding 时，查询结果会返回相似度最高的 3 个描述文本向量。

根据向量对应的图片名称进行查询，并可视化显示。看到“Two dogs playing in the snow”的查询结果。

更多语义搜图结果展示：

“Two dogs playing in the snow”

A dog in a park

A girl

two girls

A girl with sunflowers

a girl is smiling

a girl who back to us

结语

本文探究了基于 CLIP 模型和图生文模型的语义搜图检索方案，适用于以图搜图做相似商品查找、基于文搜图来找寻所需的图像素材等众多场景。本文所探究的方案具有灵活构建，并可以快速构建大规模生产应用的优势，有助于客户提高生产效率。如果您是亚马逊云科技企业级支持客户，欢迎联系您的 TAM，让我们为您提供专属技术支持。如果您有任何相关的问题或需求，都欢迎随时联系我们进一步交流。

Reference

Radford, Alec, et al. “Learning transferable visual models from natural language supervision.” International conference on machine learning. PMLR, 2021.
https://huggingface.co/docs/transformers/model_doc/vision-encoder-decoder

Kevin Du and Ananya Roy“Implement unified text and image search with a CLIP model using Amazon SageMaker and Amazon OpenSearch Service” https://aws.amazon.com/cn/blogs/machine-learning/implement-unified-text-and-image-search-with-a-clip-model-using-amazon-sagemaker-and-amazon-opensearch-service/

亚马逊AWS官方博客