基於 Hugging Face Datasets 和 Transformers 的影像相似性搜尋

基於 HuggingFace Datasets 和 Transformers 的影像相似性搜尋

透過本文，你將學習使用 ? Transformers 構建影像相似性搜尋系統。找出查詢影像和潛在候選影像之間的相似性是資訊檢索系統的一個重要用例，例如反向影像搜尋 (即找出查詢影像的原圖)。此類系統試圖解答的問題是，給定一個查詢影像和一組候選影像，找出候選影像中哪些影像與查詢影像最相似。

我們將使用 ? datasets 庫，因為它無縫支援並行處理，這在構建系統時會派上用場。

儘管這篇文章使用了基於 ViT 的模型 (nateraw/vit-base-beans) 和特定的 (Beans) 資料集，但它可以擴充套件到其他支援視覺模態的模型，也可以擴充套件到其他影像資料集。你可以嘗試的一些著名模型有:

此外，文章中介紹的方法也有可能擴充套件到其他模態。

要研究完整的影像相似度系統，你可以參考這個 Colab Notebook。

我們如何定義相似性？

要構建這個系統，我們首先需要定義我們想要如何計算兩個影像之間的相似度。一種廣泛流行的做法是先計算給定影像的稠密表徵 (即嵌入 (embedding))，然後使用餘弦相似性度量 (cosine similarity metric) 來確定兩幅影像的相似程度。

在本文中，我們將使用 “嵌入” 來表示向量空間中的影像。它為我們提供了一種將影像從高維畫素空間 (例如 224 × 224 × 3) 有意義地壓縮到一個低得多的維度 (例如 768) 的好方法。這樣做的主要優點是減少了後續步驟中的計算時間。

計算嵌入

為了計算影像的嵌入，我們需要使用一個視覺模型，該模型知道如何在向量空間中表示輸入影像。這種型別的模型通常也稱為影像編碼器 (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。

在本文中，我們快速入門並構建了一個影像相似度系統。如果你覺得這篇文章很有趣，我們強烈建議你基於我們討論的概念繼續構建你的系統，這樣你就可以更加熟悉內部工作原理。

還想了解更多嗎？以下是一些可能對你有用的其他資源:

英文原文: https://hf.co/blog/image-simi...
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。
審校、排版: zhongdongy (阿東)

基於 Hugging Face Datasets 和 Transformers 的影像相似性搜尋