依据 HuggingFace Datasets 和 Transformers 的图画类似性查找

经过本文，你将学习运用 Transformers 构建图画类似性查找体系。找出查询图画和潜在候选图画之间的类似性是信息检索体系的一个重要用例，例如反向图画查找 (即找出查询图画的原图)。此类体系试图回答的问题是，给定一个查询图画和一组候选图画，找出候选图画中哪些图画与查询图画最类似。

咱们将运用 datasets 库，因为它无缝支持并行处理，这在构建体系时会派上用场。

尽管这篇文章运用了依据 ViT 的模型 (nateraw/vit-base-beans) 和特定的 (Beans) 数据集，但它能够扩展到其他支持视觉模态的模型，也能够扩展到其他图画数据集。你能够尝试的一些闻名模型有:

Swin Transformer
ConvNeXT
RegNet

此外，文章中介绍的办法也有或许扩展到其他模态。

要研究完整的图画类似度体系，你能够参考这个 Colab Notebook。

咱们怎么界说类似性？

要构建这个体系，咱们首要需求界说咱们想要怎么核算两个图画之间的类似度。一种广泛盛行的做法是先核算给定图画的稠密表征 (即嵌入 (embedding))，然后运用余弦类似性衡量 (cosine similarity metric) 来确认两幅图画的类似程度。

在本文中，咱们将运用 “嵌入” 来表示向量空间中的图画。它为咱们供给了一种将图画从高维像素空间 (例如 224 224 3) 有含义地压缩到一个低得多的维度 (例如 768) 的好办法。这样做的首要优点是减少了后续进程中的核算时间。

基于 Hugging Face Datasets 和 Transformers 的图像相似性搜索

核算嵌入

为了核算图画的嵌入，咱们需求运用一个视觉模型，该模型知道怎么在向量空间中表示输入图画。这种类型的模型一般也称为图画编码器 (image encoder)。

咱们运用 AutoModel 类来加载模型。它为咱们供给了一个接口，能够从 HuggingFace Hub 加载任何兼容的模型 checkpoint。除了模型，咱们还会加载与模型关联的处理器 (processor) 以进行数据预处理。

from transformers import AutoFeatureExtractor, AutoModel
model_ckpt = "nateraw/vit-base-beans"
extractor = AutoFeatureExtractor.from_pretrained (model_ckpt)
model = AutoModel.from_pretrained (model_ckpt)

本例中运用的 checkpoint 是一个在 beans 数据集上微调过的 ViT 模型。

这儿或许你会问一些问题:

Q1: 为什么咱们不运用 AutoModelForImageClassification？

这是因为咱们想要取得图画的稠密表征，而 AutoModelForImageClassification 只能输出离散类别。

Q2: 为什么运用这个特定的 checkpoint？

如前所述，咱们运用特定的数据集来构建体系。因此，与其运用通用模型 (例如在 ImageNet-1k 数据集上练习的模型)，不如运用运用已针对所用数据集微调过的模型。这样，模型能更好地了解输入图画。

注意你还能够运用经过自监督预练习取得的 checkpoint, 不必得由有监督学习练习而得。事实上，假如预练习妥当，自监督模型能够取得令人形象深刻的检索功能。

现在咱们有了一个用于核算嵌入的模型，咱们需求一些候选图画来被查询。

加载候选图画数据集

后面，咱们会构建将候选图画映射到哈希值的哈希表。在查询时，咱们会运用到这些哈希表，详细评论的评论稍后进行。现在，咱们先运用 beans 数据集中的练习集来获取一组候选图画。

from datasets import load_dataset
dataset = load_dataset ("beans")

以下展示了练习会集的一个样本:

该数据集的三个 features 如下:

dataset ["train"].features
>>> {'image_file_path': Value (dtype='string', id=None),
 'image': Image (decode=True, id=None),
 'labels': ClassLabel (names=['angular_leaf_spot', 'bean_rust', 'healthy'], id=None)}

为了使图画类似性体系可演示，体系的整体运转时间需求比较短，因此咱们这儿只运用候选图画数据会集的 100 张图画。

num_samples = 100
seed = 42
candidate_subset = dataset ["train"].shuffle (seed=seed).select (range (num_samples))

寻觅类似图片的进程

下图展示了获取类似图画的基本进程。

略微拆解一下上图，咱们分为 4 步走:

从候选图画 (candidate_subset) 中提取嵌入，将它们存储在一个矩阵中。
获取查询图画并提取其嵌入。
遍历嵌入矩阵 (进程 1 中得到的) 并核算查询嵌入和当时候选嵌入之间的类似度得分。咱们一般保护一个类似字典的映射，来保护候选图画的 ID 与类似性分数之间的对应联系。
依据类似度得分进行排序并回来相应的图画 ID。最后，运用这些 ID 来获取候选图画。

咱们能够编写一个简单的东西函数用于核算嵌入并运用 map() 办法将其作用于候选图画数据集的每张图画，以有效地核算嵌入。

import torch
def extract_embeddings (model: torch.nn.Module):
    """Utility to compute embeddings."""
    device = model.device
    def pp (batch):
        images = batch ["image"]
        # `transformation_chain` is a compostion of preprocessing
        # transformations we apply to the input images to prepare them
        # for the model. For more details, check out the accompanying Colab Notebook.
        image_batch_transformed = torch.stack (
            [transformation_chain (image) for image in images]
        )
        new_batch = {"pixel_values": image_batch_transformed.to (device)}
        with torch.no_grad ():
            embeddings = model (**new_batch).last_hidden_state [:, 0].cpu ()
        return {"embeddings": embeddings}
    return pp

咱们能够像这样映射 extract_embeddings():

device = "cuda" if torch.cuda.is_available () else "cpu"
extract_fn = extract_embeddings (model.to (device))
candidate_subset_emb = candidate_subset.map (extract_fn, batched=True, batch_size=batch_size)

接下来，为便利起见，咱们创建一个候选图画 ID 的列表。

candidate_ids = []
for id in tqdm (range (len (candidate_subset_emb))):
    label = candidate_subset_emb [id]["labels"]
    # Create a unique indentifier.
    entry = str (id) + "_" + str (label)
    candidate_ids.append (entry)

咱们用包括所有候选图画的嵌入矩阵来核算与查询图画的类似度分数。咱们之前现已核算了候选图画嵌入，在这儿咱们仅仅将它们会集到一个矩阵中。

all_candidate_embeddings = np.array (candidate_subset_emb ["embeddings"])
all_candidate_embeddings = torch.from_numpy (all_candidate_embeddings)

咱们将运用余弦类似度来核算两个嵌入向量之间的类似度分数。然后，咱们用它来获取给定查询图画的类似候选图画。

def compute_scores (emb_one, emb_two):
    """Computes cosine similarity between two vectors."""
    scores = torch.nn.functional.cosine_similarity (emb_one, emb_two)
    return scores.numpy ().tolist ()
def fetch_similar (image, top_k=5):
    """Fetches the`top_k`similar images with`image`as the query."""
    # Prepare the input query image for embedding computation.
    image_transformed = transformation_chain (image).unsqueeze (0)
    new_batch = {"pixel_values": image_transformed.to (device)}
    # Comute the embedding.
    with torch.no_grad ():
        query_embeddings = model (**new_batch).last_hidden_state [:, 0].cpu ()
    # Compute similarity scores with all the candidate images at one go.
    # We also create a mapping between the candidate image identifiers
    # and their similarity scores with the query image.
    sim_scores = compute_scores (all_candidate_embeddings, query_embeddings)
    similarity_mapping = dict (zip (candidate_ids, sim_scores))
    # Sort the mapping dictionary and return `top_k` candidates.
    similarity_mapping_sorted = dict (
        sorted (similarity_mapping.items (), key=lambda x: x [1], reverse=True)
    )
    id_entries = list (similarity_mapping_sorted.keys ())[:top_k]
    ids = list (map (lambda x: int (x.split ("_")[0]), id_entries))
    labels = list (map (lambda x: int (x.split ("_")[-1]), id_entries))
    return ids, labels

执行查询

经过以上预备，咱们能够进行类似性查找了。咱们从 beans 数据集的测试会集选取一张查询图画来查找:

test_idx = np.random.choice (len (dataset ["test"]))
test_sample = dataset ["test"][test_idx]["image"]
test_label = dataset ["test"][test_idx]["labels"]
sim_ids, sim_labels = fetch_similar (test_sample)
print (f"Query label: {test_label}")
print (f"Top 5 candidate labels: {sim_labels}")

结果为:

Query label: 0
Top 5 candidate labels: [0, 0, 0, 0, 0]

看起来咱们的体系得到了一组正确的类似图画。将结果可视化，如下:

进一步扩展与结论

现在，咱们有了一个可用的图画类似度体系。但实际体系需求处理比这多得多的候选图画。考虑到这一点，咱们目前的程序有不少缺陷:

假如咱们按原样存储嵌入，内存需求会迅速增加，尤其是在处理数百万张候选图画时。在咱们的比如中嵌入是 768 维，这即使对大规模体系而言或许也是相对比较高的维度。
高维的嵌入对检索部分涉及的后续核算有直接影响。

假如咱们能以某种方式降低嵌入的维度而不影响它们的含义，咱们依然能够在速度和检索质量之间坚持良好的折衷。本文附带的 Colab Notebook 完成并演示了怎么经过随机投影 (random projection) 和位置灵敏哈希 (locality-sensitive hashing，LSH) 这两种办法来取得折衷。

Datasets 供给与 FAISS 的直接集成，进一步简化了构建类似性体系的进程。假设你现已提取了候选图画的嵌入 (beans 数据集) 并把他们存储在称为 embedding 的 feature 中。你现在能够轻松地运用 dataset 的 add_faiss_index() 办法来构建稠密索引:

dataset_with_embeddings.add_faiss_index (column="embeddings")

树立索引后，能够运用 dataset_with_embeddings 模块的 get_nearest_examples() 办法为给定查询嵌入检索最近邻:

scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples (
    "embeddings", qi_embedding, k=top_k
)

该办法回来检索分数及其对应的图画。要了解更多信息，你能够检查官方文档和这个 notebook。

在本文中，咱们快速入门并构建了一个图画类似度体系。假如你觉得这篇文章很风趣，咱们强烈建议你依据咱们评论的概念继续构建你的体系，这样你就能够更加熟悉内部作业原理。

还想了解更多吗？以下是一些或许对你有用的其他资源:

Faiss: 高效类似性查找库
ScaNN: 高效向量类似性查找
在移动使用程序中集成图画查找引擎

英文原文: hf.co/blog/image-…

译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，作业方向为 transformer-family 模型在各模态数据上的使用及大规模模型的练习推理。

审校、排版: zhongdongy (阿东)

基于 Hugging Face Datasets 和 Transformers 的图像相似性搜索

依据 HuggingFace Datasets 和 Transformers 的图画类似性查找

咱们怎么界说类似性？

核算嵌入

加载候选图画数据集

寻觅类似图片的进程

执行查询

进一步扩展与结论

相关文章

写作选题太难？五个标准帮你找准方向

云原生周刊：开源“赢了”，但它可持续吗？

TECH WORKS: 如何打造实用工程师的职业生涯

我的 2022：「无」

作者信息